Fix tests, mypy

Author: pamelafox

src/backend/fastapi_app/api_models.py

@@ -1,7 +1,7 @@
 from enum import Enum
 from typing import Any, Optional
 
-from openai.types.chat import ChatCompletionMessageParam
+from openai.types.responses import ResponseInputItemParam
 from pydantic import BaseModel, Field
 
 
@@ -36,7 +36,7 @@ class ChatRequestContext(BaseModel):
 
 
 class ChatRequest(BaseModel):
-    messages: list[ChatCompletionMessageParam]
+    messages: list[ResponseInputItemParam]
     context: ChatRequestContext
     sessionState: Optional[Any] = None
 
@@ -95,7 +95,7 @@ class ChatParams(ChatRequestOverrides):
     enable_text_search: bool
     enable_vector_search: bool
     original_user_query: str
-    past_messages: list[ChatCompletionMessageParam]
+    past_messages: list[ResponseInputItemParam]
 
 
 class Filter(BaseModel):

src/backend/fastapi_app/prompts/query_fewshots.json

@@ -1,34 +1,36 @@
 [
-  {"role": "user", "content": "good options for climbing gear that can be used outside?"},
-  {"role": "assistant", "tool_calls": [
-      {
-          "id": "call_abc123",
-          "type": "function",
-          "function": {
-              "arguments": "{\"search_query\":\"climbing gear outside\"}",
-              "name": "search_database"
-          }
-      }
-  ]},
   {
-      "role": "tool",
-      "tool_call_id": "call_abc123",
-      "content": "Search results for climbing gear that can be used outside: ..."
+    "role": "user",
+    "content": "good options for climbing gear that can be used outside?"
   },
-  {"role": "user", "content": "are there any shoes less than $50?"},
-  {"role": "assistant", "tool_calls": [
-      {
-          "id": "call_abc456",
-          "type": "function",
-          "function": {
-              "arguments": "{\"search_query\":\"shoes\",\"price_filter\":{\"comparison_operator\":\"<\",\"value\":50}}",
-              "name": "search_database"
-          }
-      }
-  ]},
   {
-      "role": "tool",
-      "tool_call_id": "call_abc456",
-      "content": "Search results for shoes cheaper than 50: ..."
+    "id": "madeup",
+    "call_id": "call_abc123",
+    "name": "search_database",
+    "arguments": "{\"search_query\":\"climbing gear outside\"}",
+    "type": "function_call"
+  },
+  {
+    "id": "madeupoutput",
+    "call_id": "call_abc123",
+    "output": "Search results for climbing gear that can be used outside: ...",
+    "type": "function_call_output"
+  },
+  {
+    "role": "user",
+    "content": "are there any shoes less than $50?"
+  },
+  {
+    "id": "madeup",
+    "call_id": "call_abc456",
+    "name": "search_database",
+    "arguments": "{\"search_query\":\"shoes\",\"price_filter\":{\"comparison_operator\":\"<\",\"value\":50}}",
+    "type": "function_call"
+  },
+  {
+    "id": "madeupoutput",
+    "call_id": "call_abc456",
+    "output": "Search results for shoes cheaper than 50: ...",
+    "type": "function_call_output"
   }
 ]

src/backend/fastapi_app/rag_advanced.py

@@ -2,17 +2,17 @@
 from collections.abc import AsyncGenerator
 from typing import Optional, Union
 
-from agents import Agent, ModelSettings, OpenAIChatCompletionsModel, Runner, function_tool, set_tracing_disabled
-from openai import AsyncAzureOpenAI, AsyncOpenAI
-from openai.types.chat import (
-    ChatCompletionMessageParam,
-)
-from openai.types.responses import (
-    EasyInputMessageParam,
-    ResponseFunctionToolCallParam,
-    ResponseTextDeltaEvent,
+from agents import (
+    Agent,
+    ModelSettings,
+    OpenAIChatCompletionsModel,
+    Runner,
+    ToolCallOutputItem,
+    function_tool,
+    set_tracing_disabled,
 )
-from openai.types.responses.response_input_item_param import FunctionCallOutput
+from openai import AsyncAzureOpenAI, AsyncOpenAI
+from openai.types.responses import EasyInputMessageParam, ResponseInputItemParam, ResponseTextDeltaEvent
 
 from fastapi_app.api_models import (
     AIChatRoles,
@@ -41,7 +41,7 @@ class AdvancedRAGChat(RAGChatBase):
     def __init__(
         self,
         *,
-        messages: list[ChatCompletionMessageParam],
+        messages: list[ResponseInputItemParam],
         overrides: ChatRequestOverrides,
         searcher: PostgresSearcher,
         openai_chat_client: Union[AsyncOpenAI, AsyncAzureOpenAI],
@@ -109,34 +109,17 @@ async def search_database(
         )
 
     async def prepare_context(self) -> tuple[list[ItemPublic], list[ThoughtStep]]:
-        few_shots = json.loads(self.query_fewshots)
-        few_shot_inputs = []
-        for few_shot in few_shots:
-            if few_shot["role"] == "user":
-                message = EasyInputMessageParam(role="user", content=few_shot["content"])
-            elif few_shot["role"] == "assistant" and few_shot["tool_calls"] is not None:
-                message = ResponseFunctionToolCallParam(
-                    id="madeup",
-                    call_id=few_shot["tool_calls"][0]["id"],
-                    name=few_shot["tool_calls"][0]["function"]["name"],
-                    arguments=few_shot["tool_calls"][0]["function"]["arguments"],
-                    type="function_call",
-                )
-            elif few_shot["role"] == "tool" and few_shot["tool_call_id"] is not None:
-                message = FunctionCallOutput(
-                    id="madeupoutput",
-                    call_id=few_shot["tool_call_id"],
-                    output=few_shot["content"],
-                    type="function_call_output",
-                )
-            few_shot_inputs.append(message)
-
+        few_shots: list[ResponseInputItemParam] = json.loads(self.query_fewshots)
         user_query = f"Find search results for user query: {self.chat_params.original_user_query}"
         new_user_message = EasyInputMessageParam(role="user", content=user_query)
-        all_messages = few_shot_inputs + self.chat_params.past_messages + [new_user_message]
+        all_messages = few_shots + self.chat_params.past_messages + [new_user_message]
 
         run_results = await Runner.run(self.search_agent, input=all_messages)
-        search_results = run_results.new_items[-1].output
+        most_recent_response = run_results.new_items[-1]
+        if isinstance(most_recent_response, ToolCallOutputItem):
+            search_results = most_recent_response.output
+        else:
+            raise ValueError("Error retrieving search results, model did not call tool properly")
 
         thoughts = [
             ThoughtStep(

src/backend/fastapi_app/rag_base.py

@@ -2,7 +2,7 @@
 from abc import ABC, abstractmethod
 from collections.abc import AsyncGenerator
 
-from openai.types.chat import ChatCompletionMessageParam
+from openai.types.responses import ResponseInputItemParam
 
 from fastapi_app.api_models import (
     ChatParams,
@@ -18,16 +18,14 @@ class RAGChatBase(ABC):
     prompts_dir = pathlib.Path(__file__).parent / "prompts/"
     answer_prompt_template = open(prompts_dir / "answer.txt").read()
 
-    def get_chat_params(
-        self, messages: list[ChatCompletionMessageParam], overrides: ChatRequestOverrides
-    ) -> ChatParams:
+    def get_chat_params(self, messages: list[ResponseInputItemParam], overrides: ChatRequestOverrides) -> ChatParams:
         response_token_limit = 1024
         prompt_template = overrides.prompt_template or self.answer_prompt_template
 
         enable_text_search = overrides.retrieval_mode in ["text", "hybrid", None]
         enable_vector_search = overrides.retrieval_mode in ["vectors", "hybrid", None]
 
-        original_user_query = messages[-1]["content"]
+        original_user_query = messages[-1].get("content")
         if not isinstance(original_user_query, str):
             raise ValueError("The most recent message content must be a string.")
 

src/backend/fastapi_app/rag_simple.py

@@ -3,8 +3,7 @@
 
 from agents import Agent, ModelSettings, OpenAIChatCompletionsModel, Runner, set_tracing_disabled
 from openai import AsyncAzureOpenAI, AsyncOpenAI
-from openai.types.chat import ChatCompletionMessageParam
-from openai.types.responses import ResponseTextDeltaEvent
+from openai.types.responses import ResponseInputItemParam, ResponseTextDeltaEvent
 
 from fastapi_app.api_models import (
     AIChatRoles,
@@ -26,7 +25,7 @@ class SimpleRAGChat(RAGChatBase):
     def __init__(
         self,
         *,
-        messages: list[ChatCompletionMessageParam],
+        messages: list[ResponseInputItemParam],
         overrides: ChatRequestOverrides,
         searcher: PostgresSearcher,
         openai_chat_client: Union[AsyncOpenAI, AsyncAzureOpenAI],