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 <think>/</think> 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 <think> 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).main
Lasse Server
5 days ago
parent
87a6cd28bc
commit
7ea92c9366
12 changed files with 975 additions and 37 deletions
@ -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 |
||||
@ -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 |
||||
# <think> 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="<think>reasoning</think>answer")])) |
||||
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="<thi"), _chunk(content="nk>reasoning here</th"), _chunk(content="ink>answer text")] |
||||
acc = StreamAccumulator(iter(chunks)) |
||||
list(acc) |
||||
assert acc.message.content == "answer text" |
||||
assert acc.message.reasoning_content == "reasoning here" |
||||
|
||||
def test_reasoning_content_field_no_inline_think(self): |
||||
acc = StreamAccumulator(iter([_chunk(reasoning="thinking..."), _chunk(content="the answer")])) |
||||
events = list(acc) |
||||
assert ("thinking", "thinking...") in events |
||||
assert acc.message.content == "the answer" |
||||
assert acc.message.reasoning_content == "thinking..." |
||||
|
||||
def test_truncated_stream_mid_think_block_never_leaks_as_content(self): |
||||
# A stream that ends while still inside <think>...</think> (e.g. connection |
||||
# dropped) must never surface the unfinished reasoning as answer prose. |
||||
acc = StreamAccumulator(iter([_chunk(content="<think>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 <think>middle</think> 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" |
||||
@ -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() |
||||
Loading…
Reference in new issue