From 7ea92c936680af4eeb46564a7628fa2407f6d10c Mon Sep 17 00:00:00 2001 From: Lasse Server Date: Thu, 6 Aug 2026 09:52:09 +0200 Subject: [PATCH] Stream the final chat answer instead of one blocking blob MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit packages/llm/client.py:_read_stream yielded token-by-token deltas but was dead code — nothing called generate(stream=True). The chat tool loop always blocked on the final answer generation, so a ~40s answer arrived as a single SSE event and read as a hang, even though the "card stack" already streamed tool-selection progress. The final answer is structurally just whichever tool-loop iteration returns content with no tool_calls — there's no way to know in advance which iteration that'll be, and a turn can carry short narration before a tool call in the very same response. So streaming has to be speculative: - packages/llm/client.py: _read_stream becomes StreamAccumulator, which both streams live (thinking/content/tool_call) deltas *and* reconstructs a full ChatCompletionMessage (content, tool_calls, reasoning_content) once the stream ends. Tool-call argument fragments are accumulated by index, same pattern every OpenAI-compatible streaming client needs. Downstream tool execution code needed zero changes because of this — it just consumes the reconstructed message instead of the SDK's blocking return value. Also hardened / detection against a tag split across chunk boundaries (was a per-chunk substring check; now a carry-over buffer) — latent before since nothing exercised streaming for real. - backend/services/streaming_answer.py (new): AnswerStreamFilter holds buffered content back until it's implausible to still be one-line narration (this codebase already caps narration at 150 chars elsewhere), and separately holds back any trailing bracket run that could still grow into a [src:ID]/[[src:ID]] marker — citation validity isn't known until parse_and_renumber_citations runs on the complete text, so a raw marker, hallucinated or not, must never reach the client even for a frame. run_streaming_iteration drives one tool-loop iteration through this filter and emits "answer_delta" events; if a tool_call shows up after content already cleared the gate (narration grew unusually long before the model pivoted), it emits one "answer_delta_retract" so the UI doesn't strand a stale preview. - backend/services/chat.py, mp_chat.py: thread an explicit stream_answer flag through to _run_tool_loop's single generate() call site, rather than inferring streaming from event_callback being non-None — eval logging's ConversationRecorder.wrap() always returns a non-None callable, even for the plain blocking /api/chat endpoint, so that inference would have been wrong. Only stream_chat_response (the two SSE routes) sets it; /api/chat is untouched. - Frontend: both SSE consumers (ChatPanel.tsx's card stack, MpChatPanel.tsx's simpler turn list) grow "answer_delta"/"answer_delta_retract" handling. The live preview renders as plain, React-escaped text (no dangerouslySetInnerHTML, no markdown parsing mid-stream — that's a one-shot full-document transform, unsafe on partial input) and is fully replaced by the authoritative, citation-renumbered/person-linked/ language-polished HTML once the terminal "answer" event lands. Both event types were previously-unhandled and silently ignored, so the backend and frontend halves are independently safe to deploy. Added tests/test_llm_client_stream.py and tests/test_streaming_answer.py covering the reconstruction (content, tool-call fragment accumulation, split tags, truncated-stream handling) and the filter (gate timing, citation-marker safety across chunk boundaries, the retract path, error propagation). 5 pre-existing failures in test_prompts_golden.py/ test_provenance.py are unrelated (reproduce identically on main). --- backend/routes/chat.py | 7 +- backend/services/chat.py | 39 +++- backend/services/mp_chat.py | 13 +- backend/services/streaming_answer.py | 226 ++++++++++++++++++++ frontend/src/components/ChatPanel.tsx | 79 ++++++- frontend/src/components/MpChatPanel.tsx | 29 ++- frontend/src/styles.css | 31 +++ frontend/src/types.ts | 5 +- packages/llm/__init__.py | 3 +- packages/llm/client.py | 154 +++++++++++-- tests/test_llm_client_stream.py | 153 +++++++++++++ tests/test_streaming_answer.py | 273 ++++++++++++++++++++++++ 12 files changed, 975 insertions(+), 37 deletions(-) create mode 100644 backend/services/streaming_answer.py create mode 100644 tests/test_llm_client_stream.py create mode 100644 tests/test_streaming_answer.py diff --git a/backend/routes/chat.py b/backend/routes/chat.py index 2c5aff8..bf0abde 100644 --- a/backend/routes/chat.py +++ b/backend/routes/chat.py @@ -139,7 +139,10 @@ def chat_stream_endpoint(payload: ChatRequest, request: Request) -> StreamingRes """ SSE endpoint: streams tool-call progress events followed by the final answer. Each event is a line of the form: data: \\n\\n - Event types: "tool_call", "status", "answer", "error". + Event types: "tool_call", "status", "answer_delta", "answer_delta_retract", + "answer", "error". "answer_delta" pieces are a provisional, speculative + preview of the final answer as it's generated — see + backend/services/streaming_answer.py. Using streaming avoids Cloudflare's 100-second proxy timeout for long-running queries. """ messages = [msg.model_dump() for msg in payload.messages] @@ -322,6 +325,8 @@ def mp_chat_stream_endpoint(payload: MpChatRequest) -> StreamingResponse: """ SSE streaming endpoint for chatting with an MP persona. The LLM role-plays as the specified person, grounded in their actual speeches. + Same event types as /chat/stream, including the "answer_delta"/ + "answer_delta_retract" live-preview events. """ from backend.services.mp_chat import MpChatService diff --git a/backend/services/chat.py b/backend/services/chat.py index 6b5f556..270457d 100644 --- a/backend/services/chat.py +++ b/backend/services/chat.py @@ -30,6 +30,7 @@ from backend.services.provenance import ( parse_and_renumber_citations, _SRC_PATTERN, ) +from backend.services.streaming_answer import run_streaming_iteration from backend.services.research_models import ( ResearchRequest, ResearchReport, @@ -306,6 +307,14 @@ class ChatService: Yields dicts with a "type" key: {"type": "tool_call", "tool": ""} – a tool is about to run {"type": "status", "message": ""} – generic progress note + {"type": "answer_delta", "text": ""} – a live piece of the final + answer, streamed speculatively as it's generated. Provisional: the + terminal "answer" event's text may still differ (citation + renumbering, person-link injection, attribution fixes and the + language pass all run once, after the fact, on the complete text). + {"type": "answer_delta_retract"} – the iteration that looked like the + final answer turned out to end in a tool call after all; discard + any "answer_delta" text shown for it. {"type": "answer", "answer": "...", "sources": [...], ...} – final answer {"type": "error", "message": ""} – unhandled exception """ @@ -321,7 +330,7 @@ class ChatService: result = self.get_chat_response( messages, top_k=top_k, focus_ids=focus_ids, event_callback=emit, provider_override=provider_override, use_editor=use_editor, - quick=quick, session_id=session_id, + quick=quick, session_id=session_id, stream_answer=True, ) event_queue.put({"type": "answer", **result}) except Exception as exc: @@ -351,6 +360,7 @@ class ChatService: use_editor: bool = False, quick: bool = False, session_id: Optional[str] = None, + stream_answer: bool = False, ) -> ChatResponse: """ Public entry point. If the first user message starts with "TEST ", @@ -358,13 +368,20 @@ class ChatService: (messages, events, tool calls/results, final answer, timings, errors) is recorded to the Postgres eval_conversations table. Normal conversations are never stored. + + stream_answer: when True, the tool loop's final (no-tool-call) answer + generation is streamed live as "answer_delta" events via + event_callback, instead of arriving as one blocking call. Only + stream_chat_response sets this — the plain /api/chat entry point + never does, so it is unaffected regardless of event_callback (which + eval logging below can make non-None even for non-streaming calls). """ messages, is_eval = detect_and_strip_test_prefix(list(messages)) if not is_eval: return self._get_chat_response_impl( messages, top_k=top_k, focus_ids=focus_ids, event_callback=event_callback, provider_override=provider_override, - use_editor=use_editor, quick=quick, + use_editor=use_editor, quick=quick, stream_answer=stream_answer, ) recorder = ConversationRecorder( session_id=session_id, @@ -384,7 +401,7 @@ class ChatService: messages, top_k=top_k, focus_ids=focus_ids, event_callback=recorder.wrap(event_callback), provider_override=provider_override, use_editor=use_editor, - quick=quick, + quick=quick, stream_answer=stream_answer, ) except Exception as exc: recorder.finish(error=exc) @@ -401,6 +418,7 @@ class ChatService: provider_override=None, use_editor: bool = False, quick: bool = False, + stream_answer: bool = False, ) -> ChatResponse: """ Generate a reply while allowing the assistant to call registered tools. @@ -514,6 +532,7 @@ class ChatService: communicator_llm=communicator_llm, supports_thinking=supports_thinking, sent_insights=sent_insights, + stream_answer=stream_answer, ) answer_text = ( response_message.final_answer @@ -1329,6 +1348,7 @@ class ChatService: communicator_llm=None, supports_thinking: bool = True, sent_insights: Optional[List[str]] = None, + stream_answer: bool = False, ) -> Tuple[FinalAnswer, List[Dict[str, Any]], List[str]]: """ Repeatedly call the smart LLM, executing tool calls as needed, until a @@ -1336,6 +1356,12 @@ class ChatService: Long tool results are compressed by the fast model before being appended to the message history, keeping the smart model's context lean. + + stream_answer: when True (and event_callback is set), each iteration's + generate() call streams instead of blocking, and content that clears + AnswerStreamFilter's confidence gate is speculatively forwarded live as + "answer_delta" events before we know for certain this iteration has no + tool_calls. See backend/services/streaming_answer.py. """ _smart = smart_llm or self.smart_llm _fast = fast_llm or self.fast_llm @@ -1411,7 +1437,12 @@ class ChatService: gen_kwargs = {"messages": current_messages, "think": think_now, "auto_execute_tools": False} if getattr(self, "tools", None): gen_kwargs["tools"] = self.tools - response: ChatCompletionMessage = _smart.generate(**gen_kwargs) + if stream_answer and event_callback is not None: + response: ChatCompletionMessage = run_streaming_iteration( + _smart, gen_kwargs, event_callback, iteration=i + ) + else: + response: ChatCompletionMessage = _smart.generate(**gen_kwargs) if isinstance(response, str): # _llm swallows API exceptions and returns a plain string error message. diff --git a/backend/services/mp_chat.py b/backend/services/mp_chat.py index 3cb1a74..9661e3a 100644 --- a/backend/services/mp_chat.py +++ b/backend/services/mp_chat.py @@ -30,6 +30,7 @@ from backend.services.provenance import ( ProvenanceRegistry, parse_and_renumber_citations, ) +from backend.services.streaming_answer import run_streaming_iteration ChatMessage = Dict[str, Any] ChatSource = Dict[str, Any] @@ -255,7 +256,7 @@ class MpChatService: def run() -> None: try: - result = self._get_chat_response(messages, event_callback=emit) + result = self._get_chat_response(messages, event_callback=emit, stream_answer=True) event_queue.put({"type": "answer", **result}) except Exception as exc: import traceback @@ -277,6 +278,7 @@ class MpChatService: self, messages: Sequence[ChatMessage], event_callback: Optional[Callable[[Dict[str, Any]], None]] = None, + stream_answer: bool = False, ) -> ChatResponse: name = self.person.get("name") or "" first_name = self.person.get("first_name") or name.split()[0] if name else "" @@ -317,6 +319,7 @@ class MpChatService: user_question=enriched_question, event_callback=event_callback, registry=registry, + stream_answer=stream_answer, ) answer_text = ( @@ -360,6 +363,7 @@ class MpChatService: user_question: str = "", event_callback: Optional[Callable[[Dict[str, Any]], None]] = None, registry: Optional[ProvenanceRegistry] = None, + stream_answer: bool = False, ) -> Tuple[FinalAnswer, List[ChatMessage]]: current_messages: List[ChatMessage] = list(messages) @@ -374,7 +378,12 @@ class MpChatService: gen_kwargs = {"messages": current_messages, "think": think_now, "auto_execute_tools": False} if getattr(self, "tools", None): gen_kwargs["tools"] = self.tools - response: ChatCompletionMessage = self.smart_llm.generate(**gen_kwargs) + if stream_answer and event_callback is not None: + response: ChatCompletionMessage = run_streaming_iteration( + self.smart_llm, gen_kwargs, event_callback, iteration=i + ) + else: + response: ChatCompletionMessage = self.smart_llm.generate(**gen_kwargs) tool_calls = getattr(response, "tool_calls", None) diff --git a/backend/services/streaming_answer.py b/backend/services/streaming_answer.py new file mode 100644 index 0000000..37869e1 --- /dev/null +++ b/backend/services/streaming_answer.py @@ -0,0 +1,226 @@ +"""Speculative live-streaming of the tool loop's final answer. + +The ReAct tool loop (`ChatService`/`MpChatService`'s ``_run_tool_loop``) can't know, +before a ``generate()`` call returns, whether that call will end in ``tool_calls`` +(the loop continues) or in bare ``content`` (this iteration *is* the final answer) +— both are legitimate outcomes of the same call, and a turn can even carry short +narration content before a tool call. So there is no way to decide up front +whether a given iteration is worth streaming live to the user. + +``AnswerStreamFilter`` resolves this by buffering: it only forwards content once +it is implausible that the buffered text is a one-line narration blurb ("Jag +söker nu efter..."), and ``run_streaming_iteration`` uses it to speculatively +turn content deltas into ``answer_delta`` SSE events — retracting them with a +single ``answer_delta_retract`` event in the rare case a tool call shows up +after all. + +Citation markers (``[src:ID]``) are never forwarded raw: their validity isn't +known until ``parse_and_renumber_citations`` runs on the complete text (it can +drop hallucinated IDs), so a marker must never reach the client, even for a +frame. This module only ever *strips* markers from the live preview — it never +renumbers them; that still happens exactly once, on the complete answer, same +as today. +""" +from __future__ import annotations + +from typing import Any, Callable, Dict, Optional + +from packages.llm import ChatCompletionMessage +from backend.services.event_logger import log_error, log_event +from backend.services.provenance import _SRC_PATTERN + +EventCallback = Callable[[Dict[str, Any]], None] + + +def strip_citation_markers(text: str) -> str: + """Remove complete ``[src:ID]``/``[[src:ID]]`` tags from a text segment. + + Used only on the live preview stream. The full raw text (markers intact) + keeps accumulating server-side, untouched, for the one real + ``parse_and_renumber_citations`` pass on the complete answer. + """ + return _SRC_PATTERN.sub("", text) + + +class AnswerStreamFilter: + """Buffers streamed content until it's confidently prose, not narration. + + One instance per tool-loop iteration. + + ``feed()`` is called with each raw content chunk as it arrives; it returns + whatever is now safe to forward live (or ``None``), holding back: + - text below the confidence gate — narration blurbs in this codebase are + capped at 150 chars (see ``mp_chat.py``'s status-text truncation); the + gate sits well clear of that, so a false-positive on unusually long + narration is very unlikely, while a real multi-paragraph answer still + clears it almost immediately; + - a trailing bracket run that could still grow into a ``[src:...]``/ + ``[[src:...]]`` marker. + Complete markers found in a releasable segment are stripped outright + (never rendered, never renumbered live). + """ + + def __init__(self, min_chars: int = 400, soft_min_chars: int = 200) -> None: + self._min_chars = min_chars + self._soft_min_chars = soft_min_chars + self._pending = "" # raw text not yet past the confidence gate + self._buffer = "" # gate-cleared text not yet forwarded (bracket-safety only) + self._armed = False + self.flushed_any = False + + def feed(self, text: str) -> Optional[str]: + """Add a raw content chunk; return text safe to forward live now, or None.""" + if not text: + return None + if not self._armed: + self._pending += text + if not self._gate_passes(self._pending): + return None + self._armed = True + self._buffer, self._pending = self._pending, "" + log_event("answer_stream_gate_armed", buffered_chars=len(self._buffer)) + else: + self._buffer += text + return self._release_safe_segment() + + def flush_remainder(self) -> Optional[str]: + """Release whatever's still held back. + + Call once, after the final chunk of an iteration confirmed to be the + answer (no ``tool_calls``) — safe now that the full text is known, so + the confidence gate no longer applies. Still strips citation markers. + + A bracket run left unclosed at this point (e.g. generation was cut + off mid-tag by a token limit) can never complete now — it's dropped + outright rather than ever being released raw, same principle as the + mid-stream case, just with "wait for more" replaced by "there is no + more". + """ + remainder = self._pending + self._buffer + self._pending = "" + self._buffer = "" + if not remainder: + return None + remainder = remainder[: _last_unmatched_bracket_start(remainder)] + if not remainder: + return None + piece = strip_citation_markers(remainder) + if piece: + self.flushed_any = True + return piece or None + + def _gate_passes(self, text: str) -> bool: + if len(text) >= self._min_chars: + return True + if len(text) >= self._soft_min_chars: + # Already into a second sentence/clause — a narration blurb is + # always a single short sentence, so more text after the first + # sentence boundary is a strong "this is the real answer" signal. + boundary = _first_sentence_boundary(text) + if boundary is not None and boundary < len(text.rstrip()) - 1: + return True + return False + + def _release_safe_segment(self) -> Optional[str]: + """Release ``self._buffer`` up to the start of any trailing, still-open + bracket run — a citation marker's validity depends on characters that + may not have arrived yet, so it (and its double-bracket partner, if + any) is held back until it either completes or is provably not one.""" + safe_end = _last_unmatched_bracket_start(self._buffer) + segment, self._buffer = self._buffer[:safe_end], self._buffer[safe_end:] + if not segment: + return None + piece = strip_citation_markers(segment) + if piece: + self.flushed_any = True + return piece or None + + +def _first_sentence_boundary(text: str) -> Optional[int]: + for i, ch in enumerate(text): + if ch in ".!?" and i + 1 < len(text) and text[i + 1] in " \n": + return i + return None + + +def _last_unmatched_bracket_start(text: str) -> int: + """Index from which the text might still be an incomplete bracket marker. + + Compares the position of the last ``[`` against the last ``]``: if the + last ``[`` comes after the last ``]`` (or there's no ``]`` at all), that + bracket run is still open — hold from there. Also grabs an immediately + preceding ``[`` so a ``[[src:...`` double-bracket run isn't split, leaving + an orphaned single ``[`` on screen while its partner is held back. + """ + last_open = text.rfind("[") + if last_open == -1: + return len(text) + last_close = text.rfind("]") + if last_close > last_open: + return len(text) + cutoff = last_open + if cutoff > 0 and text[cutoff - 1] == "[": + cutoff -= 1 + return cutoff + + +def run_streaming_iteration( + llm, + gen_kwargs: Dict[str, Any], + event_callback: Optional[EventCallback], + iteration: Optional[int] = None, +) -> ChatCompletionMessage: + """Run one tool-loop iteration with ``stream=True``, speculatively + forwarding the final answer live as ``answer_delta`` SSE events. + + Returns a ``ChatCompletionMessage`` shape-compatible with what a blocking + ``generate()`` call returns — the tool loop's existing branching logic + (tool execution, citation-hallucination retry, empty-response handling) + needs no changes to consume it; it's a drop-in replacement for the + blocking call at the single call site in ``_run_tool_loop``. + + Raises on API/stream failure — mirrors the blocking path's error-string + contract (``isinstance(response, str)``) being converted to a raised + exception, so a mid-stream failure surfaces the same way: as an SSE + ``error`` event, via ``stream_chat_response``'s existing exception + handling. + """ + accumulator = llm.generate(**gen_kwargs, stream=True) + if isinstance(accumulator, str): + # generate() caught the error before the stream ever started. + exc = RuntimeError(f"LLM API error: {accumulator}") + log_error("llm_api_failure", exc, model=getattr(llm, "model", None), iteration=iteration) + raise exc + + answer_filter = AnswerStreamFilter() + retracted = False + try: + for kind, text in accumulator: + if kind == "content": + piece = answer_filter.feed(text) + if piece and event_callback: + event_callback({"type": "answer_delta", "text": piece}) + elif kind == "tool_call": + # The model has started emitting a tool call — whatever content + # came before was narration, not the final answer. If any of it + # already cleared the confidence gate and reached the client, + # tell it to discard the speculative preview. + if answer_filter.flushed_any and not retracted and event_callback: + event_callback({"type": "answer_delta_retract"}) + retracted = True + # "thinking" deltas aren't surfaced to the SSE consumer today, + # matching the blocking path (reasoning_content is only logged). + except Exception as exc: + wrapped = RuntimeError(f"LLM API error: {exc}") + log_error("llm_api_failure", wrapped, model=getattr(llm, "model", None), iteration=iteration) + raise wrapped from exc + + message = accumulator.message + if not message.tool_calls: + # Confirmed: this iteration is the final answer. Release anything + # still held back — the confidence gate no longer matters once the + # full text is known. + piece = answer_filter.flush_remainder() + if piece and event_callback: + event_callback({"type": "answer_delta", "text": piece}) + return message diff --git a/frontend/src/components/ChatPanel.tsx b/frontend/src/components/ChatPanel.tsx index 19b5280..9044c75 100644 --- a/frontend/src/components/ChatPanel.tsx +++ b/frontend/src/components/ChatPanel.tsx @@ -197,6 +197,19 @@ const AnswerText = ({ html }: { html: string }) => { ); }; +/** + * Renders the live, speculative answer preview while it streams in. + * Plain JSX text child — not dangerouslySetInnerHTML — so React's default + * escaping is a second line of defense against any citation-marker fragment + * the backend filter missed; no markdown is parsed mid-stream (that only + * happens once, on the complete answer, via AnswerText above). + */ +const StreamingAnswerText = ({ text }: { text: string }) => ( +
+ {text} +
+); + const ResearchCardView = ({ card, isActive, @@ -268,6 +281,11 @@ const ResearchCardView = ({ ); } + // Live streaming-answer preview — provisional, growing prose text. + if (card.isStreaming) { + return ; + } + // Thinking card — show spinner if active if (isActive) { return ( @@ -403,9 +421,13 @@ export const ChatPanel = forwardRef(function ChatPanel( // ── Helpers for card state machine ────────────────────────────────────── - /** True if the latest (index 0) card is still "thinking" (no result, no answer). */ + /** True if the latest (index 0) card is still "thinking" (no result, no answer, not streaming). */ const latestIsThinking = (cards: ResearchCard[]) => - cards.length > 0 && cards[0].result === undefined && !cards[0].isAnswer; + cards.length > 0 && cards[0].result === undefined && !cards[0].isAnswer && !cards[0].isStreaming; + + /** True if the latest (index 0) card is a live, speculative answer preview. */ + const latestIsStreamingAnswer = (cards: ResearchCard[]) => + cards.length > 0 && !!cards[0].isStreaming; /** Create a new thinking card and prepend it. */ const prependCard = (cards: ResearchCard[], message: string): ResearchCard[] => @@ -417,6 +439,9 @@ export const ChatPanel = forwardRef(function ChatPanel( * result card looking like the final answer while the LLM decides its * next step. If the current top card is already an empty placeholder * (leftover from a prior result), drop it so we don't stack placeholders. + * Never drops a live streaming-answer preview, even though it also has an + * empty `message` — the shadow communicator runs on its own thread and + * its cards can arrive while a later iteration's answer is mid-stream. */ const prependResultCard = ( cards: ResearchCard[], @@ -426,6 +451,7 @@ export const ChatPanel = forwardRef(function ChatPanel( const base = cards.length > 0 && cards[0].result === undefined && !cards[0].isAnswer + && !cards[0].isStreaming && !cards[0].message ? cards.slice(1) : cards; @@ -546,6 +572,41 @@ export const ChatPanel = forwardRef(function ChatPanel( }; setResearchCards(prev => prependResultCard(prev, result, "")); + } else if (event.type === "answer_delta") { + // Live, speculative preview of the final answer — provisional text, + // superseded by the fully-processed HTML on the terminal "answer" event. + const chunk: string = event.text ?? ""; + if (!chunk) continue; + setResearchCards(prev => { + if (latestIsStreamingAnswer(prev)) { + return [ + { ...prev[0], streamingText: (prev[0].streamingText ?? "") + chunk }, + ...prev.slice(1), + ]; + } + if (latestIsThinking(prev)) { + return [ + { ...prev[0], isStreaming: true, streamingText: chunk }, + ...prev.slice(1), + ]; + } + return [ + { id: newCardId(), message: "", isAnswer: false, isStreaming: true, streamingText: chunk }, + ...prev, + ]; + }); + + } else if (event.type === "answer_delta_retract") { + // The iteration that looked like the final answer turned out to end + // in a tool call after all — revert to a plain (now-empty) thinking + // card. The tool loop's own "status"/"tool_call" event for this same + // iteration follows immediately and will populate its message normally. + setResearchCards(prev => + latestIsStreamingAnswer(prev) + ? [{ ...prev[0], isStreaming: false, streamingText: undefined, message: "" }, ...prev.slice(1)] + : prev + ); + } else if (event.type === "answer") { return event as ChatResponse; @@ -591,6 +652,7 @@ export const ChatPanel = forwardRef(function ChatPanel( // Merge the final answer into the card stack: // - If there are prior research cards, upgrade the latest thinking card + // (or live streaming-answer preview) to the answer // (or create a new card if the latest already has a result). // - If there are no prior cards (direct answer, no tools), leave the // card array empty and render the answer text directly. @@ -598,9 +660,16 @@ export const ChatPanel = forwardRef(function ChatPanel( let finalCards: ResearchCard[]; if (prior.length === 0) { finalCards = []; - } else if (latestIsThinking(prior)) { - // Upgrade latest thinking card to answer - finalCards = [{ ...prior[0], isAnswer: true, answerHtml }, ...prior.slice(1)]; + } else if (latestIsThinking(prior) || latestIsStreamingAnswer(prior)) { + // Upgrade latest thinking/streaming card to the answer. Whatever + // speculative text was streamed is discarded in favor of the + // authoritative, fully-processed HTML — citation renumbering, + // person-link injection, attribution fixes and the language + // pass can all still have reworded it since the last delta. + finalCards = [ + { ...prior[0], isAnswer: true, answerHtml, isStreaming: false, streamingText: undefined }, + ...prior.slice(1), + ]; } else { // Latest already has a result → prepend a new answer card finalCards = [{ id: newCardId(), message: "", isAnswer: true, answerHtml }, ...prior]; diff --git a/frontend/src/components/MpChatPanel.tsx b/frontend/src/components/MpChatPanel.tsx index 615884c..e7288fc 100644 --- a/frontend/src/components/MpChatPanel.tsx +++ b/frontend/src/components/MpChatPanel.tsx @@ -236,6 +236,9 @@ export function MpChatPanel({ person, initialTalkId, sessionId }: Props) { const [messages, setMessages] = useState([]); const [input, setInput] = useState(""); const [isPending, setIsPending] = useState(false); + // Live, speculative preview of the answer currently streaming in — provisional, + // superseded by the fully-processed answerHtml once the "answer" event lands. + const [streamingText, setStreamingText] = useState(""); const [shareToast, setShareToast] = useState<"copying" | "copied" | "error" | null>(null); const messagesEndRef = useRef(null); const textareaRef = useRef(null); @@ -339,6 +342,7 @@ export function MpChatPanel({ person, initialTalkId, sessionId }: Props) { setMessages(nextMessages); setTurns(prev => [...prev, { id: turnId, question: trimmed, status: "pending" }]); setIsPending(true); + setStreamingText(""); setInput(""); if (textareaRef.current) textareaRef.current.style.height = "auto"; @@ -383,9 +387,17 @@ export function MpChatPanel({ person, initialTalkId, sessionId }: Props) { if (!line.startsWith("data: ")) continue; const event = JSON.parse(line.slice(6)); - if (event.type === "answer") { + if (event.type === "answer_delta") { + const chunk: string = event.text ?? ""; + if (chunk) setStreamingText(prev => prev + chunk); + } else if (event.type === "answer_delta_retract") { + // The iteration that looked like the final answer turned out to + // end in a tool call after all — discard the speculative preview. + setStreamingText(""); + } else if (event.type === "answer") { const sources: ChatSource[] = event.sources ?? []; const answerHtml = convertMarkdownToHtml(event.answer ?? "", sources); + setStreamingText(""); setTurns(prev => prev.map(t => t.id === turnId @@ -414,6 +426,7 @@ export function MpChatPanel({ person, initialTalkId, sessionId }: Props) { } finally { clearTimeout(timeoutId); setIsPending(false); + setStreamingText(""); } }, [isPending, messages, person.person_id, initialTalkId, providerOverride]); @@ -510,9 +523,17 @@ export function MpChatPanel({ person, initialTalkId, sessionId }: Props) { alt="" onError={(e) => { (e.currentTarget as HTMLImageElement).style.display = "none"; }} /> -
- -
+ {streamingText ? ( + // Live, speculative preview — plain text, no markdown yet; + // replaced by the real MpBubble once "answer" arrives. +
+ {streamingText} +
+ ) : ( +
+ +
+ )} )} diff --git a/frontend/src/styles.css b/frontend/src/styles.css index 7317096..95df5b2 100644 --- a/frontend/src/styles.css +++ b/frontend/src/styles.css @@ -878,6 +878,23 @@ form[aria-busy="true"] .search-button[disabled] { overflow-x: auto; } +/* Live, speculative answer preview — plain text while it streams in, before + the fully-processed HTML (citations, person links) replaces it. */ +.chat-view__answerText--streaming { + opacity: 0.85; +} + +.chat-view__answerText--streaming::after { + content: "▍"; + display: inline-block; + margin-left: 1px; + animation: chat-stream-cursor-blink 1s step-end infinite; +} + +@keyframes chat-stream-cursor-blink { + 50% { opacity: 0; } +} + .chat-view__answerText table { width: 100%; margin: 1rem 0; @@ -2793,6 +2810,20 @@ a.rc-insight__footnote:hover { padding: 0.25rem 0; } +/* Live, speculative answer preview shown while the MP's reply streams in. */ +.mp-chat__typing-preview { + flex: 1; + min-width: 0; + opacity: 0.85; +} + +.mp-chat__typing-preview::after { + content: "▍"; + display: inline-block; + margin-left: 1px; + animation: chat-stream-cursor-blink 1s step-end infinite; +} + .mp-chat__error { color: #c0392b; font-size: 0.88rem; diff --git a/frontend/src/types.ts b/frontend/src/types.ts index a89330a..9b42d8e 100644 --- a/frontend/src/types.ts +++ b/frontend/src/types.ts @@ -227,7 +227,8 @@ export type LiveCard = LiveSearchCard | LiveStatsCard | LiveInsightCard; /** * One research step shown to the user during (and after) LLM research. * Starts as a "thinking" card with just a message, then gets upgraded to - * a result card (search/stats) or the final answer card. + * a result card (search/stats), a live streaming-answer preview, or the + * final answer card. */ export type ResearchCard = { id: string; @@ -235,6 +236,8 @@ export type ResearchCard = { result?: LiveCard; // upgraded to this when surface_results fires isAnswer: boolean; // upgraded to true when final answer arrives answerHtml?: string; // rendered answer HTML (when isAnswer=true) + isStreaming?: boolean; // true while this card is a live, speculative answer preview + streamingText?: string; // accumulated plain-text preview (not markdown-rendered yet) }; export type PersonRef = { diff --git a/packages/llm/__init__.py b/packages/llm/__init__.py index 3338c64..01ae306 100644 --- a/packages/llm/__init__.py +++ b/packages/llm/__init__.py @@ -17,7 +17,7 @@ to a private server at import time and so could not be run by anyone else. llm = LLM(base_url=..., model=..., tools=get_tools()) reply = llm.generate(messages=[{"role": "user", "content": "..."}]) """ -from .client import LLM, ChatCompletionMessage +from .client import LLM, ChatCompletionMessage, StreamAccumulator from .config import LLMConfig from .tools import ( TOOL_REGISTRY, @@ -31,6 +31,7 @@ __all__ = [ "LLM", "LLMConfig", "ChatCompletionMessage", + "StreamAccumulator", "register_tool", "get_tools", "execute_tool", diff --git a/packages/llm/client.py b/packages/llm/client.py index 4bfab88..b5da77c 100644 --- a/packages/llm/client.py +++ b/packages/llm/client.py @@ -29,6 +29,10 @@ from openai.types.chat.chat_completion import Choice from openai.types.chat.chat_completion_message import ( ChatCompletionMessage as _OpenAIChatCompletionMessage, ) +from openai.types.chat.chat_completion_message_tool_call import ( + ChatCompletionMessageToolCall, + Function as ToolCallFunction, +) from pydantic import BaseModel from .tools import execute_tool, parse_function_call_arguments @@ -233,7 +237,7 @@ class LLM: request["tools"] = tools_to_use if stream: request["stream"] = True - return self._read_stream(self.client.chat.completions.create(**request)) + return StreamAccumulator(self.client.chat.completions.create(**request)) message = self._create(**request).choices[0].message @@ -314,38 +318,150 @@ class LLM: content = json.dumps({"error": str(exc)}, ensure_ascii=False) self.messages.append({"role": "tool", "name": name or "unknown", "content": content}) - # -- streaming ------------------------------------------------------------ - def _read_stream(self, response) -> Generator[Tuple[str, str], None, None]: - """Yield ``("thinking" | "content", text)`` pairs as they arrive. +# -- streaming ------------------------------------------------------------ - Reasoning arrives either in ``reasoning_content`` or inline as ```` - blocks; both are surfaced as "thinking" so callers can render or drop them. - """ - in_think_block = False - for chunk in response: + +class StreamAccumulator: + """Wraps a raw streaming ``ChatCompletion`` response. + + Iterate it for ``("thinking" | "content" | "tool_call", text)`` events as + they arrive — ``text`` is ``None`` for a ``"tool_call"`` event, which just + signals that the model has started emitting a tool call (the payload + itself is accumulated internally; read it from ``.message`` once the + iterator is exhausted). It fires once per stream, on the first tool-call + delta seen, which is what a caller speculatively streaming ``content`` + live needs to know: the instant tool-call deltas start arriving, whatever + ``content`` came before was narration, not a final answer. + + Once the iterator is exhausted, ``.message`` returns a reconstructed + :class:`ChatCompletionMessage` — the same shape a blocking + ``generate()`` call returns (``content``, ``tool_calls``, + ``reasoning_content``) — so callers can treat a streamed and a blocking + call identically once the stream is done. + """ + + # Neither nor can be split into more pieces than their + # own length, so this is the longest prefix of either tag we might need + # to hold back across a chunk boundary while waiting to see the rest. + _MAX_TAG_LEN = max(len(""), len("")) + + def __init__(self, response) -> None: + self._response = response + self._exhausted = False + self._content_parts: List[str] = [] + self._reasoning_parts: List[str] = [] + self._tool_calls: Dict[int, Dict[str, Any]] = {} + self._pending = "" # carry-over buffer, tag-boundary-safe + self._in_think_block = False + + def __iter__(self) -> Generator[Tuple[str, Optional[str]], None, None]: + for chunk in self._response: if not chunk.choices: continue delta = chunk.choices[0].delta reasoning = getattr(delta, "reasoning_content", None) if reasoning: + self._reasoning_parts.append(reasoning) yield "thinking", reasoning + tool_call_deltas = getattr(delta, "tool_calls", None) + if tool_call_deltas: + is_first_sighting = not self._tool_calls + for tc in tool_call_deltas: + self._accumulate_tool_call(tc) + if is_first_sighting: + yield "tool_call", None + text = getattr(delta, "content", None) - if not text: + if text: + yield from self._feed_content(text) + + yield from self._flush_pending() + self._exhausted = True + + def _accumulate_tool_call(self, tc) -> None: + entry = self._tool_calls.setdefault(tc.index, {"id": None, "name": None, "arguments": ""}) + if getattr(tc, "id", None): + entry["id"] = tc.id + fn = getattr(tc, "function", None) + if fn is not None: + if getattr(fn, "name", None): + entry["name"] = fn.name + if getattr(fn, "arguments", None): + entry["arguments"] += fn.arguments + + def _feed_content(self, text: str) -> Generator[Tuple[str, str], None, None]: + self._pending += text + while True: + piece = self._extract_safe_piece() + if piece is None: + break + kind, safe_text = piece + if not safe_text: continue + (self._reasoning_parts if kind == "thinking" else self._content_parts).append(safe_text) + yield kind, safe_text - if "" in text: - in_think_block = True - text = text.split("", 1)[0] - if "" in text: - in_think_block = False - text = text.split("", 1)[1] + def _extract_safe_piece(self) -> Optional[Tuple[str, str]]: + """Pull one provably-safe ``(kind, text)`` piece off ``self._pending``. - if not text: - continue - yield ("thinking" if in_think_block else "content"), text + Returns ``None`` if what remains might still be a partial ````/ + ```` tag — i.e. there's nothing safe to release yet. + """ + tag = "" if self._in_think_block else "" + idx = self._pending.find(tag) + if idx != -1: + before, after = self._pending[:idx], self._pending[idx + len(tag):] + self._pending = after + kind = "thinking" if self._in_think_block else "content" + self._in_think_block = not self._in_think_block + return kind, before + + # No full tag in the buffer yet — release everything except a tail + # that could still grow into one on the next chunk. + safe_len = len(self._pending) - (self._MAX_TAG_LEN - 1) + if safe_len <= 0: + return None + kind = "thinking" if self._in_think_block else "content" + safe_text, self._pending = self._pending[:safe_len], self._pending[safe_len:] + return kind, safe_text + + def _flush_pending(self) -> Generator[Tuple[str, str], None, None]: + if self._pending: + # A still-open think block at EOF means a truncated/malformed + # stream — surface the remainder as "thinking" so it can never + # leak as prose, and let normal error handling deal with the + # truncation itself. + kind = "thinking" if self._in_think_block else "content" + (self._reasoning_parts if kind == "thinking" else self._content_parts).append(self._pending) + yield kind, self._pending + self._pending = "" + + @property + def message(self) -> ChatCompletionMessage: + if not self._exhausted: + raise RuntimeError( + "StreamAccumulator.message read before the stream was exhausted — " + "iterate the accumulator fully first." + ) + content = _strip_think("".join(self._content_parts)) if self._content_parts else None + tool_calls = None + if self._tool_calls: + tool_calls = [ + ChatCompletionMessageToolCall( + id=entry["id"] or f"call_{index}", + type="function", + function=ToolCallFunction(name=entry["name"] or "", arguments=entry["arguments"]), + ) + for index, entry in sorted(self._tool_calls.items()) + ] + message = ChatCompletionMessage.model_construct( + role="assistant", content=content, tool_calls=tool_calls + ) + message.reasoning_content = "".join(self._reasoning_parts) or None + return message # -- module helpers ----------------------------------------------------------- diff --git a/tests/test_llm_client_stream.py b/tests/test_llm_client_stream.py new file mode 100644 index 0000000..62c9d19 --- /dev/null +++ b/tests/test_llm_client_stream.py @@ -0,0 +1,153 @@ +"""Tests for StreamAccumulator: reconstructing a full ChatCompletionMessage +from a raw streaming ChatCompletion response, chunk by chunk.""" + +from types import SimpleNamespace + +import pytest + +from packages.llm.client import StreamAccumulator + + +def _chunk(content=None, reasoning=None, tool_calls=None): + delta = SimpleNamespace(content=content, reasoning_content=reasoning, tool_calls=tool_calls) + return SimpleNamespace(choices=[SimpleNamespace(delta=delta)]) + + +def _empty_chunk(): + """A chunk with no choices at all — providers send these sometimes.""" + return SimpleNamespace(choices=[]) + + +def _tc(index, id=None, name=None, arguments=None): + fn = SimpleNamespace(name=name, arguments=arguments) + return SimpleNamespace(index=index, id=id, function=fn) + + +class TestContentOnly: + def test_joins_content_across_chunks(self): + # Individual events may be split differently than the input chunks — + # StreamAccumulator holds back a tail that could still be a partial + # tag (see TestThinkTagHandling) — so only the joined result + # and .message are part of the contract, not the exact chunking. + acc = StreamAccumulator(iter([_chunk(content="Hello "), _chunk(content="world.")])) + events = list(acc) + assert all(kind == "content" for kind, _ in events) + assert "".join(text for _, text in events) == "Hello world." + assert acc.message.content == "Hello world." + assert acc.message.tool_calls is None + + def test_no_content_yields_none(self): + acc = StreamAccumulator(iter([_chunk(content=None)])) + list(acc) + assert acc.message.content is None + + def test_skips_chunks_with_no_choices(self): + acc = StreamAccumulator(iter([_empty_chunk(), _chunk(content="ok")])) + events = list(acc) + assert events == [("content", "ok")] + assert acc.message.content == "ok" + + +class TestThinkTagHandling: + def test_think_block_at_start(self): + acc = StreamAccumulator(iter([_chunk(content="reasoninganswer")])) + list(acc) + assert acc.message.content == "answer" + assert acc.message.reasoning_content == "reasoning" + + def test_think_open_tag_split_across_chunks(self): + chunks = [_chunk(content="... (e.g. connection + # dropped) must never surface the unfinished reasoning as answer prose. + acc = StreamAccumulator(iter([_chunk(content="unfinished reasoning")])) + events = list(acc) + assert all(kind == "thinking" for kind, _ in events) + assert acc.message.content is None + assert acc.message.reasoning_content == "unfinished reasoning" + + def test_content_before_and_after_think_block(self): + acc = StreamAccumulator(iter([_chunk(content="before middle after")])) + list(acc) + assert acc.message.content == "before after" + assert acc.message.reasoning_content == "middle" + + +class TestToolCalls: + def test_single_tool_call_arguments_split_across_fragments(self): + chunks = [ + _chunk(tool_calls=[_tc(0, id="call_1", name="search_speeches")]), + _chunk(tool_calls=[_tc(0, arguments='{"query":')]), + _chunk(tool_calls=[_tc(0, arguments='"AI"}')]), + ] + acc = StreamAccumulator(iter(chunks)) + events = list(acc) + assert events == [("tool_call", None)] # signalled once, on first sighting + msg = acc.message + assert msg.content is None + assert len(msg.tool_calls) == 1 + assert msg.tool_calls[0].id == "call_1" + assert msg.tool_calls[0].function.name == "search_speeches" + assert msg.tool_calls[0].function.arguments == '{"query":"AI"}' + + def test_two_parallel_tool_calls_by_index(self): + chunks = [ + _chunk(tool_calls=[_tc(0, id="call_1", name="search")]), + _chunk(tool_calls=[_tc(0, arguments='{"q":"AI"}'), _tc(1, id="call_2", name="fetch")]), + _chunk(tool_calls=[_tc(1, arguments='{"id":1}')]), + ] + acc = StreamAccumulator(iter(chunks)) + list(acc) + msg = acc.message + assert [tc.id for tc in msg.tool_calls] == ["call_1", "call_2"] + assert msg.tool_calls[0].function.arguments == '{"q":"AI"}' + assert msg.tool_calls[1].function.name == "fetch" + assert msg.tool_calls[1].function.arguments == '{"id":1}' + + def test_tool_call_event_fires_once_per_stream(self): + chunks = [ + _chunk(tool_calls=[_tc(0, id="call_1", name="search")]), + _chunk(tool_calls=[_tc(0, arguments="{}"), _tc(1, id="call_2", name="fetch")]), + ] + acc = StreamAccumulator(iter(chunks)) + events = list(acc) + assert events.count(("tool_call", None)) == 1 + + def test_narration_content_before_tool_calls(self): + chunks = [ + _chunk(content="Jag söker nu efter tal om AI."), + _chunk(tool_calls=[_tc(0, id="call_1", name="search_speeches", arguments="{}")]), + ] + acc = StreamAccumulator(iter(chunks)) + events = list(acc) + content_text = "".join(text for kind, text in events if kind == "content") + assert content_text == "Jag söker nu efter tal om AI." + assert ("tool_call", None) in events + msg = acc.message + assert msg.content == "Jag söker nu efter tal om AI." + assert msg.tool_calls[0].function.name == "search_speeches" + + +class TestMessageAccess: + def test_raises_if_read_before_exhaustion(self): + acc = StreamAccumulator(iter([_chunk(content="x")])) + with pytest.raises(RuntimeError): + _ = acc.message + + def test_readable_after_full_iteration(self): + acc = StreamAccumulator(iter([_chunk(content="x")])) + for _ in acc: + pass + assert acc.message.content == "x" diff --git a/tests/test_streaming_answer.py b/tests/test_streaming_answer.py new file mode 100644 index 0000000..04131b8 --- /dev/null +++ b/tests/test_streaming_answer.py @@ -0,0 +1,273 @@ +"""Tests for the speculative final-answer streaming filter and driver.""" + +from types import SimpleNamespace + +import pytest + +from backend.services.streaming_answer import ( + AnswerStreamFilter, + _last_unmatched_bracket_start, + run_streaming_iteration, + strip_citation_markers, +) +from packages.llm.client import StreamAccumulator + + +# --------------------------------------------------------------------------- +# _last_unmatched_bracket_start +# --------------------------------------------------------------------------- + + +class TestLastUnmatchedBracketStart: + def test_no_brackets(self): + text = "hello world" + assert _last_unmatched_bracket_start(text) == len(text) + + def test_complete_single_marker(self): + text = "Claim[src:H40911] more" + assert _last_unmatched_bracket_start(text) == len(text) + + def test_incomplete_single_marker(self): + text = "Claim[src:H409" + assert _last_unmatched_bracket_start(text) == len("Claim") + + def test_incomplete_double_bracket_holds_both_brackets(self): + text = "text [[src:AB" + assert _last_unmatched_bracket_start(text) == len("text ") + + def test_complete_double_bracket_marker(self): + text = "text [[src:ABC]] more" + assert _last_unmatched_bracket_start(text) == len(text) + + def test_preserves_non_citation_bracket_once_closed(self): + text = "Array [0] and citation[src:H" + assert _last_unmatched_bracket_start(text) == len("Array [0] and citation") + + def test_orphan_open_bracket_before_unrelated_close_is_not_held(self): + # A stray '[' that's already followed by a later, unrelated ']' can + # never grow into a marker anymore — nothing to hold back. + text = "text [ [nested] stray" + assert _last_unmatched_bracket_start(text) == len(text) + + +# --------------------------------------------------------------------------- +# strip_citation_markers +# --------------------------------------------------------------------------- + + +class TestStripCitationMarkers: + def test_strips_single_bracket(self): + assert strip_citation_markers("Claim[src:H40911] more") == "Claim more" + + def test_strips_double_bracket(self): + assert strip_citation_markers("Claim[[src:H40911]] more") == "Claim more" + + def test_leaves_non_citation_brackets(self): + assert strip_citation_markers("Array [0] here") == "Array [0] here" + + +# --------------------------------------------------------------------------- +# AnswerStreamFilter +# --------------------------------------------------------------------------- + + +class TestAnswerStreamFilterGate: + def test_short_narration_never_arms(self): + f = AnswerStreamFilter(min_chars=400, soft_min_chars=200) + assert f.feed("Jag söker nu efter tal om AI.") is None + assert f.flushed_any is False + + def test_hard_gate_arms_past_min_chars(self): + f = AnswerStreamFilter(min_chars=20, soft_min_chars=1000) + out = f.feed("x" * 25) + assert out is not None + assert f.flushed_any is True + + def test_soft_gate_needs_a_second_sentence(self): + f = AnswerStreamFilter(min_chars=1000, soft_min_chars=10) + # One short sentence past soft_min_chars but no second sentence yet — held back. + assert f.feed("A single narration sentence.") is None + + def test_soft_gate_arms_once_into_a_second_sentence(self): + f = AnswerStreamFilter(min_chars=1000, soft_min_chars=10) + out = f.feed("First sentence done. Second one starts here") + assert out is not None + + def test_flush_remainder_releases_unarmed_short_answer(self): + f = AnswerStreamFilter(min_chars=400, soft_min_chars=200) + assert f.feed("Nej.") is None + remainder = f.flush_remainder() + assert remainder == "Nej." + assert f.flushed_any is True + + def test_flush_remainder_empty_when_nothing_fed(self): + f = AnswerStreamFilter() + assert f.flush_remainder() is None + assert f.flushed_any is False + + +class TestAnswerStreamFilterCitationSafety: + def test_marker_split_across_feed_calls_never_leaks(self): + f = AnswerStreamFilter(min_chars=10, soft_min_chars=5) + f.feed("This is a long enough narration to arm the gate for sure. ") + collected = "" + for chunk in ["Claim one", "[src:H4", "0911]", " and more text after."]: + piece = f.feed(chunk) + if piece: + assert "[src:" not in piece + collected += piece + assert "H40911" not in collected + assert collected == "Claim one and more text after." + + def test_double_bracket_marker_split_never_leaks_orphan_bracket(self): + f = AnswerStreamFilter(min_chars=10, soft_min_chars=5) + f.feed("This is a long enough narration to arm the gate for sure. ") + collected = "" + for chunk in ["Claim ", "[[src:H4", "0911]]", " end."]: + piece = f.feed(chunk) + if piece: + assert "[" not in piece + collected += piece + assert collected == "Claim end." + + def test_incomplete_marker_at_true_end_of_stream_is_dropped_not_leaked(self): + # The trailing "[src:H409" is never going to be completed (no more + # chunks are coming) — flush_remainder must drop it outright rather + # than ever releasing a raw, unclosed citation-marker fragment. + f = AnswerStreamFilter(min_chars=10, soft_min_chars=5) + pieces = [] + pieces.append(f.feed("Long enough narration to arm the gate for sure now. ")) + pieces.append(f.feed("Trailing claim[src:H409")) # incomplete marker, stream ends here + pieces.append(f.flush_remainder()) + collected = "".join(p for p in pieces if p) + assert "[src:" not in collected + assert "H409" not in collected + assert "Trailing claim" in collected + + def test_flush_remainder_strips_markers_in_a_short_never_armed_answer(self): + # Short answers never arm the gate, so their text sits in `_pending` + # and never passes through feed()'s per-chunk bracket handling — + # flush_remainder is the only place that ever filters it. + f = AnswerStreamFilter(min_chars=400, soft_min_chars=200) + assert f.feed("Enligt[src:H40911] nej.") is None + remainder = f.flush_remainder() + assert remainder is not None + assert "[src:" not in remainder + assert "H40911" not in remainder + assert "Enligt" in remainder and "nej." in remainder + + +# --------------------------------------------------------------------------- +# run_streaming_iteration +# --------------------------------------------------------------------------- + + +def _chunk(content=None, tool_calls=None): + delta = SimpleNamespace(content=content, reasoning_content=None, tool_calls=tool_calls) + return SimpleNamespace(choices=[SimpleNamespace(delta=delta)]) + + +def _tc(index, id=None, name=None, arguments=None): + fn = SimpleNamespace(name=name, arguments=arguments) + return SimpleNamespace(index=index, id=id, function=fn) + + +class _FakeLLM: + """Stands in for packages.llm.LLM: generate(..., stream=True) returns a + StreamAccumulator over a canned chunk sequence (or an error string).""" + + model = "fake-model" + + def __init__(self, chunks=None, error=None): + self._chunks = chunks or [] + self._error = error + + def generate(self, **kwargs): + assert kwargs.get("stream") is True + if self._error is not None: + return self._error + return StreamAccumulator(iter(self._chunks)) + + +class TestRunStreamingIteration: + def test_long_answer_streams_deltas_and_returns_message(self): + long_text = "Enligt riksdagens protokoll har flera ledamoter debatterat AI. " * 3 + llm = _FakeLLM(chunks=[_chunk(content=long_text)]) + events = [] + message = run_streaming_iteration(llm, {"messages": []}, events.append) + + deltas = [e for e in events if e["type"] == "answer_delta"] + assert deltas # at least one delta was forwarded live + # feed()/flush_remainder() never strip whitespace (only citation + # markers) — the raw text reconstructs exactly, unlike .message.content + # below, which goes through the same _strip_think().strip() the + # non-streaming path always applies. + assert "".join(d["text"] for d in deltas) == long_text + assert not any(e["type"] == "answer_delta_retract" for e in events) + assert message.tool_calls is None + assert message.content == long_text.strip() + + def test_short_answer_flushes_once_at_the_end(self): + llm = _FakeLLM(chunks=[_chunk(content="Nej.")]) + events = [] + message = run_streaming_iteration(llm, {"messages": []}, events.append) + deltas = [e for e in events if e["type"] == "answer_delta"] + assert "".join(d["text"] for d in deltas) == "Nej." + assert message.content == "Nej." + + def test_citation_marker_never_reaches_an_event(self): + text = "Långt påstående med källa[src:H40911] och mer text för att passera spärren. " + llm = _FakeLLM(chunks=[_chunk(content=text)]) + events = [] + run_streaming_iteration(llm, {"messages": []}, events.append) + for e in events: + if e["type"] == "answer_delta": + assert "[src:" not in e["text"] + + def test_tool_call_after_long_narration_emits_exactly_one_retract(self): + long_narration = "Jag har läst igenom flera anföranden om detta ämne och tänker nu göra en sökning. " * 3 + llm = _FakeLLM(chunks=[ + _chunk(content=long_narration), + _chunk(tool_calls=[_tc(0, id="call_1", name="search_speeches", arguments="{}")]), + ]) + events = [] + message = run_streaming_iteration(llm, {"messages": []}, events.append) + + retracts = [e for e in events if e["type"] == "answer_delta_retract"] + deltas = [e for e in events if e["type"] == "answer_delta"] + assert deltas # the narration was long enough to have cleared the gate + assert len(retracts) == 1 + assert message.tool_calls is not None + assert message.tool_calls[0].function.name == "search_speeches" + + def test_short_narration_before_tool_call_never_retracts(self): + llm = _FakeLLM(chunks=[ + _chunk(content="Jag söker nu."), + _chunk(tool_calls=[_tc(0, id="call_1", name="search_speeches", arguments="{}")]), + ]) + events = [] + message = run_streaming_iteration(llm, {"messages": []}, events.append) + assert not any(e["type"] == "answer_delta" for e in events) + assert not any(e["type"] == "answer_delta_retract" for e in events) + assert message.tool_calls[0].function.name == "search_speeches" + + def test_generate_error_string_raises(self): + llm = _FakeLLM(error="LLM request failed: connection refused") + with pytest.raises(RuntimeError, match="connection refused"): + run_streaming_iteration(llm, {"messages": []}, lambda e: None) + + def test_mid_stream_exception_raises_runtime_error(self): + def bad_chunks(): + yield _chunk(content="partial answer text that starts fine") + raise ConnectionError("dropped") + + llm = _FakeLLM() + llm.generate = lambda **kwargs: StreamAccumulator(bad_chunks()) + with pytest.raises(RuntimeError, match="dropped"): + run_streaming_iteration(llm, {"messages": []}, lambda e: None) + + def test_none_event_callback_is_safe(self): + long_text = "Ett långt svar som absolut ska passera konfidensgränsen för strömning. " * 3 + llm = _FakeLLM(chunks=[_chunk(content=long_text)]) + message = run_streaming_iteration(llm, {"messages": []}, None) + assert message.content == long_text.strip()