add pydatnic types for Item table

Author: john0isaac

src/fastapi_app/api_models.py

@@ -30,3 +30,16 @@ class RetrievalResponse(BaseModel):
     message: Message
     context: RAGContext
     session_state: Any | None = None
+
+
+class ItemPublic(BaseModel):
+    id: int
+    type: str
+    brand: str
+    name: str
+    description: str
+    price: float
+
+
+class ItemWithDistance(ItemPublic):
+    distance: float

src/fastapi_app/routes/api_routes.py

@@ -1,10 +1,8 @@
-from typing import Any
-
 import fastapi
 from fastapi import HTTPException
 from sqlalchemy import select
 
-from fastapi_app.api_models import ChatRequest, RetrievalResponse
+from fastapi_app.api_models import ChatRequest, ItemPublic, ItemWithDistance, RetrievalResponse
 from fastapi_app.dependencies import ChatClient, CommonDeps, DBSession, EmbeddingsClient
 from fastapi_app.postgres_models import Item
 from fastapi_app.postgres_searcher import PostgresSearcher
@@ -14,17 +12,17 @@
 router = fastapi.APIRouter()
 
 
-@router.get("/items/{id}", response_model=dict[str, Any])
-async def item_handler(id: int, database_session: DBSession) -> dict[str, Any]:
+@router.get("/items/{id}", response_model=ItemPublic)
+async def item_handler(id: int, database_session: DBSession) -> 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:
         raise HTTPException(detail=f"Item with ID {id} not found.", status_code=404)
-    return item.to_dict()
+    return ItemPublic.model_validate(item.to_dict())
 
 
-@router.get("/similar", response_model=list[dict[str, Any]])
-async def similar_handler(database_session: DBSession, id: int, n: int = 5) -> list[dict[str, Any]]:
+@router.get("/similar", response_model=list[ItemWithDistance])
+async def similar_handler(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:
@@ -35,10 +33,12 @@ async def similar_handler(database_session: DBSession, id: int, n: int = 5) -> l
         .order_by(Item.embedding.l2_distance(item.embedding))
         .limit(n)
     )
-    return [item.to_dict() | {"distance": round(distance, 2)} for item, distance in closest]
+    return [
+        ItemWithDistance.model_validate(item.to_dict() | {"distance": round(distance, 2)}) for item, distance in closest
+    ]
 
 
-@router.get("/search", response_model=list[dict[str, Any]])
+@router.get("/search", response_model=list[ItemPublic])
 async def search_handler(
     context: CommonDeps,
     database_session: DBSession,
@@ -47,7 +47,7 @@ async def search_handler(
     top: int = 5,
     enable_vector_search: bool = True,
     enable_text_search: bool = True,
-) -> list[dict[str, Any]]:
+) -> list[ItemPublic]:
     """A search API to find items based on a query."""
     searcher = PostgresSearcher(
         db_session=database_session,
@@ -59,7 +59,7 @@ async def search_handler(
     results = await searcher.search_and_embed(
         query, top=top, enable_vector_search=enable_vector_search, enable_text_search=enable_text_search
     )
-    return [item.to_dict() for item in results]
+    return [ItemPublic.model_validate(item.to_dict()) for item in results]
 
 
 @router.post("/chat", response_model=RetrievalResponse)

tests/data.py

@@ -1,14 +1,7 @@
-from dataclasses import dataclass
+from fastapi_app.api_models import ItemPublic
 
 
-@dataclass
-class TestData:
-    id: int
-    type: str
-    brand: str
-    name: str
-    description: str
-    price: float
+class TestData(ItemPublic):
     embeddings: list[float]