fix type for streaming to conform with Microsoft Chat Protocol

Author: john0isaac

src/backend/fastapi_app/api_models.py

@@ -1,12 +1,19 @@
+from enum import Enum
 from typing import Any
 
 from openai.types.chat import ChatCompletionMessageParam
 from pydantic import BaseModel
 
 
+class AIChatRoles(str, Enum):
+    USER = "user"
+    ASSISTANT = "assistant"
+    SYSTEM = "system"
+
+
 class Message(BaseModel):
     content: str
-    role: str = "user"
+    role: AIChatRoles = AIChatRoles.USER
 
 
 class ChatRequest(BaseModel):
@@ -32,6 +39,12 @@ class RetrievalResponse(BaseModel):
     session_state: Any | None = None
 
 
+class RetrievalResponseDelta(BaseModel):
+    delta: Message | None = None
+    context: RAGContext | None = None
+    session_state: Any | None = None
+
+
 class ItemPublic(BaseModel):
     id: int
     type: str

src/backend/fastapi_app/rag_advanced.py

@@ -5,7 +5,14 @@
 from openai.types.chat import ChatCompletion, ChatCompletionChunk, ChatCompletionMessageParam
 from openai_messages_token_helper import build_messages, get_token_limit
 
-from fastapi_app.api_models import Message, RAGContext, RetrievalResponse, ThoughtStep
+from fastapi_app.api_models import (
+    AIChatRoles,
+    Message,
+    RAGContext,
+    RetrievalResponse,
+    RetrievalResponseDelta,
+    ThoughtStep,
+)
 from fastapi_app.postgres_models import Item
 from fastapi_app.postgres_searcher import PostgresSearcher
 from fastapi_app.query_rewriter import build_search_function, extract_search_arguments
@@ -110,10 +117,10 @@ async def run(
             stream=False,
         )
 
-        first_choice_message = chat_completion_response.choices[0].message
-
         return RetrievalResponse(
-            message=Message(content=str(first_choice_message.content), role=first_choice_message.role),
+            message=Message(
+                content=str(chat_completion_response.choices[0].message.content), role=AIChatRoles.ASSISTANT
+            ),
             context=RAGContext(
                 data_points={item.id: item.to_dict() for item in results},
                 thoughts=[
@@ -157,7 +164,7 @@ async def run_stream(
         self,
         messages: list[ChatCompletionMessageParam],
         overrides: dict[str, Any] = {},
-    ) -> AsyncGenerator[RetrievalResponse | Message, None]:
+    ) -> AsyncGenerator[RetrievalResponseDelta, None]:
         chat_params = self.get_params(messages, overrides)
 
         # Generate an optimized keyword search query based on the chat history and the last question
@@ -188,8 +195,7 @@ async def run_stream(
         # The connection closes when it returns back to the context manger in the dependencies
         await self.searcher.db_session.close()
 
-        yield RetrievalResponse(
-            message=Message(content="", role="assistant"),
+        yield RetrievalResponseDelta(
             context=RAGContext(
                 data_points={item.id: item.to_dict() for item in results},
                 thoughts=[
@@ -230,7 +236,9 @@ async def run_stream(
         )
 
         async for response_chunk in chat_completion_async_stream:
-            # first response has empty choices
-            if response_chunk.choices:
-                yield Message(content=str(response_chunk.choices[0].delta.content), role="assistant")
+            # first response has empty choices and last response has empty content
+            if response_chunk.choices and response_chunk.choices[0].delta.content:
+                yield RetrievalResponseDelta(
+                    delta=Message(content=str(response_chunk.choices[0].delta.content), role=AIChatRoles.ASSISTANT)
+                )
         return

src/backend/fastapi_app/rag_simple.py

@@ -8,7 +8,14 @@
 from openai_messages_token_helper import build_messages, get_token_limit
 from pydantic import BaseModel
 
-from fastapi_app.api_models import Message, RAGContext, RetrievalResponse, ThoughtStep
+from fastapi_app.api_models import (
+    AIChatRoles,
+    Message,
+    RAGContext,
+    RetrievalResponse,
+    RetrievalResponseDelta,
+    ThoughtStep,
+)
 from fastapi_app.postgres_models import Item
 from fastapi_app.postgres_searcher import PostgresSearcher
 
@@ -76,7 +83,7 @@ async def run_stream(
         self,
         messages: list[ChatCompletionMessageParam],
         overrides: dict[str, Any] = {},
-    ) -> AsyncGenerator[RetrievalResponse | Message, None]:
+    ) -> AsyncGenerator[RetrievalResponseDelta, None]:
         raise NotImplementedError
         if False:
             yield 0
@@ -145,10 +152,10 @@ async def run(
             stream=False,
         )
 
-        first_choice_message = chat_completion_response.choices[0].message
-
         return RetrievalResponse(
-            message=Message(content=str(first_choice_message.content), role=first_choice_message.role),
+            message=Message(
+                content=str(chat_completion_response.choices[0].message.content), role=AIChatRoles.ASSISTANT
+            ),
             context=RAGContext(
                 data_points={item.id: item.to_dict() for item in results},
                 thoughts=[
@@ -182,7 +189,7 @@ async def run_stream(
         self,
         messages: list[ChatCompletionMessageParam],
         overrides: dict[str, Any] = {},
-    ) -> AsyncGenerator[RetrievalResponse | Message, None]:
+    ) -> AsyncGenerator[RetrievalResponseDelta, None]:
         chat_params = self.get_params(messages, overrides)
 
         # Retrieve relevant items from the database
@@ -206,8 +213,7 @@ async def run_stream(
         # The connection closes when it returns back to the context manger in the dependencies
         await self.searcher.db_session.close()
 
-        yield RetrievalResponse(
-            message=Message(content="", role="assistant"),
+        yield RetrievalResponseDelta(
             context=RAGContext(
                 data_points={item.id: item.to_dict() for item in results},
                 thoughts=[
@@ -237,7 +243,9 @@ async def run_stream(
             ),
         )
         async for response_chunk in chat_completion_async_stream:
-            # first response has empty choices
-            if response_chunk.choices:
-                yield Message(content=str(response_chunk.choices[0].delta.content), role="assistant")
+            # first response has empty choices and last response has empty content
+            if response_chunk.choices and response_chunk.choices[0].delta.content:
+                yield RetrievalResponseDelta(
+                    delta=Message(content=str(response_chunk.choices[0].delta.content), role=AIChatRoles.ASSISTANT)
+                )
         return

src/backend/fastapi_app/routes/api_routes.py

@@ -7,7 +7,13 @@
 from fastapi.responses import StreamingResponse
 from sqlalchemy import select
 
-from fastapi_app.api_models import ChatRequest, ItemPublic, ItemWithDistance, Message, RetrievalResponse
+from fastapi_app.api_models import (
+    ChatRequest,
+    ItemPublic,
+    ItemWithDistance,
+    RetrievalResponse,
+    RetrievalResponseDelta,
+)
 from fastapi_app.dependencies import ChatClient, CommonDeps, DBSession, EmbeddingsClient
 from fastapi_app.postgres_models import Item
 from fastapi_app.postgres_searcher import PostgresSearcher
@@ -17,13 +23,13 @@
 router = fastapi.APIRouter()
 
 
-async def format_as_ndjson(r: AsyncGenerator[RetrievalResponse | Message, None]) -> AsyncGenerator[str, None]:
+async def format_as_ndjson(r: AsyncGenerator[RetrievalResponseDelta, None]) -> AsyncGenerator[str, None]:
     """
     Format the response as NDJSON
     """
     try:
         async for event in r:
-            yield json.dumps(event.model_dump(), ensure_ascii=False) + "\n"
+            yield event.model_dump_json() + "\n"
     except Exception as error:
         logging.exception("Exception while generating response stream: %s", error)
         yield json.dumps({"error": str(error)}, ensure_ascii=False) + "\n"