Make mypy happy

pamelafox · pamelafox · commit ae0871dd8250 · 2024-08-03T12:35:19.000Z
diff --git a/src/backend/fastapi_app/embeddings.py b/src/backend/fastapi_app/embeddings.py
@@ -10,7 +10,7 @@ async def compute_text_embedding(
     openai_client: AsyncOpenAI | AsyncAzureOpenAI,
     embed_model: str,
     embed_deployment: str | None = None,
-    embedding_dimensions: int = 1536,
+    embedding_dimensions: int | None = None,
 ) -> list[float]:
     SUPPORTED_DIMENSIONS_MODEL = {
         "text-embedding-ada-002": False,
@@ -21,9 +21,12 @@ async def compute_text_embedding(
     class ExtraArgs(TypedDict, total=False):
         dimensions: int
 
-    dimensions_args: ExtraArgs = (
-        {"dimensions": embedding_dimensions} if SUPPORTED_DIMENSIONS_MODEL.get(embed_model) else {}
-    )
+    dimensions_args: ExtraArgs = {}
+    if SUPPORTED_DIMENSIONS_MODEL.get(embed_model):
+        if embedding_dimensions is None:
+            raise ValueError(f"Model {embed_model} requires embedding dimensions")
+        else:
+            dimensions_args = {"dimensions": embedding_dimensions}
 
     embedding = await openai_client.embeddings.create(
         # Azure OpenAI takes the deployment name as the model name
diff --git a/src/backend/fastapi_app/postgres_engine.py b/src/backend/fastapi_app/postgres_engine.py
@@ -30,7 +30,7 @@ def get_password_from_azure_credential():
 
     engine = create_async_engine(
         DATABASE_URI,
-        echo=False,
+        echo=True,
     )
 
     @event.listens_for(engine.sync_engine, "do_connect")
diff --git a/src/backend/fastapi_app/postgres_models.py b/src/backend/fastapi_app/postgres_models.py
@@ -21,7 +21,7 @@ class Item(Base):
     description: Mapped[str] = mapped_column()
     price: Mapped[float] = mapped_column()
     embedding_ada002: Mapped[Vector] = mapped_column(Vector(1536))  # ada-002
-    embedding_nomic: Mapped[Vector] | None = mapped_column(Vector(768), nullable=True)  # nomic-embed-text
+    embedding_nomic: Mapped[Vector] = mapped_column(Vector(768))  # nomic-embed-text
 
     def to_dict(self, include_embedding: bool = False):
         model_dict = asdict(self)
diff --git a/src/backend/fastapi_app/postgres_searcher.py b/src/backend/fastapi_app/postgres_searcher.py
@@ -14,7 +14,7 @@ def __init__(
         openai_embed_client: AsyncOpenAI | AsyncAzureOpenAI,
         embed_deployment: str | None,  # Not needed for non-Azure OpenAI or for retrieval_mode="text"
         embed_model: str,
-        embed_dimensions: int,
+        embed_dimensions: int | None,
         embedding_column: str,
     ):
         self.db_session = db_session
diff --git a/src/backend/fastapi_app/routes/api_routes.py b/src/backend/fastapi_app/routes/api_routes.py
@@ -45,15 +45,18 @@ async def item_handler(database_session: DBSession, id: int) -> ItemPublic:
 
 
 @router.get("/similar", response_model=list[ItemWithDistance])
-async def similar_handler(database_session: DBSession, id: int, n: int = 5) -> list[ItemWithDistance]:
+async def similar_handler(
+    context: CommonDeps, database_session: DBSession, id: int, n: int = 5
+) -> list[ItemWithDistance]:
     """A similarity API to find items similar to items with given ID."""
     item = (await database_session.scalars(select(Item).where(Item.id == id))).first()
     if not item:
         raise HTTPException(detail=f"Item with ID {id} not found.", status_code=404)
+
     closest = await database_session.execute(
-        select(Item, Item.embedding.l2_distance(item.embedding))
+        select(Item, Item.embedding_ada002.l2_distance(item.embedding_ada002))
         .filter(Item.id != id)
-        .order_by(Item.embedding.l2_distance(item.embedding))
+        .order_by(Item.embedding_ada002.l2_distance(item.embedding_ada002))
         .limit(n)
     )
     return [
@@ -78,6 +81,7 @@ async def search_handler(
         embed_deployment=context.openai_embed_deployment,
         embed_model=context.openai_embed_model,
         embed_dimensions=context.openai_embed_dimensions,
+        embedding_column=context.embedding_column,
     )
     results = await searcher.search_and_embed(
         query, top=top, enable_vector_search=enable_vector_search, enable_text_search=enable_text_search
diff --git a/src/backend/fastapi_app/update_embeddings.py b/src/backend/fastapi_app/update_embeddings.py
@@ -63,9 +63,9 @@ async def update_embeddings(in_seed_data=False):
 
     async with async_sessionmaker(engine, expire_on_commit=False)() as session:
         async with session.begin():
-            items = (await session.scalars(select(Item))).all()
+            items_to_update = (await session.scalars(select(Item))).all()
 
-            for item in items:
+            for item in items_to_update:
                 setattr(
                     item,
                     embedding_column,
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -263,4 +263,5 @@ async def postgres_searcher(mock_session_env, mock_default_azure_credential, db_
         embed_deployment="text-embedding-ada-002",
         embed_model="text-embedding-ada-002",
         embed_dimensions=1536,
+        embedding_column="embedding_ada002",
     )

Original file line number	Diff line number	Diff line change
`@@ -30,7 +30,7 @@ def get_password_from_azure_credential():`
`30`	`30`
`31`	`31`	`engine = create_async_engine(`
`32`	`32`	`DATABASE_URI,`
`33`		`- echo=False,`
	`33`	`+ echo=True,`
`34`	`34`	`)`
`35`	`35`
`36`	`36`	`@event.listens_for(engine.sync_engine, "do_connect")`
Original file line number	Diff line number	Diff line change
`@@ -263,4 +263,5 @@ async def postgres_searcher(mock_session_env, mock_default_azure_credential, db_`
`263`	`263`	`embed_deployment="text-embedding-ada-002",`
`264`	`264`	`embed_model="text-embedding-ada-002",`
`265`	`265`	`embed_dimensions=1536,`
	`266`	`+ embedding_column="embedding_ada002",`
`266`	`267`	`)`