More Pydantic-AI usage

pamelafox · pamelafox · commit b1b8746611fb · 2025-05-05T22:35:57.000Z
diff --git a/src/backend/fastapi_app/__init__.py b/src/backend/fastapi_app/__init__.py
@@ -34,7 +34,13 @@ class State(TypedDict):
 @asynccontextmanager
 async def lifespan(app: fastapi.FastAPI) -> AsyncIterator[State]:
     context = await common_parameters()
-    azure_credential = await get_azure_credential()
+    azure_credential = None
+    if (
+        os.getenv("OPENAI_CHAT_HOST") == "azure"
+        or os.getenv("OPENAI_EMBED_HOST") == "azure"
+        or os.getenv("POSTGRES_HOST").endswith(".database.azure.com")
+    ):
+        azure_credential = await get_azure_credential()
     engine = await create_postgres_engine_from_env(azure_credential)
     sessionmaker = await create_async_sessionmaker(engine)
     chat_client = await create_openai_chat_client(azure_credential)
diff --git a/src/backend/fastapi_app/api_models.py b/src/backend/fastapi_app/api_models.py
@@ -77,6 +77,9 @@ class ItemPublic(BaseModel):
     description: str
     price: float
 
+    def to_str_for_rag(self):
+        return f"Name:{self.name} Description:{self.description} Price:{self.price} Brand:{self.brand} Type:{self.type}"
+
 
 class ItemWithDistance(ItemPublic):
     distance: float
diff --git a/src/backend/fastapi_app/rag_advanced.py b/src/backend/fastapi_app/rag_advanced.py
@@ -12,6 +12,7 @@
 
 from fastapi_app.api_models import (
     AIChatRoles,
+    ItemPublic,
     Message,
     RAGContext,
     RetrievalResponse,
@@ -50,6 +51,14 @@ class BrandFilter(TypedDict):
     """The brand name to compare against (e.g., 'AirStrider')"""
 
 
+class SearchResults(TypedDict):
+    items: list[ItemPublic]
+    """List of items that match the search query and filters"""
+
+    filters: list[Union[PriceFilter, BrandFilter]]
+    """List of filters applied to the search results"""
+
+
 class AdvancedRAGChat(RAGChatBase):
     def __init__(
         self,
@@ -71,7 +80,7 @@ async def search_database(
         search_query: str,
         price_filter: Optional[PriceFilter] = None,
         brand_filter: Optional[BrandFilter] = None,
-    ) -> list[str]:
+    ) -> SearchResults:
         """
         Search PostgreSQL database for relevant products based on user query
 
@@ -83,7 +92,6 @@ async def search_database(
         Returns:
             List of formatted items that match the search query and filters
         """
-        print(search_query, price_filter, brand_filter)
         # Only send non-None filters
         filters = []
         if price_filter:
@@ -97,9 +105,9 @@ async def search_database(
             enable_text_search=ctx.deps.enable_text_search,
             filters=filters,
         )
-        return [f"[{(item.id)}]:{item.to_str_for_rag()}\n\n" for item in results]
+        return SearchResults(items=[ItemPublic.model_validate(item.to_dict()) for item in results], filters=filters)
 
-    async def prepare_context(self, chat_params: ChatParams) -> tuple[str, list[Item], list[ThoughtStep]]:
+    async def prepare_context(self, chat_params: ChatParams) -> tuple[list[ItemPublic], list[ThoughtStep]]:
         model = OpenAIModel(
             os.environ["AZURE_OPENAI_CHAT_DEPLOYMENT"], provider=OpenAIProvider(openai_client=self.openai_chat_client)
         )
@@ -108,17 +116,15 @@ async def prepare_context(self, chat_params: ChatParams) -> tuple[str, list[Item
             model_settings=ModelSettings(temperature=0.0, max_tokens=500, seed=chat_params.seed),
             system_prompt=self.query_prompt_template,
             tools=[self.search_database],
-            output_type=list[str],
+            output_type=SearchResults,
         )
         # TODO: Provide few-shot examples
         results = await agent.run(
             f"Find search results for user query: {chat_params.original_user_query}",
             # message_history=chat_params.past_messages, # TODO
             deps=chat_params,
         )
-        if not isinstance(results, list):
-            raise ValueError("Search results should be a list of strings")
-
+        items = results.output.items
         thoughts = [
             ThoughtStep(
                 title="Prompt to generate search arguments",
@@ -144,12 +150,12 @@ async def prepare_context(self, chat_params: ChatParams) -> tuple[str, list[Item
                 description="",  # TODO
             ),
         ]
-        return results, thoughts
+        return items, thoughts
 
     async def answer(
         self,
         chat_params: ChatParams,
-        results: list[str],
+        items: list[ItemPublic],
         earlier_thoughts: list[ThoughtStep],
     ) -> RetrievalResponse:
         agent = Agent(
@@ -163,15 +169,16 @@ async def answer(
             ),
         )
 
+        item_references = [item.to_str_for_rag() for item in items]
         response = await agent.run(
-            user_prompt=chat_params.original_user_query + "Sources:\n" + "\n".join(results),
+            user_prompt=chat_params.original_user_query + "Sources:\n" + "\n".join(item_references),
             message_history=chat_params.past_messages,
         )
 
         return RetrievalResponse(
             message=Message(content=str(response.output), role=AIChatRoles.ASSISTANT),
             context=RAGContext(
-                data_points={item.id: item.to_dict() for item in results},
+                data_points={},  # TODO
                 thoughts=earlier_thoughts
                 + [
                     ThoughtStep(
diff --git a/src/backend/fastapi_app/routes/api_routes.py b/src/backend/fastapi_app/routes/api_routes.py
@@ -136,8 +136,8 @@ async def chat_handler(
 
         chat_params = rag_flow.get_params(chat_request.messages, chat_request.context.overrides)
 
-        results, thoughts = await rag_flow.prepare_context(chat_params)
-        response = await rag_flow.answer(chat_params=chat_params, results=results, earlier_thoughts=thoughts)
+        items, thoughts = await rag_flow.prepare_context(chat_params)
+        response = await rag_flow.answer(chat_params=chat_params, items=items, earlier_thoughts=thoughts)
         return response
     except Exception as e:
         if isinstance(e, APIError) and e.code == "content_filter":
