Coverage for src / lexigram / ai / relay / stream / state.py: 96%
194 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-08-08 23:08 +0800
1"""Mutable stream session state and shared lifecycle rules.
3The session is the state machine shared by every target emitter. It
4accepts one source wire event at a time (producing zero or more
5``StreamDelta`` objects), applies deltas in order, hands each delta to
6the configured target emitter, and finalizes idempotently when the
7upstream stream truncates or completes.
8"""
10from __future__ import annotations
12from dataclasses import dataclass
13from typing import Any, Protocol, runtime_checkable
15from lexigram.ai.relay.errors import stream_already_finalized, stream_state_invalid
16from lexigram.contracts.ai.exceptions import RelayError, RelayErrorCode
17from lexigram.contracts.ai.relay.ir import StreamDelta, StreamState
18from lexigram.contracts.ai.relay.types import RelayFormat, RelayLoss, RelayUsage
19from lexigram.contracts.core.result import Result
21__all__ = [
22 "StreamEmitter",
23 "StreamNormalizer",
24 "StreamSession",
25 "StreamSnapshot",
26 "StreamToolCallRecord",
27]
30def _pre_text(state: StreamSnapshot, delta: StreamDelta) -> str:
31 """Reconstruct the text accumulated before a delta was applied."""
32 if delta.kind == "content" and delta.content and state.text.endswith(delta.content):
33 return state.text[: -len(delta.content)]
34 return state.text
37def _pre_thinking(state: StreamSnapshot, delta: StreamDelta) -> str:
38 """Reconstruct the thinking text before a delta was applied."""
39 if (
40 delta.kind == "thinking"
41 and delta.thinking_delta
42 and state.thinking_text.endswith(delta.thinking_delta)
43 ):
44 return state.thinking_text[: -len(delta.thinking_delta)]
45 return state.thinking_text
48def _first_tool_contribution(state: StreamSnapshot, delta: StreamDelta) -> bool:
49 """Whether a tool-call delta is the first contribution for its index.
51 A tool block/item opens on the first delta that carries any fragment
52 for a source-side index. Later fragments change the accumulated
53 record, so equality against the joined record identifies the first.
54 """
55 index = delta.tool_call_index
56 if index is None:
57 return False
58 record: StreamToolCallRecord | None = None
59 for candidate in state.tool_calls:
60 if candidate.index == index:
61 record = candidate
62 break
63 if record is None:
64 return False
65 contributed = (
66 delta.tool_call_id is not None
67 or delta.tool_call_name is not None
68 or delta.tool_call_arguments is not None
69 )
70 return contributed and (
71 record.id == (delta.tool_call_id or "")
72 and record.name == (delta.tool_call_name or "")
73 and record.arguments == (delta.tool_call_arguments or "")
74 )
77def _pre_tool_indices(state: StreamSnapshot, delta: StreamDelta) -> set[int]:
78 """Tool call indices that existed before a delta was applied."""
79 indices = {record.index for record in state.tool_calls}
80 if delta.kind == "tool_call" and delta.tool_call_index is not None:
81 if _first_tool_contribution(state, delta):
82 indices.discard(delta.tool_call_index)
83 return indices
86def _started_pre(state: StreamSnapshot, delta: StreamDelta) -> bool:
87 """Whether the stream had accumulated any state before a delta."""
88 if state.role is not None and delta.kind != "role":
89 return True
90 if _pre_text(state, delta):
91 return True
92 if _pre_thinking(state, delta):
93 return True
94 if _pre_tool_indices(state, delta):
95 return True
96 if state.usage is not None and state.usage is not delta.usage:
97 return True
98 if state.finish_reason is not None and not (
99 delta.kind == "finish" and state.finish_reason == delta.finish_reason
100 ):
101 return True
102 return state.status is not None and not (
103 delta.kind == "status" and state.status == delta.status
104 )
107@runtime_checkable
108class StreamNormalizer(Protocol):
109 """A source mapper's ``stream_to_delta`` viewed as a callable.
111 Maps one source wire event into zero or more canonical
112 ``StreamDelta`` objects. The session hands the current accumulated
113 ``StreamState`` so a mapper can ground decisions on prior events.
114 """
116 def __call__(
117 self, event: Any, *, state: StreamState
118 ) -> Result[tuple[StreamDelta, ...], RelayError]:
119 """Convert one source wire event into canonical deltas."""
120 ...
123@runtime_checkable
124class StreamEmitter(Protocol):
125 """A target emitter's ``delta_to_stream`` viewed as a callable.
127 Maps one canonical ``StreamDelta`` into zero or more target wire
128 events. The session hands the accumulated :class:`StreamSnapshot`
129 so the emitter can stamp metadata and close blocks with full data.
130 """
132 def __call__(
133 self, delta: StreamDelta, *, state: StreamSnapshot
134 ) -> Result[tuple[Any, ...], RelayError]:
135 """Convert one canonical delta into target wire events."""
136 ...
139@dataclass(frozen=True)
140class StreamToolCallRecord:
141 """Accumulated state of one tool call inside a stream session.
143 Attributes:
144 index: Source-side tool call index (stable across fragments).
145 id: Fragmented id, joined verbatim so far.
146 name: Fragmented function name, joined verbatim so far.
147 arguments: Fragmented JSON argument text, kept raw until the
148 stream ends. Invalid/incomplete JSON is never parsed.
149 """
151 index: int
152 id: str = ""
153 name: str = ""
154 arguments: str = ""
157@dataclass(frozen=True)
158class StreamSnapshot:
159 """Read-only snapshot of a stream session.
161 Attributes:
162 source: Upstream wire format.
163 target: Downstream wire format.
164 model: Model name stamped on emitted chunks.
165 stream_id: Upstream stream id, or ``None``.
166 created: Epoch seconds the stream started, or ``None``.
167 include_usage: Whether a final usage event is requested.
168 text: Accumulated assistant text.
169 role: Announced assistant role, or ``None``.
170 thinking_text: Accumulated thinking/reasoning text.
171 thinking_signatures: Thinking signatures taken verbatim, in order.
172 tool_calls: Accumulated tool call fragments by source index.
173 usage: Latest usage seen (may be the only content of a chunk),
174 or ``None``.
175 finish_reason: Raw finish reason from a terminal event, or ``None``.
176 status: Last target status value, or ``None``.
177 open_blocks: Content block indices opened but not yet closed.
178 next_output_index: Next target output index to assign.
179 is_done: Whether the stream has been finalized.
180 warnings: Human-readable loss messages accumulated so far.
181 losses: Machine-readable ``RelayLoss`` records accumulated while
182 converting streams.
183 """
185 source: RelayFormat
186 target: RelayFormat
187 model: str
188 stream_id: str | None
189 created: int | None
190 include_usage: bool
191 text: str
192 role: str | None
193 thinking_text: str
194 thinking_signatures: tuple[str, ...]
195 tool_calls: tuple[StreamToolCallRecord, ...]
196 usage: RelayUsage | None
197 finish_reason: str | None
198 status: str | None
199 open_blocks: tuple[int, ...]
200 next_output_index: int
201 is_done: bool
202 warnings: tuple[str, ...] = ()
203 losses: tuple[RelayLoss, ...] = ()
206class StreamSession:
207 """One upstream stream's mutable conversion state.
209 The session buffers fragmented tool calls by source index, preserves
210 invalid/incomplete argument JSON verbatim until the stream ends, and
211 remembers the latest usage event (including usage-only terminal
212 chunks). Mutable state stays private; callers observe it through
213 :meth:`snapshot` and never through the protocol.
214 """
216 def __init__(
217 self,
218 *,
219 source: RelayFormat,
220 target: RelayFormat,
221 model: str,
222 normalizer: StreamNormalizer,
223 emitter: StreamEmitter,
224 stream_id: str | None = None,
225 created: int | None = None,
226 include_usage: bool = False,
227 ) -> None:
228 self._source = source
229 self._target = target
230 self._model = model
231 self._stream_id = stream_id
232 self._created = created
233 self._include_usage = include_usage
234 self.normalizer = normalizer
235 self.emitter = emitter
237 self._text = ""
238 self._role: str | None = None
239 self._thinking_parts: list[str] = []
240 self._thinking_signatures: list[str] = []
241 self._tool_calls: dict[int, StreamToolCallRecord] = {}
242 self._usage: RelayUsage | None = None
243 self._finish_reason: str | None = None
244 self._status: str | None = None
245 self._open_blocks: list[int] = []
246 self._next_output_index = 0
247 self._finalized = False
248 self._warnings: list[str] = []
249 self._losses: list[RelayLoss] = []
251 def accept(self, event: Any) -> tuple[Any, ...]:
252 """Accept one source wire event and emit target events.
254 Args:
255 event: One source wire event (DTO or raw dict per mapper).
257 Returns:
258 Zero, one, or many target wire events emitted for this event.
260 Raises:
261 RelayError: Wrong source format (``stream_state_invalid``),
262 already finalized (``stream_already_finalized``), or a
263 malformed event.
264 """
265 if self._finalized:
266 raise stream_already_finalized(
267 f"cannot accept event on finalized stream {self._stream_id!r}"
268 )
269 state = self._stream_state()
270 result = self.normalizer(event, state=state)
271 if result.is_err():
272 error = result.unwrap_err()
273 if error.code == RelayErrorCode.UNSUPPORTED_FORMAT.value:
274 raise stream_state_invalid(str(error))
275 raise error
276 emitted: list[Any] = []
277 for delta in result.unwrap():
278 self._apply(delta)
279 emission = self.emitter(delta, state=self.snapshot())
280 if emission.is_err():
281 raise emission.unwrap_err()
282 emitted.extend(emission.unwrap())
283 return tuple(emitted)
285 def finalize(self) -> tuple[Any, ...]:
286 """Close the stream deterministically and return terminal events.
288 A stream that never saw a terminal event is closed with a safe
289 ``finish``/``stop`` delta; a requested usage event is appended
290 when usage was observed. Repeated calls return an empty tuple
291 without further mutation.
293 Returns:
294 Target terminal events; empty when already finalized.
295 """
296 if self._finalized:
297 return ()
298 self._finalized = True
299 terminal: list[Any] = []
300 if self._finish_reason is None:
301 safe_stop = StreamDelta(kind="finish", finish_reason="stop")
302 self._finish_reason = "stop"
303 terminal.extend(self._emit(safe_stop))
304 if self._include_usage and self._usage is not None:
305 usage_delta = StreamDelta(kind="usage", usage=self._usage)
306 terminal.extend(self._emit(usage_delta))
307 return tuple(terminal)
309 def snapshot(self) -> StreamSnapshot:
310 """Return a read-only snapshot of the accumulated session state."""
311 return StreamSnapshot(
312 source=self._source,
313 target=self._target,
314 model=self._model,
315 stream_id=self._stream_id,
316 created=self._created,
317 include_usage=self._include_usage,
318 text=self._text,
319 role=self._role,
320 thinking_text=self._thinking_text(),
321 thinking_signatures=tuple(self._thinking_signatures),
322 tool_calls=tuple(
323 self._tool_calls[index] for index in self._tool_call_order()
324 ),
325 usage=self._usage,
326 finish_reason=self._finish_reason,
327 status=self._status,
328 open_blocks=tuple(self._open_blocks),
329 next_output_index=self._next_output_index,
330 is_done=self._finalized,
331 warnings=tuple(self._warnings),
332 losses=tuple(self._losses),
333 )
335 def _stream_state(self) -> StreamState:
336 """Build the immutable mapper-facing state for this call."""
337 return StreamState(
338 source=self._source,
339 target=self._target,
340 model=self._model,
341 include_usage=self._include_usage,
342 tool_calls=[],
343 thinking_signatures=list(self._thinking_signatures),
344 is_done=self._finalized,
345 usage=self._usage,
346 )
348 def record_loss(
349 self, *, field: str, reason: str, severity: str = "warning"
350 ) -> None:
351 """Record a semantic loss and its rendered warning on the session.
353 Emitters call this when the target format cannot represent a
354 source feature. The snapshot exposes the accumulated records and
355 warnings; nothing is raised.
357 Args:
358 field: Source wire field (or feature) that was adapted.
359 reason: Machine-readable reason (e.g. ``thinking_not_supported``).
360 severity: ``error``, ``warning``, or ``info``.
361 """
362 loss = RelayLoss(
363 field=field, target=self._target, reason=reason, severity=severity
364 )
365 self._losses.append(loss)
366 self._warnings.append(f"{field}: {reason} ({self._target.value}, {severity})")
368 def _thinking_text(self) -> str:
369 return "".join(self._thinking_parts)
371 def _tool_call_order(self) -> list[int]:
372 return list(self._tool_calls)
374 def _apply(self, delta: StreamDelta) -> None:
375 """Fold one canonical delta into the accumulated state."""
376 if delta.kind == "content":
377 if delta.content:
378 self._text += delta.content
379 if delta.block_index is not None:
380 if delta.block_index not in self._open_blocks:
381 self._open_blocks.append(delta.block_index)
382 if delta.output_index is not None:
383 self._next_output_index = max(
384 self._next_output_index, delta.output_index + 1
385 )
386 elif delta.kind == "role":
387 if delta.role is not None:
388 self._role = delta.role
389 elif delta.kind == "thinking":
390 if delta.thinking_delta:
391 self._thinking_parts.append(delta.thinking_delta)
392 signature = delta.passthrough.get("signature")
393 if isinstance(signature, str) and signature:
394 self._thinking_signatures.append(signature)
395 elif delta.kind == "tool_call":
396 if delta.tool_call_index is None:
397 return
398 record = self._tool_calls.setdefault(
399 delta.tool_call_index,
400 StreamToolCallRecord(index=delta.tool_call_index),
401 )
402 if delta.tool_call_id:
403 record = StreamToolCallRecord(
404 index=record.index,
405 id=record.id + delta.tool_call_id,
406 name=record.name,
407 arguments=record.arguments,
408 )
409 if delta.tool_call_name:
410 record = StreamToolCallRecord(
411 index=record.index,
412 id=record.id,
413 name=record.name + delta.tool_call_name,
414 arguments=record.arguments,
415 )
416 if delta.tool_call_arguments:
417 record = StreamToolCallRecord(
418 index=record.index,
419 id=record.id,
420 name=record.name,
421 arguments=record.arguments + delta.tool_call_arguments,
422 )
423 self._tool_calls[delta.tool_call_index] = record
424 elif delta.kind == "usage":
425 if delta.usage is not None:
426 self._usage = delta.usage
427 elif delta.kind == "finish":
428 if delta.finish_reason is not None:
429 self._finish_reason = delta.finish_reason
430 elif delta.kind == "status":
431 if delta.status is not None:
432 self._status = delta.status
433 if delta.usage is not None:
434 self._usage = delta.usage
436 def _emit(self, delta: StreamDelta) -> tuple[Any, ...]:
437 """Route one delta through the target emitter."""
438 emission = self.emitter(delta, state=self.snapshot())
439 if emission.is_err():
440 raise emission.unwrap_err()
441 return tuple(emission.unwrap())