use basemodel instead of dataclass to match the other models and for validation

Author: john0isaac

src/fastapi_app/api_models.py

@@ -1,4 +1,3 @@
-from dataclasses import dataclass
 from typing import Any
 
 from pydantic import BaseModel
@@ -20,15 +19,13 @@ class ThoughtStep(BaseModel):
     props: dict = {}
 
 
-@dataclass
-class RAGContext:
+class RAGContext(BaseModel):
     data_points: dict[int, dict[str, Any]]
     thoughts: list[ThoughtStep]
     followup_questions: list[str] | None = None
 
 
-@dataclass
-class RetrievalResponse:
+class RetrievalResponse(BaseModel):
     message: Message
     context: RAGContext
     session_state: Any | None = None

src/fastapi_app/rag_advanced.py

@@ -8,7 +8,7 @@
 from openai.types.chat import ChatCompletion, ChatCompletionMessageParam
 from openai_messages_token_helper import build_messages, get_token_limit
 
-from .api_models import RAGContext, RetrievalResponse, ThoughtStep
+from .api_models import Message, RAGContext, RetrievalResponse, ThoughtStep
 from .postgres_searcher import PostgresSearcher
 from .query_rewriter import build_search_function, extract_search_arguments
 
@@ -99,7 +99,7 @@ async def run(
         )
         first_choice = chat_completion_response.choices[0]
         return RetrievalResponse(
-            message=first_choice.message,
+            message=Message(content=first_choice.message.content, role=first_choice.message.role),
             context=RAGContext(
                 data_points={item.id: item.to_dict() for item in results},
                 thoughts=[

src/fastapi_app/rag_simple.py

@@ -8,7 +8,7 @@
 from openai.types.chat import ChatCompletion, ChatCompletionMessageParam
 from openai_messages_token_helper import build_messages, get_token_limit
 
-from .api_models import RAGContext, RetrievalResponse, ThoughtStep
+from .api_models import Message, RAGContext, RetrievalResponse, ThoughtStep
 from .postgres_searcher import PostgresSearcher
 
 
@@ -69,7 +69,7 @@ async def run(
         )
         first_choice = chat_completion_response.choices[0]
         return RetrievalResponse(
-            message=first_choice.message,
+            message=Message(content=first_choice.message.content, role=first_choice.message.role),
             context=RAGContext(
                 data_points={item.id: item.to_dict() for item in results},
                 thoughts=[