diff --git a/assistant_agent.py b/assistant_agent.py index d948626..94396ac 100644 --- a/assistant_agent.py +++ b/assistant_agent.py @@ -6,6 +6,8 @@ import urllib.request from collections.abc import Iterator from typing import Any +from llm_stream import OpenAIStreamAccumulator + class ReviewAssistantError(RuntimeError): pass @@ -41,6 +43,7 @@ def stream_review_assistant( try: with urllib.request.urlopen(request, timeout=timeout) as response: yielded = False + accumulator = OpenAIStreamAccumulator() for raw_line in response: line = raw_line.decode("utf-8", errors="replace").strip() if not line or line.startswith(":"): @@ -57,10 +60,7 @@ def stream_review_assistant( if not choices: continue choice = choices[0] or {} - delta = choice.get("delta") or {} - content = delta.get("content") - if content is None: - content = (choice.get("message") or {}).get("content") + content = accumulator.feed(choice) if content: yielded = True yield str(content) diff --git a/llm_stream.py b/llm_stream.py new file mode 100644 index 0000000..3759fbe --- /dev/null +++ b/llm_stream.py @@ -0,0 +1,40 @@ +from __future__ import annotations + +from typing import Any + + +class OpenAIStreamAccumulator: + """Normalize incremental deltas and provider-specific full-message snapshots.""" + + def __init__(self) -> None: + self.text = "" + self.saw_delta = False + + def feed(self, choice: dict[str, Any]) -> str: + delta = choice.get("delta") + if isinstance(delta, dict) and delta.get("content") is not None: + chunk = str(delta.get("content") or "") + if chunk: + self.saw_delta = True + self.text += chunk + return chunk + + message = choice.get("message") + if not isinstance(message, dict) or message.get("content") is None: + return "" + snapshot = str(message.get("content") or "") + if not snapshot: + return "" + if not self.text: + self.text = snapshot + return snapshot + if snapshot == self.text or self.text.startswith(snapshot): + return "" + if snapshot.startswith(self.text): + suffix = snapshot[len(self.text):] + self.text = snapshot + return suffix + if self.saw_delta: + # A final full snapshot cannot safely replace chunks already delivered. + return "" + return "" diff --git a/mentor_agent.py b/mentor_agent.py index 58cac2a..0c0432e 100644 --- a/mentor_agent.py +++ b/mentor_agent.py @@ -10,6 +10,8 @@ from dataclasses import dataclass from pathlib import Path from typing import Any +from llm_stream import OpenAIStreamAccumulator + class MentorAgentError(RuntimeError): pass @@ -211,6 +213,7 @@ def stream_with_mentor( try: with urllib.request.urlopen(request, timeout=timeout) as response: yielded = False + accumulator = OpenAIStreamAccumulator() for raw_line in response: line = raw_line.decode("utf-8", errors="replace").strip() if not line or line.startswith(":"): @@ -227,9 +230,7 @@ def stream_with_mentor( if not choices: continue choice = choices[0] or {} - content = (choice.get("delta") or {}).get("content") - if content is None: - content = (choice.get("message") or {}).get("content") + content = accumulator.feed(choice) if content: yielded = True yield str(content) diff --git a/tests/test_llm_stream.py b/tests/test_llm_stream.py new file mode 100644 index 0000000..a92f1e3 --- /dev/null +++ b/tests/test_llm_stream.py @@ -0,0 +1,35 @@ +from __future__ import annotations + +import unittest + +from llm_stream import OpenAIStreamAccumulator + + +class OpenAIStreamAccumulatorTests(unittest.TestCase): + def test_repeated_delta_chunks_are_preserved_as_model_output(self) -> None: + accumulator = OpenAIStreamAccumulator() + self.assertEqual(accumulator.feed({"delta": {"content": "yes"}}), "yes") + self.assertEqual(accumulator.feed({"delta": {"content": "yes"}}), "yes") + self.assertEqual(accumulator.text, "yesyes") + + def test_final_snapshot_can_add_a_missing_suffix(self) -> None: + accumulator = OpenAIStreamAccumulator() + self.assertEqual(accumulator.feed({"delta": {"content": "first"}}), "first") + self.assertEqual( + accumulator.feed({"message": {"content": "first second"}}), + " second", + ) + self.assertEqual(accumulator.text, "first second") + + def test_incompatible_final_snapshot_is_not_appended_twice(self) -> None: + accumulator = OpenAIStreamAccumulator() + accumulator.feed({"delta": {"content": "streamed answer"}}) + self.assertEqual( + accumulator.feed({"message": {"content": "rewritten final answer"}}), + "", + ) + self.assertEqual(accumulator.text, "streamed answer") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_mentor_stream.py b/tests/test_mentor_stream.py index 003bed0..1dad1ab 100644 --- a/tests/test_mentor_stream.py +++ b/tests/test_mentor_stream.py @@ -66,6 +66,43 @@ class MentorStreamTests(unittest.TestCase): ) self.assertEqual(result["answer"], "first second") + def test_final_full_message_does_not_duplicate_streamed_deltas(self): + lines = [ + b'data: {"choices":[{"delta":{"content":"first"}}]}\n', + b'data: {"choices":[{"delta":{"content":" second"}}]}\n', + b'data: {"choices":[{"message":{"content":"first second"}}]}\n', + b"data: [DONE]\n", + ] + with patch( + "mentor_agent.urllib.request.urlopen", + return_value=FakeStreamResponse(lines), + ): + chunks = list( + stream_with_mentor( + self.skill, {}, "question", [], "key", + "https://example.test/v1", "model", + ) + ) + self.assertEqual(chunks, ["first", " second"]) + + def test_snapshot_only_stream_emits_only_new_suffix(self): + lines = [ + b'data: {"choices":[{"message":{"content":"first"}}]}\n', + b'data: {"choices":[{"message":{"content":"first second"}}]}\n', + b"data: [DONE]\n", + ] + with patch( + "mentor_agent.urllib.request.urlopen", + return_value=FakeStreamResponse(lines), + ): + chunks = list( + stream_with_mentor( + self.skill, {}, "question", [], "key", + "https://example.test/v1", "model", + ) + ) + self.assertEqual(chunks, ["first", " second"]) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_review_assistant.py b/tests/test_review_assistant.py index dd767ba..b470516 100644 --- a/tests/test_review_assistant.py +++ b/tests/test_review_assistant.py @@ -45,6 +45,23 @@ class ReviewAssistantStreamingTests(unittest.TestCase): stream_review_assistant({}, "question", [], "key", "https://example.test/v1", "model") ) + def test_final_full_message_does_not_duplicate_streamed_deltas(self): + response = StreamingResponse( + [ + b'data: {"choices":[{"delta":{"content":"first"}}]}\n', + b'data: {"choices":[{"delta":{"content":" second"}}]}\n', + b'data: {"choices":[{"message":{"content":"first second"}}]}\n', + b"data: [DONE]\n", + ] + ) + with patch("assistant_agent.urllib.request.urlopen", return_value=response): + chunks = list( + stream_review_assistant( + {}, "question", [], "key", "https://example.test/v1", "model" + ) + ) + self.assertEqual(chunks, ["first", " second"]) + def test_ndjson_event_writer_flushes_complete_line(self): handler = RequestHandler.__new__(RequestHandler) handler.wfile = io.BytesIO()