From b91852281638452b6a0bceeca61bd9bd634ebf07 Mon Sep 17 00:00:00 2001 From: Ryan Chen Date: Wed, 9 Sep 2026 23:14:08 -0400 Subject: [PATCH] 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 --- STREAMING_IMPLEMENTATION.md | 79 +++++++++++++++++++ blueprints/conversation/__init__.py | 9 +++ raggr-frontend/src/api/conversationService.ts | 1 + raggr-frontend/src/components/ChatScreen.tsx | 5 +- raggr-frontend/src/hooks/useChat.ts | 49 +++++++++++- 5 files changed, 139 insertions(+), 4 deletions(-) create mode 100644 STREAMING_IMPLEMENTATION.md diff --git a/STREAMING_IMPLEMENTATION.md b/STREAMING_IMPLEMENTATION.md new file mode 100644 index 0000000..d1f0e77 --- /dev/null +++ b/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": }`), from `on_tool_start`. +- `tool_end` — a tool finished, from `on_tool_end`. +- `content` — an incremental answer chunk (`{ "type": "content", "delta": }`), from + `on_chat_model_stream`. Chunks are accumulated server-side into `streamed`. +- `response` — the final, authoritative message (`{ "type": "response", "message": }`). + 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": }`). 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 +``` diff --git a/blueprints/conversation/__init__.py b/blueprints/conversation/__init__.py index 272900d..07f0c00 100644 --- a/blueprints/conversation/__init__.py +++ b/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, diff --git a/raggr-frontend/src/api/conversationService.ts b/raggr-frontend/src/api/conversationService.ts index ae205ba..7baab48 100644 --- a/raggr-frontend/src/api/conversationService.ts +++ b/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 }; diff --git a/raggr-frontend/src/components/ChatScreen.tsx b/raggr-frontend/src/components/ChatScreen.tsx index 9fb3b7c..d848350 100644 --- a/raggr-frontend/src/components/ChatScreen.tsx +++ b/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 ; })} - {(isLoading || messagesLoading) && } + {((isLoading && !streaming) || messagesLoading) && ( + + )} {error && (
diff --git a/raggr-frontend/src/hooks/useChat.ts b/raggr-frontend/src/hooks/useChat.ts index 13730d3..c97024b 100644 --- a/raggr-frontend/src/hooks/useChat.ts +++ b/raggr-frontend/src/hooks/useChat.ts @@ -48,8 +48,12 @@ export function useChat({ const [messagesLoading, setMessagesLoading] = useState(false); const [pendingImage, setPendingImage] = useState(null); const [error, setError] = useState(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(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") { - updateMessages((prev) => - prev.concat([{ text: event.message, speaker: "simba" }]), - ); + // 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,