More Pydantic AI changes

pamelafox · pamelafox · commit 202fa4b8e645 · 2025-05-06T22:17:37.000Z
diff --git a/src/backend/fastapi_app/api_models.py b/src/backend/fastapi_app/api_models.py
@@ -41,14 +41,34 @@ class ChatRequest(BaseModel):
     sessionState: Optional[Any] = None
 
 
+class ItemPublic(BaseModel):
+    id: int
+    type: str
+    brand: str
+    name: str
+    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
+
+    def __init__(self, **data):
+        super().__init__(**data)
+        self.distance = round(self.distance, 2)
+
+
 class ThoughtStep(BaseModel):
     title: str
     description: Any
     props: dict = {}
 
 
 class RAGContext(BaseModel):
-    data_points: dict[int, dict[str, Any]]
+    data_points: dict[int, ItemPublic]
     thoughts: list[ThoughtStep]
     followup_questions: Optional[list[str]] = None
 
@@ -69,26 +89,6 @@ class RetrievalResponseDelta(BaseModel):
     sessionState: Optional[Any] = None
 
 
-class ItemPublic(BaseModel):
-    id: int
-    type: str
-    brand: str
-    name: str
-    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
-
-    def __init__(self, **data):
-        super().__init__(**data)
-        self.distance = round(self.distance, 2)
-
-
 class ChatParams(ChatRequestOverrides):
     prompt_template: str
     response_token_limit: int = 1024
diff --git a/src/backend/fastapi_app/openai_clients.py b/src/backend/fastapi_app/openai_clients.py
@@ -14,7 +14,7 @@ async def create_openai_chat_client(
     openai_chat_client: Union[openai.AsyncAzureOpenAI, openai.AsyncOpenAI]
     OPENAI_CHAT_HOST = os.getenv("OPENAI_CHAT_HOST")
     if OPENAI_CHAT_HOST == "azure":
-        api_version = os.environ["AZURE_OPENAI_VERSION"] or "2024-03-01-preview"
+        api_version = os.environ["AZURE_OPENAI_VERSION"] or "2024-10-21"
         azure_endpoint = os.environ["AZURE_OPENAI_ENDPOINT"]
         azure_deployment = os.environ["AZURE_OPENAI_CHAT_DEPLOYMENT"]
         if api_key := os.getenv("AZURE_OPENAI_KEY"):
diff --git a/src/backend/fastapi_app/rag_advanced.py b/src/backend/fastapi_app/rag_advanced.py
@@ -52,6 +52,9 @@ class BrandFilter(TypedDict):
 
 
 class SearchResults(TypedDict):
+    query: str
+    """The original search query"""
+
     items: list[ItemPublic]
     """List of items that match the search query and filters"""
 
@@ -105,7 +108,9 @@ async def search_database(
             enable_text_search=ctx.deps.enable_text_search,
             filters=filters,
         )
-        return SearchResults(items=[ItemPublic.model_validate(item.to_dict()) for item in results], filters=filters)
+        return SearchResults(
+            query=search_query, items=[ItemPublic.model_validate(item.to_dict()) for item in results], filters=filters
+        )
 
     async def prepare_context(self, chat_params: ChatParams) -> tuple[list[ItemPublic], list[ThoughtStep]]:
         model = OpenAIModel(
@@ -119,35 +124,36 @@ async def prepare_context(self, chat_params: ChatParams) -> tuple[list[ItemPubli
             output_type=SearchResults,
         )
         # TODO: Provide few-shot examples
+        user_query = f"Find search results for user query: {chat_params.original_user_query}"
         results = await agent.run(
-            f"Find search results for user query: {chat_params.original_user_query}",
-            # message_history=chat_params.past_messages, # TODO
+            user_query,
+            message_history=chat_params.past_messages,
             deps=chat_params,
         )
-        items = results.output.items
+        items = results.output["items"]
         thoughts = [
             ThoughtStep(
                 title="Prompt to generate search arguments",
-                description=chat_params.past_messages,  # TODO: update this
+                description=results.all_messages(),
                 props=(
                     {"model": self.chat_model, "deployment": self.chat_deployment}
                     if self.chat_deployment
-                    else {"model": self.chat_model}
+                    else {"model": self.chat_model}  # TODO
                 ),
             ),
             ThoughtStep(
                 title="Search using generated search arguments",
-                description=chat_params.original_user_query,  # TODO:
+                description=results.output["query"],
                 props={
                     "top": chat_params.top,
                     "vector_search": chat_params.enable_vector_search,
                     "text_search": chat_params.enable_text_search,
-                    "filters": [],  # TODO
+                    "filters": results.output["filters"],
                 },
             ),
             ThoughtStep(
                 title="Search results",
-                description="",  # TODO
+                description=items,
             ),
         ]
         return items, thoughts
@@ -178,12 +184,12 @@ async def answer(
         return RetrievalResponse(
             message=Message(content=str(response.output), role=AIChatRoles.ASSISTANT),
             context=RAGContext(
-                data_points={},  # TODO
+                data_points={item.id: item for item in items},
                 thoughts=earlier_thoughts
                 + [
                     ThoughtStep(
                         title="Prompt to generate answer",
-                        description="",  # TODO: update
+                        description=response.all_messages(),
                         props=(
                             {"model": self.chat_model, "deployment": self.chat_deployment}
                             if self.chat_deployment
