implement chat/stream/ endpoint and fix empty choices error

Author: john0isaac

src/backend/fastapi_app/rag_simple.py

@@ -232,5 +232,7 @@ async def run_stream(
             ),
         )
         async for response_chunk in chat_completion_async_stream:
-            yield Message(content=str(response_chunk.choices[0].delta.content), role="assistant")
+            # first response has empty choices
+            if response_chunk.choices:
+                yield Message(content=str(response_chunk.choices[0].delta.content), role="assistant")
         return

src/backend/fastapi_app/routes/api_routes.py

@@ -1,8 +1,13 @@
+import json
+import logging
+from collections.abc import AsyncGenerator
+
 import fastapi
 from fastapi import HTTPException
+from fastapi.responses import StreamingResponse
 from sqlalchemy import select
 
-from fastapi_app.api_models import ChatRequest, ItemPublic, ItemWithDistance, RetrievalResponse
+from fastapi_app.api_models import ChatRequest, ItemPublic, ItemWithDistance, Message, RetrievalResponse
 from fastapi_app.dependencies import ChatClient, CommonDeps, DBSession, EmbeddingsClient
 from fastapi_app.postgres_models import Item
 from fastapi_app.postgres_searcher import PostgresSearcher
@@ -12,6 +17,18 @@
 router = fastapi.APIRouter()
 
 
+async def format_as_ndjson(r: AsyncGenerator[RetrievalResponse | Message, 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"
+    except Exception as error:
+        logging.exception("Exception while generating response stream: %s", error)
+        yield json.dumps({"error": str(error)}, ensure_ascii=False) + "\n"
+
+
 @router.get("/items/{id}", response_model=ItemPublic)
 async def item_handler(database_session: DBSession, id: int) -> ItemPublic:
     """A simple API to get an item by ID."""
@@ -96,3 +113,39 @@ async def chat_handler(
 
     response = await run_ragchat(chat_request.messages, overrides=overrides)
     return response
+
+
+@router.post("/chat/stream")
+async def chat_stream_handler(
+    context: CommonDeps,
+    database_session: DBSession,
+    openai_embed: EmbeddingsClient,
+    openai_chat: ChatClient,
+    chat_request: ChatRequest,
+):
+    overrides = chat_request.context.get("overrides", {})
+
+    searcher = PostgresSearcher(
+        db_session=database_session,
+        openai_embed_client=openai_embed.client,
+        embed_deployment=context.openai_embed_deployment,
+        embed_model=context.openai_embed_model,
+        embed_dimensions=context.openai_embed_dimensions,
+    )
+    if overrides.get("use_advanced_flow"):
+        run_ragchat = AdvancedRAGChat(
+            searcher=searcher,
+            openai_chat_client=openai_chat.client,
+            chat_model=context.openai_chat_model,
+            chat_deployment=context.openai_chat_deployment,
+        ).run_stream
+    else:
+        run_ragchat = SimpleRAGChat(
+            searcher=searcher,
+            openai_chat_client=openai_chat.client,
+            chat_model=context.openai_chat_model,
+            chat_deployment=context.openai_chat_deployment,
+        ).run_stream
+
+    result = run_ragchat(chat_request.messages, overrides=overrides)
+    return StreamingResponse(content=format_as_ndjson(result), media_type="application/x-ndjson")