You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
153 lines
6.5 KiB
153 lines
6.5 KiB
"""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"
|
|
|