diff --git a/backend/app/rag/agent.py b/backend/app/rag/agent.py index b7e91d5..8baf8e8 100644 --- 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,6 +254,14 @@ 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: @@ -258,7 +269,12 @@ def generate_answer_stream( 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 index b1e028a..e064316 100644 --- 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 index 580fe3a..1108631 100644 --- 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 | null>(null); const sharedTimeoutRef = useRef | null>(null); const utteranceRef = useRef(null); @@ -296,6 +303,36 @@ const handleBranch = () => { )} + {message.thoughts && message.thoughts.length > 0 && ( +
+ + {showThoughts && ( +
+ {message.thoughts.map((thought, idx) => ( +
+ {thought} +
+ ))} +
+ )} +
+ )} +
{message.content ? (