diff --git a/backend/app/rag/agent.py b/backend/app/rag/agent.py index b7e91d5a..8baf8e8a 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 927bcb4a..5be9677f 100644 --- a/frontend/src/components/chat/ChatPanel.tsx +++ b/frontend/src/components/chat/ChatPanel.tsx @@ -290,50 +290,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, - ), - ); - } - } else if (event.type === "sources") { + ensureAssistantCreated(); setMessages((prev) => - prev.map((m) => - m.id === assistantId - ? { ...m, sources: event.data as SourceChunk[] } - : m, - ), + 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); @@ -401,37 +389,36 @@ export default function ChatPanel({ activeDoc, onCitationClick }: Props) { 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, sources: event.data as SourceChunk[] } - : m, - ), + prev.map((m) => (m.id === assistantId ? { ...m, thoughts: [...(m.thoughts || []), event.data as string] } : m)) ); } else if (event.type === "error") { setIsTyping(false); diff --git a/frontend/src/components/chat/MessageBubble.tsx b/frontend/src/components/chat/MessageBubble.tsx index ece6302a..ef144366 100644 --- a/frontend/src/components/chat/MessageBubble.tsx +++ b/frontend/src/components/chat/MessageBubble.tsx @@ -18,6 +18,7 @@ import { Play, Pause, GitBranch, + ChevronDown, } from "lucide-react"; import { buttonVariants } from "@/components/ui/button"; import { @@ -79,6 +80,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); @@ -325,6 +333,36 @@ export default function MessageBubble({ message }: Props) { )} + {message.thoughts && message.thoughts.length > 0 && ( +
+ + {showThoughts && ( +
+ {message.thoughts.map((thought, idx) => ( +
+ {thought} +
+ ))} +
+ )} +
+ )} +
diff --git a/frontend/src/store/chat-store.ts b/frontend/src/store/chat-store.ts index 5ed6a27c..65162ff6 100644 --- 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 {