feat: convert StreamHandler to callback-based (remove Rich.Live)
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -1,41 +1,56 @@
|
||||
"""Streaming response handler — accumulates SSE chunks into a complete Message."""
|
||||
|
||||
from collections.abc import AsyncIterator
|
||||
|
||||
from rich.live import Live
|
||||
from rich.markdown import Markdown
|
||||
from rich.panel import Panel
|
||||
import time
|
||||
from collections.abc import AsyncIterator, Callable
|
||||
|
||||
from app.models.config import DisplayConfig
|
||||
from app.models.message import Message
|
||||
from app.models.tool_call import ToolCall, ToolCallFunction
|
||||
from app.utils.logging import console, get_logger
|
||||
from app.utils.logging import get_logger
|
||||
from app.utils.token_counter import TokenUsage
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
# Minimum interval between content update callbacks (seconds)
|
||||
_UPDATE_THROTTLE_INTERVAL = 0.1
|
||||
|
||||
|
||||
class StreamHandler:
|
||||
"""Processes an SSE chunk stream into a Rich live display and final Message.
|
||||
"""Processes an SSE chunk stream and produces a complete assistant Message.
|
||||
|
||||
Accumulates content deltas and tool call fragments, renders a live Markdown
|
||||
panel during streaming, and produces a complete assistant Message on finish.
|
||||
Accumulates content deltas and tool call fragments. Notifies the UI via
|
||||
optional callbacks during streaming.
|
||||
"""
|
||||
|
||||
def __init__(self, display_config: DisplayConfig) -> None:
|
||||
"""Initialize the stream handler.
|
||||
|
||||
Args:
|
||||
display_config: Display preferences (streaming toggle, etc.).
|
||||
"""
|
||||
self._display_config = display_config
|
||||
self._accumulated_content: str = ""
|
||||
self._accumulated_reasoning: str = ""
|
||||
self._tool_calls: dict[int, dict[str, str]] = {}
|
||||
self._usage: TokenUsage | None = None
|
||||
self._on_content: Callable[[str], None] | None = None
|
||||
self._on_thinking: Callable[[], None] | None = None
|
||||
self._on_done: Callable[[], None] | None = None
|
||||
|
||||
def set_callbacks(
|
||||
self,
|
||||
on_content: Callable[[str], None] | None = None,
|
||||
on_thinking: Callable[[], None] | None = None,
|
||||
on_done: Callable[[], None] | None = None,
|
||||
) -> None:
|
||||
"""Set UI callbacks for streaming updates.
|
||||
|
||||
Args:
|
||||
on_content: Called with accumulated content string (throttled to ~100ms).
|
||||
on_thinking: Called once when first reasoning token arrives.
|
||||
on_done: Called when streaming is complete.
|
||||
"""
|
||||
self._on_content = on_content
|
||||
self._on_thinking = on_thinking
|
||||
self._on_done = on_done
|
||||
|
||||
async def process_stream(self, chunk_iter: AsyncIterator[dict]) -> Message:
|
||||
"""Consume a chunk iterator, rendering live output and returning the final Message.
|
||||
"""Consume a chunk iterator and return the final Message.
|
||||
|
||||
Args:
|
||||
chunk_iter: Async iterator of parsed SSE chunk dicts.
|
||||
@@ -43,26 +58,42 @@ class StreamHandler:
|
||||
Returns:
|
||||
Complete assistant Message with accumulated content and tool calls.
|
||||
"""
|
||||
with Live(console=console, refresh_per_second=8) as live:
|
||||
async for chunk in chunk_iter:
|
||||
self._process_chunk(chunk)
|
||||
thinking_notified = False
|
||||
last_update_time = 0.0
|
||||
|
||||
# Show reasoning while waiting for content
|
||||
display_text = self._accumulated_content
|
||||
if not display_text and self._accumulated_reasoning:
|
||||
display_text = "*thinking...*"
|
||||
async for chunk in chunk_iter:
|
||||
self._process_chunk(chunk)
|
||||
|
||||
if display_text and self._display_config.stream_output:
|
||||
# Render inside the same Assistant panel used for final output
|
||||
# so the live display and final frame are visually consistent
|
||||
live.update(
|
||||
Panel(
|
||||
Markdown(display_text),
|
||||
title="Assistant",
|
||||
border_style="green",
|
||||
expand=True,
|
||||
)
|
||||
)
|
||||
if not self._display_config.stream_output:
|
||||
continue
|
||||
|
||||
# Notify thinking once
|
||||
if (
|
||||
not thinking_notified
|
||||
and not self._accumulated_content
|
||||
and self._accumulated_reasoning
|
||||
and self._on_thinking is not None
|
||||
):
|
||||
self._on_thinking()
|
||||
thinking_notified = True
|
||||
|
||||
# Throttled content updates
|
||||
if self._accumulated_content and self._on_content is not None:
|
||||
now = time.monotonic()
|
||||
if now - last_update_time >= _UPDATE_THROTTLE_INTERVAL:
|
||||
self._on_content(self._accumulated_content)
|
||||
last_update_time = now
|
||||
|
||||
# Final content update (ensures last chunk is shown)
|
||||
if (
|
||||
self._display_config.stream_output
|
||||
and self._accumulated_content
|
||||
and self._on_content is not None
|
||||
):
|
||||
self._on_content(self._accumulated_content)
|
||||
|
||||
if self._on_done is not None:
|
||||
self._on_done()
|
||||
|
||||
tool_calls = self._build_tool_calls() or None
|
||||
return Message(
|
||||
@@ -72,12 +103,7 @@ class StreamHandler:
|
||||
)
|
||||
|
||||
def _process_chunk(self, chunk: dict) -> None:
|
||||
"""Extract content, tool calls, and usage from a single SSE chunk.
|
||||
|
||||
Args:
|
||||
chunk: Parsed JSON dict from one SSE data line.
|
||||
"""
|
||||
# Content delta
|
||||
"""Extract content, tool calls, and usage from a single SSE chunk."""
|
||||
choices = chunk.get("choices", [])
|
||||
if choices:
|
||||
delta = choices[0].get("delta", {})
|
||||
@@ -86,12 +112,10 @@ class StreamHandler:
|
||||
if content_piece:
|
||||
self._accumulated_content += content_piece
|
||||
|
||||
# Reasoning tokens (e.g. qwen3.5 thinking mode)
|
||||
reasoning_piece = delta.get("reasoning")
|
||||
if reasoning_piece:
|
||||
self._accumulated_reasoning += reasoning_piece
|
||||
|
||||
# Tool call deltas (accumulated by index)
|
||||
for tc_delta in delta.get("tool_calls", []):
|
||||
idx = tc_delta.get("index", 0)
|
||||
if idx not in self._tool_calls:
|
||||
@@ -109,7 +133,6 @@ class StreamHandler:
|
||||
if func.get("arguments"):
|
||||
entry["arguments"] += func["arguments"]
|
||||
|
||||
# Token usage (typically in the final chunk)
|
||||
usage_data = chunk.get("usage")
|
||||
if usage_data:
|
||||
self._usage = TokenUsage(
|
||||
@@ -119,11 +142,7 @@ class StreamHandler:
|
||||
)
|
||||
|
||||
def _build_tool_calls(self) -> list[ToolCall]:
|
||||
"""Convert accumulated tool call fragments into sorted ToolCall list.
|
||||
|
||||
Returns:
|
||||
List of ToolCall objects sorted by stream index.
|
||||
"""
|
||||
"""Convert accumulated tool call fragments into sorted ToolCall list."""
|
||||
if not self._tool_calls:
|
||||
return []
|
||||
|
||||
@@ -143,14 +162,7 @@ class StreamHandler:
|
||||
return result
|
||||
|
||||
def get_partial_message(self) -> Message | None:
|
||||
"""Return whatever content/tool_calls have been accumulated so far.
|
||||
|
||||
Useful for mid-stream interruption recovery — returns None if nothing
|
||||
has been accumulated yet.
|
||||
|
||||
Returns:
|
||||
Partial assistant Message, or None if no content accumulated.
|
||||
"""
|
||||
"""Return whatever content/tool_calls have been accumulated so far."""
|
||||
tool_calls = self._build_tool_calls() or None
|
||||
if not self._accumulated_content and not tool_calls:
|
||||
return None
|
||||
@@ -176,3 +188,6 @@ class StreamHandler:
|
||||
self._accumulated_reasoning = ""
|
||||
self._tool_calls.clear()
|
||||
self._usage = None
|
||||
self._on_content = None
|
||||
self._on_thinking = None
|
||||
self._on_done = None
|
||||
|
||||
111
tests/unit/test_streaming.py
Normal file
111
tests/unit/test_streaming.py
Normal file
@@ -0,0 +1,111 @@
|
||||
"""Tests for callback-based StreamHandler."""
|
||||
|
||||
import asyncio
|
||||
from unittest.mock import MagicMock, call
|
||||
|
||||
import pytest
|
||||
|
||||
from app.models.config import DisplayConfig
|
||||
from app.services.streaming import StreamHandler
|
||||
|
||||
|
||||
def _make_chunk(content: str | None = None, reasoning: str | None = None) -> dict:
|
||||
"""Helper to create a fake SSE chunk."""
|
||||
delta: dict = {}
|
||||
if content is not None:
|
||||
delta["content"] = content
|
||||
if reasoning is not None:
|
||||
delta["reasoning"] = reasoning
|
||||
return {"choices": [{"delta": delta}]}
|
||||
|
||||
|
||||
def _make_tool_call_chunk(index: int, tc_id: str = "", name: str = "", args: str = "") -> dict:
|
||||
"""Helper to create a fake tool call chunk."""
|
||||
tc_delta: dict = {"index": index}
|
||||
if tc_id:
|
||||
tc_delta["id"] = tc_id
|
||||
func: dict = {}
|
||||
if name:
|
||||
func["name"] = name
|
||||
if args:
|
||||
func["arguments"] = args
|
||||
if func:
|
||||
tc_delta["function"] = func
|
||||
return {"choices": [{"delta": {"tool_calls": [tc_delta]}}]}
|
||||
|
||||
|
||||
async def _async_iter(items: list[dict]):
|
||||
for item in items:
|
||||
yield item
|
||||
|
||||
|
||||
class TestStreamHandlerCallbacks:
|
||||
@pytest.mark.asyncio
|
||||
async def test_on_content_called_with_accumulated_text(self) -> None:
|
||||
handler = StreamHandler(DisplayConfig(stream_output=True))
|
||||
on_content = MagicMock()
|
||||
on_thinking = MagicMock()
|
||||
on_done = MagicMock()
|
||||
handler.set_callbacks(on_content=on_content, on_thinking=on_thinking, on_done=on_done)
|
||||
|
||||
chunks = [_make_chunk(content="Hello"), _make_chunk(content=" world")]
|
||||
msg = await handler.process_stream(_async_iter(chunks))
|
||||
|
||||
assert msg.content == "Hello world"
|
||||
assert on_content.call_count >= 1
|
||||
# Last call should have full accumulated content
|
||||
last_content = on_content.call_args_list[-1][0][0]
|
||||
assert last_content == "Hello world"
|
||||
on_done.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_on_thinking_called_for_reasoning(self) -> None:
|
||||
handler = StreamHandler(DisplayConfig(stream_output=True))
|
||||
on_content = MagicMock()
|
||||
on_thinking = MagicMock()
|
||||
on_done = MagicMock()
|
||||
handler.set_callbacks(on_content=on_content, on_thinking=on_thinking, on_done=on_done)
|
||||
|
||||
chunks = [_make_chunk(reasoning="let me think")]
|
||||
msg = await handler.process_stream(_async_iter(chunks))
|
||||
|
||||
on_thinking.assert_called_once()
|
||||
on_content.assert_not_called()
|
||||
on_done.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_display_callbacks_skip_when_stream_output_disabled(self) -> None:
|
||||
"""on_content and on_thinking are suppressed, but on_done always fires."""
|
||||
handler = StreamHandler(DisplayConfig(stream_output=False))
|
||||
on_content = MagicMock()
|
||||
on_thinking = MagicMock()
|
||||
on_done = MagicMock()
|
||||
handler.set_callbacks(on_content=on_content, on_thinking=on_thinking, on_done=on_done)
|
||||
|
||||
chunks = [_make_chunk(content="Hello")]
|
||||
msg = await handler.process_stream(_async_iter(chunks))
|
||||
|
||||
assert msg.content == "Hello"
|
||||
on_content.assert_not_called()
|
||||
on_thinking.assert_not_called()
|
||||
on_done.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_callbacks_by_default(self) -> None:
|
||||
"""process_stream works without set_callbacks (backward compat)."""
|
||||
handler = StreamHandler(DisplayConfig(stream_output=True))
|
||||
chunks = [_make_chunk(content="Hello")]
|
||||
msg = await handler.process_stream(_async_iter(chunks))
|
||||
assert msg.content == "Hello"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_calls_still_accumulated(self) -> None:
|
||||
handler = StreamHandler(DisplayConfig(stream_output=True))
|
||||
chunks = [
|
||||
_make_tool_call_chunk(0, tc_id="call_1", name="read_file"),
|
||||
_make_tool_call_chunk(0, args='{"path": "foo.py"}'),
|
||||
]
|
||||
msg = await handler.process_stream(_async_iter(chunks))
|
||||
assert msg.tool_calls is not None
|
||||
assert len(msg.tool_calls) == 1
|
||||
assert msg.tool_calls[0].function.name == "read_file"
|
||||
Reference in New Issue
Block a user