Stream assistant answers token-by-token over SSE

The /stream-query endpoint previously emitted only tool-lifecycle events plus
one final complete message, so the answer still appeared all at once. Handle
on_chat_model_stream events to emit incremental `content` deltas, accumulate
them for persistence fallback, and consume them on the frontend to build a live
assistant bubble (reconciled with the authoritative final text). Hide the
redundant loading indicator once tokens start arriving.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

STREAMING_IMPLEMENTATION.md

@@ -0,0 +1,79 @@
+# Streaming Implementation Summary
+
+## Overview
+The `/stream-query` endpoint uses Server-Sent Events (SSE) to give real-time feedback
+during AI query processing: tool-execution status **and** token-by-token streaming of the
+assistant's answer.
+
+## Backend (`blueprints/conversation/__init__.py`)
+
+### `POST /api/conversation/stream-query`
+Streams SSE frames (`data: {json}\n\n`) by iterating `main_agent.astream_events(..., version="v2")`.
+Event types emitted to the client:
+
+- `tool_start` — a tool began (`{ "type": "tool_start", "tool": <name> }`), from `on_tool_start`.
+- `tool_end` — a tool finished, from `on_tool_end`.
+- `content` — an incremental answer chunk (`{ "type": "content", "delta": <text> }`), from
+  `on_chat_model_stream`. Chunks are accumulated server-side into `streamed`.
+- `response` — the final, authoritative message (`{ "type": "response", "message": <full text> }`).
+  Sourced from `on_chain_end` (falling back to the accumulated `streamed` tokens). This is the
+  text that is persisted to the DB via `add_message_to_conversation(..., speaker="simba")`.
+- `error` — an error message.
+- The stream terminates with `data: [DONE]`.
+
+The user's message is persisted before streaming begins; the assistant message is persisted once
+the stream completes.
+
+### `POST /api/conversation/query` (legacy, non-streaming)
+Backwards-compatible endpoint that returns the complete response in one JSON payload
+(`{ "response": <text> }`). Kept for compatibility; the frontend uses `/stream-query`.
+
+## Frontend
+
+### API service (`src/api/conversationService.ts`)
+- `streamQuery(query, conversation_id, onEvent, signal?, imageKey?)` — POSTs to `/stream-query`
+  and drives the response through `_readSSEStream()`.
+- `_readSSEStream()` — reads the `ReadableStream`, splits on `\n\n`, parses `data:` frames into
+  `SSEEvent` objects, and stops on `[DONE]`.
+- `SSEEvent` union: `tool_start | tool_end | content | response | error`.
+
+### Chat state (`src/hooks/useChat.ts`)
+`sendMessage()` consumes the stream and maintains a unified `messages` array:
+
+- `tool_start` → appends a `tool` bubble (mapped to a friendly label via `TOOL_MESSAGES`) and
+  resets the streaming pointer so any text after a tool call starts a fresh bubble.
+- `content` → on the first delta, appends a new `simba` bubble and flips `streamingActiveRef` /
+  `streaming`; subsequent deltas append to that bubble's text.
+- `response` → overwrites the live `simba` bubble with the authoritative message (guaranteeing
+  parity with what was persisted), or appends it if nothing streamed (e.g. an instant answer).
+- `error` → surfaces a user-facing error.
+
+A `streaming` boolean is exposed so the UI can hide the redundant loading indicator once tokens
+start arriving.
+
+### Components
+- `ChatScreen.tsx` — renders `messages` (`ToolBubble` / `AnswerBubble` / `QuestionBubble`) and
+  shows the standalone loading-dots `AnswerBubble` only before the first token
+  (`(isLoading && !streaming) || messagesLoading`).
+- `AnswerBubble.tsx` — renders the assistant message as markdown, or animated loading dots when
+  `loading` is set.
+- `ToolBubble.tsx` — renders tool-status messages as inline badges.
+
+## User Experience
+1. User sends a message.
+2. Loading dots appear.
+3. Tool status bubbles appear as tools run (e.g. "Searching Simba's records...").
+4. The answer streams in word-by-word into a single assistant bubble.
+5. On completion the bubble is reconciled with the persisted final text.
+
+## Testing
+```bash
+# Backend
+make test                       # or: .venv/bin/python -m pytest tests/
+
+# Frontend
+cd raggr-frontend && yarn build
+
+# Full stack (manual)
+docker compose up --build
+```

blueprints/conversation/__init__.py

@@ -199,6 +199,7 @@ async def stream_query():
 
     async def event_generator():
         final_message = None
+        streamed = ""
         try:
             async for event in main_agent.astream_events(
                 payload, version="v2", config=agent_config
@@ -211,6 +212,13 @@ async def stream_query():
                 elif event_type == "on_tool_end":
                     yield f"data: {json.dumps({'type': 'tool_end', 'tool': event['name']})}\n\n"
 
+                elif event_type == "on_chat_model_stream":
+                    chunk = event.get("data", {}).get("chunk")
+                    content = getattr(chunk, "content", None)
+                    if isinstance(content, str) and content:
+                        streamed += content
+                        yield f"data: {json.dumps({'type': 'content', 'delta': content})}\n\n"
+
                 elif event_type == "on_chain_end":
                     output = event.get("data", {}).get("output")
                     if isinstance(output, dict):
@@ -224,6 +232,7 @@ async def stream_query():
         except Exception as e:
             yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n"
 
+        final_message = final_message or streamed
         if final_message:
             await add_message_to_conversation(
                 conversation=conversation,

raggr-frontend/src/api/conversationService.ts

@@ -3,6 +3,7 @@ import { userService } from "./userService";
 export type SSEEvent =
   | { type: "tool_start"; tool: string }
   | { type: "tool_end"; tool: string }
+  | { type: "content"; delta: string }
   | { type: "response"; message: string }
   | { type: "error"; message: string };
 

raggr-frontend/src/components/ChatScreen.tsx

@@ -52,6 +52,7 @@ export const ChatScreen = ({ setAuthenticated, isAdmin }: ChatScreenProps) => {
   const {
     messages,
     isLoading,
+    streaming,
     messagesLoading,
     pendingImage,
     setPendingImage,
@@ -247,7 +248,9 @@ export const ChatScreen = ({ setAuthenticated, isAdmin }: ChatScreenProps) => {
                   return <QuestionBubble key={index} text={msg.text} image_key={msg.image_key} />;
                 })}
 
-                {(isLoading || messagesLoading) && <AnswerBubble text="" loading={true} />}
+                {((isLoading && !streaming) || messagesLoading) && (
+                  <AnswerBubble text="" loading={true} />
+                )}
 
                 {error && (
                   <div className="flex justify-center message-enter">

raggr-frontend/src/hooks/useChat.ts

@@ -48,8 +48,12 @@ export function useChat({
   const [messagesLoading, setMessagesLoading] = useState(false);
   const [pendingImage, setPendingImage] = useState<File | null>(null);
   const [error, setError] = useState<string | null>(null);
+  const [streaming, setStreaming] = useState(false);
 
   const isMountedRef = useRef(true);
+  // True once tokens for the current assistant reply have started arriving, so
+  // new content appends to the live bubble instead of creating another one.
+  const streamingActiveRef = useRef(false);
   const abortControllerRef = useRef<AbortController | null>(null);
   // Set to a conversation id we just created locally, so the route-change
   // loader below doesn't wipe the optimistic message + in-flight stream.
@@ -135,6 +139,8 @@ export function useChat({
       const imageFile = pendingImage;
 
       setError(null);
+      streamingActiveRef.current = false;
+      setStreaming(false);
       updateMessages((prev) => prev.concat([{ text: query, speaker: "user" }]));
       setPendingImage(null);
       setIsLoading(true);
@@ -179,15 +185,49 @@ export function useChat({
           (event) => {
             if (!isMountedRef.current) return;
             if (event.type === "tool_start") {
+              // Any text after a tool call belongs in a fresh bubble.
+              streamingActiveRef.current = false;
               const friendly =
                 TOOL_MESSAGES[event.tool] ?? `Using ${event.tool}...`;
               updateMessages((prev) =>
                 prev.concat([{ text: friendly, speaker: "tool" }]),
               );
+            } else if (event.type === "content") {
+              if (streamingActiveRef.current) {
+                updateMessages((prev) => {
+                  const updated = [...prev];
+                  const last = updated[updated.length - 1];
+                  updated[updated.length - 1] = {
+                    ...last,
+                    text: last.text + event.delta,
+                  };
+                  return updated;
+                });
+              } else {
+                streamingActiveRef.current = true;
+                setStreaming(true);
+                updateMessages((prev) =>
+                  prev.concat([{ text: event.delta, speaker: "simba" }]),
+                );
+              }
             } else if (event.type === "response") {
+              // Reconcile the streamed bubble with the authoritative (persisted)
+              // text; if nothing streamed, append it as a new message.
+              if (streamingActiveRef.current) {
+                updateMessages((prev) => {
+                  const updated = [...prev];
+                  updated[updated.length - 1] = {
+                    ...updated[updated.length - 1],
+                    text: event.message,
+                  };
+                  return updated;
+                });
+                streamingActiveRef.current = false;
+              } else {
                 updateMessages((prev) =>
                   prev.concat([{ text: event.message, speaker: "simba" }]),
                 );
+              }
             } else if (event.type === "error") {
               console.error("Stream error:", event.message);
               if (isMountedRef.current) {
@@ -215,10 +255,12 @@ export function useChat({
       } finally {
         if (isMountedRef.current) {
           setIsLoading(false);
+          setStreaming(false);
           if (createdNew) {
             refreshConversations();
           }
         }
+        streamingActiveRef.current = false;
         abortControllerRef.current = null;
       }
     },
@@ -244,6 +286,7 @@ export function useChat({
     messages,
     setMessages: updateMessages,
     isLoading,
+    streaming,
     messagesLoading,
     pendingImage,
     setPendingImage,