174 lines
6.8 KiB
Python
174 lines
6.8 KiB
Python
"""Tests for how a turn's streamed assistant text is built up.
|
|
|
|
The text used to be grown with ``+=`` on an attribute. Python cannot grow a
|
|
string in place there, so every delta copied the whole thing again and a long
|
|
reply cost the square of its length in copying. The text is now held as a list
|
|
of pieces and joined when something reads it.
|
|
|
|
These tests cover the behaviour callers depend on, plus a check on the stored
|
|
pieces that fails if the copying ever comes back.
|
|
"""
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
|
|
def _make_agent():
|
|
from run_agent import AIAgent
|
|
|
|
agent = AIAgent(
|
|
api_key="test-key",
|
|
base_url="https://openrouter.ai/api/v1",
|
|
model="test/model",
|
|
quiet_mode=True,
|
|
skip_context_files=True,
|
|
skip_memory=True,
|
|
)
|
|
agent.api_mode = "chat_completions"
|
|
agent._interrupt_requested = False
|
|
return agent
|
|
|
|
|
|
class TestStreamedTextValue:
|
|
"""The value callers read must not change."""
|
|
|
|
def test_starts_empty(self):
|
|
agent = _make_agent()
|
|
assert agent._current_streamed_assistant_text == ""
|
|
|
|
def test_deltas_join_in_order(self):
|
|
agent = _make_agent()
|
|
for piece in ["Hello", ", ", "world", "!"]:
|
|
agent._record_streamed_assistant_text(piece)
|
|
assert agent._current_streamed_assistant_text == "Hello, world!"
|
|
|
|
def test_reading_twice_gives_the_same_answer(self):
|
|
agent = _make_agent()
|
|
agent._record_streamed_assistant_text("one ")
|
|
agent._record_streamed_assistant_text("two")
|
|
first = agent._current_streamed_assistant_text
|
|
second = agent._current_streamed_assistant_text
|
|
assert first == second == "one two"
|
|
|
|
def test_reading_does_not_stop_later_deltas(self):
|
|
agent = _make_agent()
|
|
agent._record_streamed_assistant_text("before ")
|
|
assert agent._current_streamed_assistant_text == "before "
|
|
agent._record_streamed_assistant_text("after")
|
|
assert agent._current_streamed_assistant_text == "before after"
|
|
|
|
def test_direct_assignment_still_works(self):
|
|
# Several call sites set this attribute straight, both to seed a value
|
|
# and to clear it between turns.
|
|
agent = _make_agent()
|
|
agent._record_streamed_assistant_text("thrown away")
|
|
agent._current_streamed_assistant_text = "set by hand"
|
|
assert agent._current_streamed_assistant_text == "set by hand"
|
|
agent._record_streamed_assistant_text(" plus more")
|
|
assert agent._current_streamed_assistant_text == "set by hand plus more"
|
|
|
|
def test_clearing_resets_to_empty(self):
|
|
agent = _make_agent()
|
|
agent._record_streamed_assistant_text("left over")
|
|
agent._current_streamed_assistant_text = ""
|
|
assert agent._current_streamed_assistant_text == ""
|
|
agent._record_streamed_assistant_text("new turn")
|
|
assert agent._current_streamed_assistant_text == "new turn"
|
|
|
|
def test_empty_and_non_string_deltas_are_ignored(self):
|
|
agent = _make_agent()
|
|
agent._record_streamed_assistant_text("keep")
|
|
agent._record_streamed_assistant_text("")
|
|
agent._record_streamed_assistant_text(None) # type: ignore[arg-type]
|
|
agent._record_streamed_assistant_text(12345) # type: ignore[arg-type]
|
|
assert agent._current_streamed_assistant_text == "keep"
|
|
|
|
def test_superseded_writer_is_still_fenced_out(self):
|
|
# The single-writer guard (#65991) must keep working now that the
|
|
# text is stored as pieces.
|
|
agent = _make_agent()
|
|
agent._record_streamed_assistant_text("allowed")
|
|
with patch.object(agent, "_stream_writer_superseded", return_value=True):
|
|
agent._record_streamed_assistant_text("blocked")
|
|
assert agent._current_streamed_assistant_text == "allowed"
|
|
|
|
|
|
class TestStreamedTextCost:
|
|
"""Adding a delta must not touch the text already collected.
|
|
|
|
Checked by looking at the stored pieces rather than by timing, so the
|
|
test gives the same answer on a busy CI box as it does on a quiet one.
|
|
"""
|
|
|
|
def test_each_delta_is_stored_as_its_own_piece(self):
|
|
agent = _make_agent()
|
|
for i in range(500):
|
|
agent._record_streamed_assistant_text(f"delta-{i} ")
|
|
# One piece per delta means nothing joined or copied the text that was
|
|
# already there. If a delta ever rebuilds the whole string again, this
|
|
# collapses to a single piece and the test fails.
|
|
assert len(agent._streamed_assistant_text_parts) == 500
|
|
|
|
def test_reading_the_text_does_not_collapse_the_pieces(self):
|
|
# Collapsing on read would drop any delta that lands between the join
|
|
# and the write back, so reading has to leave the pieces alone.
|
|
agent = _make_agent()
|
|
for i in range(10):
|
|
agent._record_streamed_assistant_text(str(i))
|
|
assert agent._current_streamed_assistant_text == "0123456789"
|
|
assert len(agent._streamed_assistant_text_parts) == 10
|
|
|
|
def test_a_long_reply_is_assembled_correctly(self):
|
|
agent = _make_agent()
|
|
delta = "x" * 8
|
|
for _ in range(20000):
|
|
agent._record_streamed_assistant_text(delta)
|
|
assert agent._current_streamed_assistant_text == delta * 20000
|
|
assert len(agent._streamed_assistant_text_parts) == 20000
|
|
|
|
|
|
def _agent_with_sink():
|
|
agent = _make_agent()
|
|
delivered = []
|
|
agent.stream_delta_callback = delivered.append
|
|
agent._stream_callback = None
|
|
return agent, delivered
|
|
|
|
|
|
class TestFireStreamDeltaEmptiness:
|
|
"""_fire_stream_delta used to join the whole reply on every token just
|
|
to decide whether to strip leading newlines. That check now looks at
|
|
the parts list.
|
|
"""
|
|
|
|
def test_first_delta_strips_leading_newlines(self):
|
|
agent, delivered = _agent_with_sink()
|
|
agent._fire_stream_delta("\n\nhello")
|
|
assert delivered == ["hello"]
|
|
assert agent._current_streamed_assistant_text == "hello"
|
|
|
|
def test_later_delta_keeps_leading_newlines(self):
|
|
agent, delivered = _agent_with_sink()
|
|
agent._fire_stream_delta("hello")
|
|
agent._fire_stream_delta("\n\nworld")
|
|
assert delivered == ["hello", "\n\nworld"]
|
|
assert agent._current_streamed_assistant_text == "hello\n\nworld"
|
|
|
|
def test_after_clear_the_next_delta_strips_again(self):
|
|
agent, delivered = _agent_with_sink()
|
|
agent._fire_stream_delta("hello")
|
|
agent._current_streamed_assistant_text = ""
|
|
agent._fire_stream_delta("\n\nagain")
|
|
assert delivered[-1] == "again"
|
|
assert agent._current_streamed_assistant_text == "again"
|
|
|
|
def test_fire_path_stores_one_piece_per_delta(self):
|
|
agent, _delivered = _agent_with_sink()
|
|
for i in range(200):
|
|
agent._fire_stream_delta(f"d{i} ")
|
|
assert len(agent._streamed_assistant_text_parts) == 200
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(pytest.main([__file__, "-q"]))
|