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