mirror of
https://github.com/openswarm-ai/openswarm.git
synced 2026-08-20 11:42:23 +02:00
71 lines
3.6 KiB
Python
71 lines
3.6 KiB
Python
"""Mutable per-turn state for the agent streaming loop. Replaces the loop's `nonlocal`
|
|
locals with a captured object so the streaming/thinking closures can eventually move out of
|
|
agent_manager (a closure that mutates `state.field` needs no `nonlocal`)."""
|
|
|
|
import asyncio
|
|
from typing import Dict, List, Optional
|
|
|
|
from pydantic import BaseModel, ConfigDict, InstanceOf
|
|
|
|
|
|
class ThinkingState(BaseModel):
|
|
"""The consolidated-thinking side-channel for one turn: the live 'Thought for Ns ·
|
|
N tokens · N tools' pill. A single persisted message id is reused across a multi-step
|
|
turn so the bubble updates in place; everything resets at the next turn boundary."""
|
|
|
|
model_config = ConfigDict(validate_assignment=True, arbitrary_types_allowed=True)
|
|
|
|
# block index -> wall-clock start (s); popped to accumulate total_ms when a block ends.
|
|
block_starts: Dict[int, float] = {}
|
|
total_ms: int = 0
|
|
# Stable id for the turn's single thinking message (frontend dedupe replaces in place).
|
|
msg_id: Optional[str] = None
|
|
text_parts: List[str] = []
|
|
# Latest Gemini thoughtSignature, re-attached on later turns for reasoning continuity.
|
|
thought_signature: Optional[str] = None
|
|
# Background ticker handle; re-emits the pill every 1s so the elapsed counter keeps moving.
|
|
ticker_task: Optional[InstanceOf[asyncio.Task]] = None
|
|
|
|
|
|
class TurnState(BaseModel):
|
|
"""Mutable per-turn streaming state: the live streaming-message ids, the accumulated
|
|
assistant text, and the running token/char/timing counters. Reset at each turn boundary.
|
|
(validate_assignment runs per SDK event, not per token, so the cost is negligible against
|
|
a multi-second turn.)"""
|
|
|
|
model_config = ConfigDict(validate_assignment=True)
|
|
|
|
stream_text_msg_id: Optional[str] = None
|
|
stream_tool_msg_ids_ordered: List[str] = []
|
|
stream_block_index_map: Dict[int, str] = {}
|
|
stream_text_accum: str = ""
|
|
current_turn_emitted: bool = False
|
|
number: int = 0
|
|
first_event: bool = True
|
|
tool_count: int = 0
|
|
started_ts: Optional[float] = None
|
|
total_ms: int = 0
|
|
# ACTIVE time: sum of inter-event deltas with stall gaps capped, so a turn that waits an hour
|
|
# on an approval or a wedged workflow books seconds of work, not the wall gap (ENG-189).
|
|
last_event_ts: Optional[float] = None
|
|
active_ms: int = 0
|
|
output_tokens: int = 0
|
|
assistant_text_chars: int = 0
|
|
tool_input_chars: int = 0
|
|
# Cumulative-token snapshot taken at turn start; subtracted at emit time so the thinking pill shows THIS turn's new tokens, not the whole session's running total.
|
|
baseline_session_in: int = 0
|
|
baseline_session_out: int = 0
|
|
baseline_children_in: int = 0
|
|
baseline_children_out: int = 0
|
|
baseline_captured: bool = False
|
|
# CLI compact_boundary events seen this turn; one plus a ProcessError = the autocompact-thrash death the context-pressure valve retries.
|
|
compact_boundaries: int = 0
|
|
# Provider 500s/429s the CLI retried on its own; the user sees only a long silence, so these are counted rather than lost.
|
|
provider_retries: int = 0
|
|
provider_retry_wait_ms: int = 0
|
|
# Mid-turn context breaker: fires once per turn, and only after a below-trigger reading (a turn that STARTS over the trigger must run, or a failed shrink would break-loop forever).
|
|
context_break_fired: bool = False
|
|
saw_input_below_trigger: bool = False
|
|
# The LAST inference step's request size (input + cache read + cache creation): the true live context. The ResultMessage's usage sums these across every step of the turn, which is billing, not context.
|
|
last_step_input: int = 0
|