feat: stream intermediate agent reasoning and token-by-token final answers

knoxiboy · knoxiboy · commit b35bda8b1642 · 2026-06-13T17:50:43.000+05:30
diff --git a/backend/app/rag/agent.py b/backend/app/rag/agent.py
@@ -233,7 +233,10 @@ def generate_answer_stream(
 
         for step in executor.stream({"input": question, "chat_history": formatted_history}):
             if "actions" in step:
-                continue
+                for action in step["actions"]:
+                    log_content = getattr(action, "log", "")
+                    if log_content:
+                        yield f"data: {json.dumps({'type': 'thought', 'data': log_content.strip()})}\n\n"
 
             elif "intermediate_steps" in step:
                 if not sources_sent and getattr(pdf_tool, "last_sources", []):
@@ -251,14 +254,27 @@ def generate_answer_stream(
                     yield f"data: {json.dumps({'type': 'sources', 'data': sources})}\n\n"
                     sources_sent = True
 
+                for agent_step in step["intermediate_steps"]:
+                    if isinstance(agent_step, tuple) and len(agent_step) >= 2:
+                        observation = agent_step[1]
+                        obs_str = str(observation)
+                        if len(obs_str) > 1000:
+                            obs_str = obs_str[:1000] + "... (truncated)"
+                        yield f"data: {json.dumps({'type': 'thought', 'data': f'Observation: {obs_str}'})}\n\n"
+
             elif "output" in step:
                 full_answer = step["output"]
                 try:
                     clean_answer = parse_agent_output(full_answer)
                 except OutputParserError as e:
                     logger.warning(f"Rejected malformed streamed LLM output: {e}")
                     clean_answer = MALFORMED_OUTPUT_MESSAGE
-                yield f"data: {json.dumps({'type': 'token', 'data': clean_answer})}\n\n"
+
+                import time
+                chunk_size = 4
+                for i in range(0, len(clean_answer), chunk_size):
+                    yield f"data: {json.dumps({'type': 'token', 'data': clean_answer[i:i+chunk_size]})}\n\n"
+                    time.sleep(0.01)
 
     except Exception as e:
         logger.error(f"Agent streaming error: {e}")
diff --git a/frontend/src/components/chat/ChatPanel.tsx b/frontend/src/components/chat/ChatPanel.tsx
@@ -275,36 +275,38 @@ export default function ChatPanel({ activeDoc, onCitationClick }: Props) {
           reject(new Error("WebSocket connection timeout"));
         }, 800);
 
+        const ensureAssistantCreated = (initialThoughts?: string[], initialSources?: SourceChunk[]) => {
+          if (!assistantCreated) {
+            assistantCreated = true;
+            setIsTyping(false);
+            const assistantMsg: ChatMsg = {
+              id: assistantId,
+              role: "assistant",
+              content: "",
+              sources: initialSources || [],
+              isStreaming: true,
+              thoughts: initialThoughts || [],
+            };
+            setMessages((prev) => [...prev, assistantMsg]);
+          }
+        };
+
         ws.onmessage = (ev) => {
           clearTimeout(connectTimeout);
           try {
             const event = JSON.parse(ev.data);
             if (event.type === "token") {
-              if (!assistantCreated) {
-                assistantCreated = true;
-                setIsTyping(false);
-
-                const assistantMsg: ChatMsg = {
-                  id: assistantId,
-                  role: "assistant",
-                  content: event.data as string,
-                  sources: [],
-                  isStreaming: true,
-                };
-
-                setMessages((prev) => [...prev, assistantMsg]);
-              } else {
-                setMessages((prev) =>
-                  prev.map((m) => (m.id === assistantId ? { ...m, content: m.content + (event.data as string) } : m))
-                );
-              }
+              ensureAssistantCreated();
+              setMessages((prev) =>
+                prev.map((m) => (m.id === assistantId ? { ...m, content: m.content + (event.data as string) } : m))
+              );
             } else if (event.type === "sources") {
+              ensureAssistantCreated(undefined, event.data as SourceChunk[]);
               setMessages((prev) => prev.map((m) => (m.id === assistantId ? { ...m, sources: event.data as SourceChunk[] } : m)));
             } else if (event.type === "thought") {
-              // Append thoughts as a temporary assistant note (optional UI handling)
-              // For simplicity, add to assistant message content in brackets
+              ensureAssistantCreated([event.data as string]);
               setMessages((prev) =>
-                prev.map((m) => (m.id === assistantId ? { ...m, content: m.content + `\n[thought] ${event.data}` } : m))
+                prev.map((m) => (m.id === assistantId ? { ...m, thoughts: [...(m.thoughts || []), event.data as string] } : m))
               );
             } else if (event.type === "error") {
               setIsTyping(false);
@@ -354,28 +356,37 @@ export default function ChatPanel({ activeDoc, onCitationClick }: Props) {
           session_id: activeSessionId,
         }, abortController.signal);
 
+        let sseAssistantCreated = false;
+        const ensureSseAssistantCreated = (initialThoughts?: string[], initialSources?: SourceChunk[]) => {
+          if (!sseAssistantCreated) {
+            sseAssistantCreated = true;
+            setIsTyping(false);
+            const assistantMsg: ChatMsg = {
+              id: assistantId,
+              role: "assistant",
+              content: "",
+              sources: initialSources || [],
+              isStreaming: true,
+              thoughts: initialThoughts || [],
+            };
+            setMessages((prev) => [...prev, assistantMsg]);
+          }
+        };
+
         for await (const event of stream) {
           if (event.type === "token") {
-            if (!assistantCreated) {
-              assistantCreated = true;
-              setIsTyping(false);
-
-              const assistantMsg: ChatMsg = {
-                id: assistantId,
-                role: "assistant",
-                content: event.data as string,
-                sources: [],
-                isStreaming: true,
-              };
-
-              setMessages((prev) => [...prev, assistantMsg]);
-            } else {
-              setMessages((prev) =>
-                prev.map((m) => (m.id === assistantId ? { ...m, content: m.content + (event.data as string) } : m))
-              );
-            }
+            ensureSseAssistantCreated();
+            setMessages((prev) =>
+              prev.map((m) => (m.id === assistantId ? { ...m, content: m.content + (event.data as string) } : m))
+            );
           } else if (event.type === "sources") {
+            ensureSseAssistantCreated(undefined, event.data as SourceChunk[]);
             setMessages((prev) => prev.map((m) => (m.id === assistantId ? { ...m, sources: event.data as SourceChunk[] } : m)));
+          } else if (event.type === "thought") {
+            ensureSseAssistantCreated([event.data as string]);
+            setMessages((prev) =>
+              prev.map((m) => (m.id === assistantId ? { ...m, thoughts: [...(m.thoughts || []), event.data as string] } : m))
+            );
           } else if (event.type === "error") {
             setIsTyping(false);
             setMessages((prev) => prev.map((m) => (m.id === assistantId ? { ...m, content: `Error: ${event.data}`, isStreaming: false } : m)));
diff --git a/frontend/src/components/chat/MessageBubble.tsx b/frontend/src/components/chat/MessageBubble.tsx
@@ -7,7 +7,7 @@ import rehypeHighlight from "rehype-highlight";
 import remarkGfm from "remark-gfm";
 import type { ChatMsg } from "@/store/chat-store";
 import { api } from "@/lib/api";
-import { Brain, User, Copy, Check, Share2, Link2, X, Play, Pause, GitBranch } from "lucide-react";
+import { Brain, User, Copy, Check, Share2, Link2, X, Play, Pause, GitBranch, ChevronDown } from "lucide-react";
 import { buttonVariants } from "@/components/ui/button";
 import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
 import { cn } from "@/lib/utils";
@@ -64,6 +64,13 @@ export default function MessageBubble({ message }: Props) {
   const [shared, setShared] = useState(false);
   const [shareFailed, setShareFailed] = useState(false);
   const [isSpeaking, setIsSpeaking] = useState(false);
+  const [showThoughts, setShowThoughts] = useState(false);
+
+  useEffect(() => {
+    if (message.isStreaming && message.thoughts && message.thoughts.length > 0) {
+      setShowThoughts(true);
+    }
+  }, [message.isStreaming, message.thoughts]);
   const copiedTimeoutRef = useRef<ReturnType<typeof setTimeout> | null>(null);
   const sharedTimeoutRef = useRef<ReturnType<typeof setTimeout> | null>(null);
   const utteranceRef = useRef<SpeechSynthesisUtterance | null>(null);
@@ -296,6 +303,36 @@ const handleBranch = () => {
               </>
             )}
 
+            {message.thoughts && message.thoughts.length > 0 && (
+              <div className="mb-3 border border-border/60 rounded-lg bg-muted/40 overflow-hidden text-xs">
+                <button
+                  type="button"
+                  onClick={() => setShowThoughts((prev) => !prev)}
+                  className="w-full flex items-center justify-between px-3 py-2 text-muted-foreground hover:text-foreground font-medium transition-colors"
+                >
+                  <div className="flex items-center gap-1.5">
+                    <Brain className="w-3.5 h-3.5 animate-pulse text-primary" />
+                    <span>{message.isStreaming ? "Thinking..." : "View reasoning steps"}</span>
+                  </div>
+                  <ChevronDown
+                    className={cn(
+                      "w-3.5 h-3.5 transition-transform duration-200",
+                      showThoughts && "transform rotate-180"
+                    )}
+                  />
+                </button>
+                {showThoughts && (
+                  <div className="px-3 pb-3 pt-1 border-t border-border/40 space-y-2 max-h-60 overflow-y-auto font-mono text-[11px] leading-relaxed text-muted-foreground">
+                    {message.thoughts.map((thought, idx) => (
+                      <div key={idx} className="whitespace-pre-wrap border-l-2 border-primary/30 pl-2">
+                        {thought}
+                      </div>
+                    ))}
+                  </div>
+                )}
+              </div>
+            )}
+
             <div className={`prose-chat ${fontSizeClass} ${message.content ? "pr-20" : ""}`}>
               {message.content ? (
                 <ReactMarkdown
diff --git a/frontend/src/store/chat-store.ts b/frontend/src/store/chat-store.ts
@@ -22,13 +22,14 @@ export interface SourceChunk {
 
 export interface ChatMsg {
   branch_id?: string;
-parent_message_id?: string;
+  parent_message_id?: string;
   id: string;
   role: "user" | "assistant";
   content: string;
   sources: SourceChunk[];
   feedback?: "up" | "down" | null;
   isStreaming?: boolean;
+  thoughts?: string[];
 }
 
 export interface ChatSession {
