Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 18 additions & 2 deletions backend/app/rag/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -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", []):
Expand All @@ -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}")
Expand Down
89 changes: 50 additions & 39 deletions frontend/src/components/chat/ChatPanel.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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)));
Expand Down
39 changes: 38 additions & 1 deletion frontend/src/components/chat/MessageBubble.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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
Expand Down
3 changes: 2 additions & 1 deletion frontend/src/store/chat-store.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
Loading