Merge pull request #6 from rohit-ganguly/newdata-tests

Tests for new data, all tests passing

Author: pamelafox

.vscode/settings.json

@@ -1,22 +1,4 @@
 {
-    "pgsql.connections": [
-        {
-            "id": "92CC1089-BAD0-44A4-B071-A50A6EC12B67",
-            "groupId": "F3347CD6-9995-4EE9-8D98-D88DB010FA5B",
-            "authenticationType": "SqlLogin",
-            "connectTimeout": 15,
-            "applicationName": "vscode-pgsql",
-            "clientEncoding": "utf8",
-            "sslmode": "prefer",
-            "server": "localhost",
-            "user": "admin",
-            "password": "",
-            "savePassword": true,
-            "database": "postgres",
-            "profileName": "local-pg",
-            "expiresOn": 0
-        }
-    ],
     "python.testing.pytestArgs": [
         "tests"
     ],

convert_csv_json.py

@@ -1,6 +1,7 @@
 import ast
 import csv
 import json
+from typing import Any
 
 # Read CSV file - Using the correct dialect to handle quotes properly
 with open("pittsburgh_restaurants.csv", encoding="utf-8") as csv_file:
@@ -15,7 +16,7 @@
     item = {}
     for i in range(len(header)):
         if i < len(row):  # Ensure we don't go out of bounds
-            value = row[i].strip()
+            value: Any = row[i].strip()
             # Check if the value looks like a JSON array
             if value.startswith("[") and value.endswith("]"):
                 try:

evals/generate_ground_truth.py

@@ -56,18 +56,17 @@ def source_retriever() -> Generator[str, None, None]:
     DATABASE_URI = f"postgresql://{DBUSER}:{DBPASS}@{DBHOST}/{DBNAME}"
     engine = create_engine(DATABASE_URI, echo=False)
     with Session(engine) as session:
-        # Fetch all products for a particular type
-        item_types = session.scalars(select(Item.type).distinct())
-        for item_type in item_types:
-            records = list(session.scalars(select(Item).filter(Item.type == item_type).order_by(Item.id)))
-            logger.info(f"Processing database records for type: {item_type}")
-            yield "\n\n".join([f"## Product ID: [{record.id}]\n" + record.to_str_for_rag() for record in records])
+        # Fetch all products for a particular type - depends on the database columns
+        # item_types = session.scalars(select(Item.type).distinct())
+        # for item_type in item_types:
+        #    records = list(session.scalars(select(Item).filter(Item.type == item_type).order_by(Item.id)))
+        #    logger.info(f"Processing database records for type: {item_type}")
+        #    yield "\n\n".join([f"## Product ID: [{record.id}]\n" + record.to_str_for_rag() for record in records])
         # Fetch each item individually
-        # records = list(session.scalars(select(Item).order_by(Item.id)))
-        # for record in records:
-        #    logger.info(f"Processing database record: {record.name}")
-        #    yield f"## Product ID: [{record.id}]\n" + record.to_str_for_rag()
-        # await self.openai_chat_client.chat.completions.create(
+        records = list(session.scalars(select(Item).order_by(Item.id)))
+        for record in records:
+            logger.info(f"Processing database record: {record.name}")
+            yield f"## Product ID: [{record.id}]\n" + record.to_str_for_rag()
 
 
 def source_to_text(source) -> str:

src/backend/fastapi_app/api_models.py

@@ -45,7 +45,7 @@ class ItemPublic(BaseModel):
     id: str
     name: str
     cuisine: str
-    rating: int
+    rating: float
     price_level: int
     review_count: int
     description: str

src/backend/fastapi_app/postgres_models.py

@@ -15,7 +15,7 @@ class Item(Base):
     id: Mapped[str] = mapped_column(primary_key=True)
     name: Mapped[str] = mapped_column()
     cuisine: Mapped[str] = mapped_column()
-    rating: Mapped[int] = mapped_column()
+    rating: Mapped[float] = mapped_column()
     price_level: Mapped[int] = mapped_column()
     review_count: Mapped[int] = mapped_column()
     description: Mapped[str] = mapped_column()

src/backend/fastapi_app/routes/api_routes.py

@@ -45,7 +45,7 @@ async def format_as_ndjson(r: AsyncGenerator[RetrievalResponseDelta, None]) -> A
 
 
 @router.get("/items/{id}", response_model=ItemPublic)
-async def item_handler(database_session: DBSession, id: int) -> ItemPublic:
+async def item_handler(database_session: DBSession, id: str) -> ItemPublic:
     """A simple API to get an item by ID."""
     item = (await database_session.scalars(select(Item).where(Item.id == id))).first()
     if not item:
@@ -55,7 +55,7 @@ async def item_handler(database_session: DBSession, id: int) -> ItemPublic:
 
 @router.get("/similar", response_model=list[ItemWithDistance])
 async def similar_handler(
-    context: CommonDeps, database_session: DBSession, id: int, n: int = 5
+    context: CommonDeps, database_session: DBSession, id: str, 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()