Adds optional streaming support for chat requests (#532)

* Quart draft

* Fix ask and test

* Quart deploying now

* Use semantic

* Get tests working

* Revert simple

* Typing fixes

* dont use pipe

* Initial draft of streaming for quart

* Get streaming working

* Revert unneeded changes

* Update version

* Use __anext__

* Update package lock

* Update ndjson output per recommendations

* Typing changes

* Split into 2 endpoints, more tests

* Formatting fixes

* Sources moved to end

* Dont use session workaround for streaming

* Ignore timeout type issue

Author: pamelafox

.pre-commit-config.yaml

@@ -1,3 +1,4 @@
+exclude: '^tests/snapshots/'
 repos:
 -   repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v4.4.0

app/backend/app.py

@@ -1,8 +1,10 @@
 import io
+import json
 import logging
 import mimetypes
 import os
 import time
+from typing import AsyncGenerator
 
 import aiohttp
 import openai
@@ -18,6 +20,7 @@
     abort,
     current_app,
     jsonify,
+    make_response,
     request,
     send_file,
     send_from_directory,
@@ -97,12 +100,36 @@ async def chat():
         # Workaround for: https://github.com/openai/openai-python/issues/371
         async with aiohttp.ClientSession() as s:
             openai.aiosession.set(s)
-            r = await impl.run(request_json["history"], request_json.get("overrides") or {})
+            r = await impl.run_without_streaming(request_json["history"], request_json.get("overrides", {}))
         return jsonify(r)
     except Exception as e:
         logging.exception("Exception in /chat")
         return jsonify({"error": str(e)}), 500
 
+
+async def format_as_ndjson(r: AsyncGenerator[dict, None]) -> AsyncGenerator[str, None]:
+    async for event in r:
+        yield json.dumps(event, ensure_ascii=False) + "\n"
+
+@bp.route("/chat_stream", methods=["POST"])
+async def chat_stream():
+    if not request.is_json:
+        return jsonify({"error": "request must be json"}), 415
+    request_json = await request.get_json()
+    approach = request_json["approach"]
+    try:
+        impl = current_app.config[CONFIG_CHAT_APPROACHES].get(approach)
+        if not impl:
+            return jsonify({"error": "unknown approach"}), 400
+        response_generator = impl.run_with_streaming(request_json["history"], request_json.get("overrides", {}))
+        response = await make_response(format_as_ndjson(response_generator))
+        response.timeout = None # type: ignore
+        return response
+    except Exception as e:
+        logging.exception("Exception in /chat")
+        return jsonify({"error": str(e)}), 500
+
+
 @bp.before_request
 async def ensure_openai_token():
     openai_token = current_app.config[CONFIG_OPENAI_TOKEN]

app/backend/approaches/approach.py

@@ -2,13 +2,7 @@
 from typing import Any
 
 
-class ChatApproach(ABC):
-    @abstractmethod
-    async def run(self, history: list[dict], overrides: dict[str, Any]) -> Any:
-        ...
-
-
 class AskApproach(ABC):
     @abstractmethod
-    async def run(self, q: str, overrides: dict[str, Any]) -> Any:
+    async def run(self, q: str, overrides: dict[str, Any]) -> dict[str, Any]:
         ...

app/backend/approaches/chatreadretrieveread.py

@@ -1,16 +1,15 @@
-from typing import Any
+from typing import Any, AsyncGenerator
 
 import openai
 from azure.search.documents.aio import SearchClient
 from azure.search.documents.models import QueryType
 
-from approaches.approach import ChatApproach
 from core.messagebuilder import MessageBuilder
 from core.modelhelper import get_token_limit
 from text import nonewlines
 
 
-class ChatReadRetrieveReadApproach(ChatApproach):
+class ChatReadRetrieveReadApproach:
     # Chat roles
     SYSTEM = "system"
     USER = "user"
@@ -57,7 +56,7 @@ def __init__(self, search_client: SearchClient, chatgpt_deployment: str, chatgpt
         self.content_field = content_field
         self.chatgpt_token_limit = get_token_limit(chatgpt_model)
 
-    async def run(self, history: list[dict[str, str]], overrides: dict[str, Any]) -> Any:
+    async def run_until_final_call(self, history: list[dict[str, str]], overrides: dict[str, Any], should_stream: bool = False) -> tuple:
         has_text = overrides.get("retrieval_mode") in ["text", "hybrid", None]
         has_vector = overrides.get("retrieval_mode") in ["vectors", "hybrid", None]
         use_semantic_captions = True if overrides.get("semantic_captions") and has_text else False
@@ -146,20 +145,31 @@ async def run(self, history: list[dict[str, str]], overrides: dict[str, Any]) ->
             history,
             history[-1]["user"]+ "\n\nSources:\n" + content, # Model does not handle lengthy system messages well. Moving sources to latest user conversation to solve follow up questions prompt.
             max_tokens=self.chatgpt_token_limit)
-
-        chat_completion = await openai.ChatCompletion.acreate(
-            deployment_id=self.chatgpt_deployment,
-            model=self.chatgpt_model,
-            messages=messages,
-            temperature=overrides.get("temperature") or 0.7,
-            max_tokens=1024,
-            n=1)
-
-        chat_content = chat_completion.choices[0].message.content
-
         msg_to_display = '\n\n'.join([str(message) for message in messages])
 
-        return {"data_points": results, "answer": chat_content, "thoughts": f"Searched for:<br>{query_text}<br><br>Conversations:<br>" + msg_to_display.replace('\n', '<br>')}
+        extra_info = {"data_points": results, "thoughts": f"Searched for:<br>{query_text}<br><br>Conversations:<br>" + msg_to_display.replace('\n', '<br>')}
+        chat_coroutine = openai.ChatCompletion.acreate(
+                deployment_id=self.chatgpt_deployment,
+                model=self.chatgpt_model,
+                messages=messages,
+                temperature=overrides.get("temperature") or 0.7,
+                max_tokens=1024,
+                n=1,
+                stream=should_stream)
+        return (extra_info, chat_coroutine)
+
+    async def run_without_streaming(self, history: list[dict[str, str]], overrides: dict[str, Any]) -> dict[str, Any]:
+        extra_info, chat_coroutine = await self.run_until_final_call(history, overrides, should_stream=False)
+        chat_content = (await chat_coroutine).choices[0].message.content
+        extra_info["answer"] = chat_content
+        return extra_info
+
+    async def run_with_streaming(self, history: list[dict[str, str]], overrides: dict[str, Any]) -> AsyncGenerator[dict, None]:
+        extra_info, chat_coroutine = await self.run_until_final_call(history, overrides, should_stream=True)
+        yield extra_info
+        async for event in await chat_coroutine:
+            yield event
+
 
     def get_messages_from_history(self, system_prompt: str, model_id: str, history: list[dict[str, str]], user_conv: str, few_shots = [], max_tokens: int = 4096) -> list:
         message_builder = MessageBuilder(system_prompt, model_id)

app/backend/approaches/readdecomposeask.py

@@ -85,7 +85,7 @@ async def lookup(self, q: str) -> Optional[str]:
             return "\n".join([d['content'] async for d in r])
         return None
 
-    async def run(self, q: str, overrides: dict[str, Any]) -> Any:
+    async def run(self, q: str, overrides: dict[str, Any]) -> dict[str, Any]:
 
         search_results = None
         async def search_and_store(q: str) -> Any:

app/backend/approaches/readretrieveread.py

@@ -99,7 +99,7 @@ async def retrieve(self, query_text: str, overrides: dict[str, Any]) -> Any:
         content = "\n".join(results)
         return results, content
 
-    async def run(self, q: str, overrides: dict[str, Any]) -> Any:
+    async def run(self, q: str, overrides: dict[str, Any]) -> dict[str, Any]:
 
         retrieve_results = None
         async def retrieve_and_store(q: str) -> Any:

app/backend/approaches/retrievethenread.py

@@ -44,7 +44,7 @@ def __init__(self, search_client: SearchClient, openai_deployment: str, chatgpt_
         self.sourcepage_field = sourcepage_field
         self.content_field = content_field
 
-    async def run(self, q: str, overrides: dict[str, Any]) -> Any:
+    async def run(self, q: str, overrides: dict[str, Any]) -> dict[str, Any]:
         has_text = overrides.get("retrieval_mode") in ["text", "hybrid", None]
         has_vector = overrides.get("retrieval_mode") in ["vectors", "hybrid", None]
         use_semantic_captions = True if overrides.get("semantic_captions") and has_text else False

app/frontend/package-lock.json

@@ -18,7 +18,8 @@
     "dompurify": "^3.0.4",
     "react": "^18.2.0",
     "react-dom": "^18.2.0",
-    "react-router-dom": "^6.14.1"
+    "react-router-dom": "^6.14.1",
+    "ndjson-readablestream": "^1.0.6"
   },
   "devDependencies": {
     "@types/dompurify": "^3.0.2",

app/frontend/package.json

@@ -31,8 +31,9 @@ export async function askApi(options: AskRequest): Promise<AskResponse> {
     return parsedResponse;
 }
 
-export async function chatApi(options: ChatRequest): Promise<AskResponse> {
-    const response = await fetch("/chat", {
+export async function chatApi(options: ChatRequest): Promise<Response> {
+    const url = options.shouldStream ? "/chat_stream" : "/chat";
+    return await fetch(url, {
         method: "POST",
         headers: {
             "Content-Type": "application/json"
@@ -54,13 +55,6 @@ export async function chatApi(options: ChatRequest): Promise<AskResponse> {
             }
         })
     });
-
-    const parsedResponse: AskResponse = await response.json();
-    if (response.status > 299 || !response.ok) {
-        throw Error(parsedResponse.error || "Unknown error");
-    }
-
-    return parsedResponse;
 }
 
 export function getCitationFilePath(citation: string): string {