feat: convert StreamHandler to callback-based (remove Rich.Live)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
2026-03-11 12:44:11 -05:00
parent 8335978583
commit 021fe340c1
2 changed files with 181 additions and 55 deletions

View File

@@ -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