Make /chat a generator: stream_chat in server_api.py, generator api_chat in app.py, chunked final-answer yields in streaming.py

app.py

@@ -113,8 +113,8 @@ def api_chat(
     history: list[dict],
     globe_state: dict | None,
     hf_token: gr.OAuthToken | None,
-) -> dict:
-    return server_api.run_chat(message, history, globe_state, hf_token)
+):
+    yield from server_api.stream_chat(message, history, globe_state, hf_token)
 
 
 app.mount("/assets", StaticFiles(directory=str(ASSETS_DIR)), name="assets")

ui/agent/streaming.py

@@ -5,13 +5,24 @@ from typing import Any
 
 from gradio import ChatMessage
 
+STREAM_CHUNK_SIZE = 40
+
+
+def _chunk_end_indices(text: str) -> list[int]:
+    if not text:
+        return []
+    indices = list(range(STREAM_CHUNK_SIZE, len(text), STREAM_CHUNK_SIZE))
+    if not indices or indices[-1] != len(text):
+        indices.append(len(text))
+    return indices
+
 
 def yield_streaming_string(text: str, globe_state: dict[str, Any]):
     if not text:
         yield "", globe_state
         return
-    for index in range(len(text)):
-        yield text[: index + 1], globe_state
+    for end in _chunk_end_indices(text):
+        yield text[:end], globe_state
 
 
 def yield_streaming_messages(
@@ -26,8 +37,8 @@ def yield_streaming_messages(
         return
 
     messages.append(ChatMessage(role="assistant", content=""))
-    for index in range(len(text)):
-        messages[-1] = ChatMessage(role="assistant", content=text[: index + 1])
+    for end in _chunk_end_indices(text):
+        messages[-1] = ChatMessage(role="assistant", content=text[:end])
         yield messages, globe_state
 
 

ui/server_api.py

@@ -85,12 +85,31 @@ def build_persona_prompt(persona_id: str) -> str:
     return ""
 
 
-def run_chat(
+def _merge_chat_history(
+    history: list[dict[str, Any]],
+    message: str,
+    ui_messages: list[ChatMessage],
+    assistant_text: str,
+) -> list[dict[str, Any]]:
+    updated_history = list(history)
+    updated_history.append({"role": "user", "content": message})
+    for ui_message in ui_messages:
+        updated_history.append(_chat_message_to_dict(ui_message))
+    if assistant_text and (
+        not ui_messages
+        or ui_messages[-1].role != "assistant"
+        or ui_messages[-1].metadata
+    ):
+        updated_history.append({"role": "assistant", "content": assistant_text})
+    return updated_history
+
+
+def stream_chat(
     message: str,
     history: list[dict[str, Any]],
     globe_state: dict[str, Any] | None,
     hf_token: gr.OAuthToken | None,
-) -> dict[str, Any]:
+):
     state = globe_state if globe_state else empty_globe_state()
     ui_messages: list[ChatMessage] = []
     assistant_text = ""
@@ -111,18 +130,16 @@ def run_chat(
         elif isinstance(payload, str):
             assistant_text = payload
 
-    updated_history = list(history)
-    updated_history.append({"role": "user", "content": message})
-    for ui_message in ui_messages:
-        updated_history.append(_chat_message_to_dict(ui_message))
-    if assistant_text and (
-        not ui_messages
-        or ui_messages[-1].role != "assistant"
-        or ui_messages[-1].metadata
-    ):
-        updated_history.append({"role": "assistant", "content": assistant_text})
+        yield {
+            "history": _merge_chat_history(
+                history, message, ui_messages, assistant_text
+            ),
+            "globe_state": state,
+            "done": False,
+        }
 
-    return {
-        "history": updated_history,
+    yield {
+        "history": _merge_chat_history(history, message, ui_messages, assistant_text),
         "globe_state": state,
+        "done": True,
     }