427 lines
16 KiB
Python
427 lines
16 KiB
Python
#!/usr/bin/env python3
|
|
"""Consume bounded model SSE streams without reusing partial responses.
|
|
|
|
The module is intentionally provider-neutral. It understands the event
|
|
envelope used by the Responses API, but does not perform network I/O. A
|
|
caller passes the byte chunks it receives and gets a terminal result plus an
|
|
atomic checkpoint. Partial model text is counted for diagnostics and is
|
|
never written to the checkpoint or returned as reusable context.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import codecs
|
|
import hashlib
|
|
import json
|
|
import os
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Iterable
|
|
|
|
from validate_model_budget import ModelBudgetError, validate_provider_response
|
|
|
|
|
|
DEFAULT_MAX_EVENT_BYTES = 256 * 1024
|
|
DEFAULT_MAX_RESPONSE_BYTES = 8 * 1024 * 1024
|
|
MAX_ERROR_TEXT = 1024
|
|
PROFILE_RETRY = {
|
|
"milestone": "contract",
|
|
"contract": "feature",
|
|
"feature": "doc",
|
|
"doc": "doc",
|
|
}
|
|
|
|
|
|
class StreamProtocolError(ValueError):
|
|
"""Raised when a stream cannot be decoded as bounded SSE."""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SSEEvent:
|
|
event: str
|
|
data: str
|
|
event_id: str | None = None
|
|
|
|
|
|
class SSEDecoder:
|
|
"""Incrementally decode UTF-8 SSE frames with hard size limits."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
max_event_bytes: int = DEFAULT_MAX_EVENT_BYTES,
|
|
max_response_bytes: int = DEFAULT_MAX_RESPONSE_BYTES,
|
|
) -> None:
|
|
if max_event_bytes <= 0 or max_response_bytes <= 0:
|
|
raise ValueError("stream byte limits must be positive")
|
|
if max_event_bytes > max_response_bytes:
|
|
raise ValueError("event limit cannot exceed response limit")
|
|
self.max_event_bytes = int(max_event_bytes)
|
|
self.max_response_bytes = int(max_response_bytes)
|
|
self.bytes_received = 0
|
|
self._decoder = codecs.getincrementaldecoder("utf-8")("strict")
|
|
self._line_buffer = ""
|
|
self._data: list[str] = []
|
|
self._event = "message"
|
|
self._event_id: str | None = None
|
|
self._event_bytes = 0
|
|
|
|
def _dispatch(self) -> SSEEvent | None:
|
|
if not self._data:
|
|
self._event = "message"
|
|
self._event_id = None
|
|
self._event_bytes = 0
|
|
return None
|
|
event = SSEEvent(self._event, "\n".join(self._data), self._event_id)
|
|
self._data = []
|
|
self._event = "message"
|
|
self._event_id = None
|
|
self._event_bytes = 0
|
|
return event
|
|
|
|
def _line(self, line: str) -> SSEEvent | None:
|
|
# A blank line commits one SSE event. Ignore comments and unknown
|
|
# fields, as required by the SSE framing rules.
|
|
if line == "":
|
|
return self._dispatch()
|
|
self._event_bytes += len(line.encode("utf-8")) + 1
|
|
if self._event_bytes > self.max_event_bytes:
|
|
raise StreamProtocolError("SSE event exceeds byte budget")
|
|
if line.startswith(":"):
|
|
return None
|
|
field, separator, value = line.partition(":")
|
|
if separator and value.startswith(" "):
|
|
value = value[1:]
|
|
if field == "event":
|
|
self._event = value or "message"
|
|
elif field == "data":
|
|
self._data.append(value)
|
|
elif field == "id":
|
|
self._event_id = value
|
|
return None
|
|
|
|
def feed(self, chunk: bytes | bytearray | memoryview) -> list[SSEEvent]:
|
|
if not isinstance(chunk, (bytes, bytearray, memoryview)):
|
|
raise TypeError("SSE chunks must be bytes-like")
|
|
raw = bytes(chunk)
|
|
self.bytes_received += len(raw)
|
|
if self.bytes_received > self.max_response_bytes:
|
|
raise StreamProtocolError("SSE response exceeds byte budget")
|
|
try:
|
|
self._line_buffer += self._decoder.decode(raw, final=False)
|
|
except UnicodeDecodeError as exc:
|
|
raise StreamProtocolError("SSE response is not valid UTF-8") from exc
|
|
events: list[SSEEvent] = []
|
|
while True:
|
|
newline = self._line_buffer.find("\n")
|
|
if newline < 0:
|
|
if len(self._line_buffer.encode("utf-8")) > self.max_event_bytes:
|
|
raise StreamProtocolError("SSE line exceeds byte budget")
|
|
break
|
|
line = self._line_buffer[:newline]
|
|
self._line_buffer = self._line_buffer[newline + 1 :]
|
|
if line.endswith("\r"):
|
|
line = line[:-1]
|
|
event = self._line(line)
|
|
if event is not None:
|
|
events.append(event)
|
|
return events
|
|
|
|
def finish(self) -> list[SSEEvent]:
|
|
try:
|
|
self._line_buffer += self._decoder.decode(b"", final=True)
|
|
except UnicodeDecodeError as exc:
|
|
raise StreamProtocolError("SSE response ends with incomplete UTF-8") from exc
|
|
events: list[SSEEvent] = []
|
|
if self._line_buffer:
|
|
line = self._line_buffer
|
|
self._line_buffer = ""
|
|
event = self._line(line.rstrip("\r"))
|
|
if event is not None:
|
|
events.append(event)
|
|
event = self._dispatch()
|
|
if event is not None:
|
|
events.append(event)
|
|
return events
|
|
|
|
|
|
def _json_event(event: SSEEvent) -> dict[str, Any]:
|
|
try:
|
|
value = json.loads(event.data)
|
|
except (TypeError, json.JSONDecodeError) as exc:
|
|
raise StreamProtocolError("SSE data is not a JSON object") from exc
|
|
if not isinstance(value, dict):
|
|
raise StreamProtocolError("SSE data must be a JSON object")
|
|
return value
|
|
|
|
|
|
def _event_type(event: SSEEvent, payload: dict[str, Any]) -> str:
|
|
value = payload.get("type")
|
|
return str(value) if value is not None else event.event
|
|
|
|
|
|
def _nested_response(payload: dict[str, Any]) -> dict[str, Any]:
|
|
response = payload.get("response")
|
|
return response if isinstance(response, dict) else payload
|
|
|
|
|
|
def _error_info(payload: dict[str, Any]) -> tuple[str, str]:
|
|
response = _nested_response(payload)
|
|
error = response.get("error")
|
|
if isinstance(error, dict):
|
|
code = str(error.get("code", "provider_error"))
|
|
message = str(error.get("message", "provider response failed"))
|
|
else:
|
|
code = str(response.get("code", "provider_error"))
|
|
message = str(response.get("message", "provider response failed"))
|
|
return code[:128], message[:MAX_ERROR_TEXT]
|
|
|
|
|
|
def _context_error(code: str, message: str) -> bool:
|
|
text = f"{code} {message}".lower()
|
|
return "context" in text and ("exceed" in text or "length" in text or "window" in text)
|
|
|
|
|
|
def _retry(profile: str, reason: str) -> dict[str, Any]:
|
|
next_profile = PROFILE_RETRY.get(profile, "doc")
|
|
context_overflow = "context" in reason.lower() and (
|
|
"exceed" in reason.lower() or "length" in reason.lower()
|
|
)
|
|
return {
|
|
"action": "split_task" if profile == "doc" and context_overflow else "rebuild_packet",
|
|
"profile": next_profile,
|
|
"reuse_packet": False,
|
|
"drop_optional_evidence": True,
|
|
"reason": reason,
|
|
}
|
|
|
|
|
|
def _checkpoint(path: Path | None, value: dict[str, Any]) -> str | None:
|
|
if path is None:
|
|
return None
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
payload = json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode("utf-8")
|
|
temporary = path.with_name(f".{path.name}.tmp")
|
|
try:
|
|
with temporary.open("wb") as handle:
|
|
handle.write(payload)
|
|
handle.flush()
|
|
os.fsync(handle.fileno())
|
|
os.replace(temporary, path)
|
|
finally:
|
|
try:
|
|
temporary.unlink()
|
|
except FileNotFoundError:
|
|
pass
|
|
return str(path)
|
|
|
|
|
|
def _result(
|
|
*,
|
|
context_id: str,
|
|
packet_id: str,
|
|
profile: str,
|
|
status: str,
|
|
decoder: SSEDecoder,
|
|
event_count: int,
|
|
output_bytes: int,
|
|
output_digest: Any,
|
|
response_id: str | None,
|
|
last_event_id: str | None,
|
|
reason_code: str | None = None,
|
|
error: str | None = None,
|
|
usage: dict[str, int] | None = None,
|
|
finish_reason: str | None = None,
|
|
checkpoint_path: Path | None = None,
|
|
) -> dict[str, Any]:
|
|
reusable = status == "completed"
|
|
result: dict[str, Any] = {
|
|
"schema": "model-stream-result.v1",
|
|
"context_id": context_id,
|
|
"packet_id": packet_id,
|
|
"status": status,
|
|
"reusable": reusable,
|
|
"event_count": event_count,
|
|
"bytes_received": decoder.bytes_received,
|
|
"output_bytes": output_bytes,
|
|
"output_sha256": output_digest.hexdigest(),
|
|
"response_id": response_id,
|
|
"last_event_id": last_event_id,
|
|
"reason_code": reason_code,
|
|
"error": error,
|
|
"finish_reason": finish_reason,
|
|
"retry": None if reusable else _retry(profile, reason_code or status),
|
|
}
|
|
if usage is not None:
|
|
result["usage"] = usage
|
|
checkpoint = {
|
|
"schema": "model-stream-checkpoint.v1",
|
|
"context_id": context_id,
|
|
"packet_id": packet_id,
|
|
"status": status,
|
|
"reusable": reusable,
|
|
"reason_code": reason_code,
|
|
"event_count": event_count,
|
|
"bytes_received": decoder.bytes_received,
|
|
"output_bytes": output_bytes,
|
|
"output_sha256": output_digest.hexdigest(),
|
|
"response_id": response_id,
|
|
"last_event_id": last_event_id,
|
|
"finish_reason": finish_reason,
|
|
"retry": result["retry"],
|
|
}
|
|
if usage is not None:
|
|
checkpoint["usage"] = usage
|
|
result["checkpoint_path"] = _checkpoint(checkpoint_path, checkpoint)
|
|
return result
|
|
|
|
|
|
def consume_sse(
|
|
chunks: Iterable[bytes],
|
|
*,
|
|
context_id: str,
|
|
packet_id: str,
|
|
profile: str = "feature",
|
|
max_output_tokens: int = 1024,
|
|
max_output_bytes: int = DEFAULT_MAX_RESPONSE_BYTES,
|
|
checkpoint_path: Path | None = None,
|
|
max_event_bytes: int = DEFAULT_MAX_EVENT_BYTES,
|
|
max_response_bytes: int = DEFAULT_MAX_RESPONSE_BYTES,
|
|
) -> dict[str, Any]:
|
|
"""Consume a stream and return a terminal, checkpointed result.
|
|
|
|
Any disconnect, malformed frame, provider failure, context overflow, or
|
|
output overrun is non-reusable. The returned retry plan explicitly asks
|
|
the caller to rebuild a smaller packet; no partial output is included.
|
|
"""
|
|
|
|
decoder = SSEDecoder(
|
|
max_event_bytes=max_event_bytes,
|
|
max_response_bytes=max_response_bytes,
|
|
)
|
|
if max_output_tokens < 0 or max_output_bytes <= 0:
|
|
raise ValueError("output limits must be non-negative/positive")
|
|
event_count = 0
|
|
output_bytes = 0
|
|
output_digest = hashlib.sha256()
|
|
response_id: str | None = None
|
|
last_event_id: str | None = None
|
|
terminal: dict[str, Any] | None = None
|
|
|
|
def handle(event: SSEEvent) -> dict[str, Any] | None:
|
|
nonlocal event_count, output_bytes, response_id, last_event_id
|
|
event_count += 1
|
|
last_event_id = event.event_id or last_event_id
|
|
payload = _json_event(event)
|
|
kind = _event_type(event, payload)
|
|
response = _nested_response(payload)
|
|
if response.get("id") is not None:
|
|
response_id = str(response["id"])
|
|
delta = payload.get("delta")
|
|
if isinstance(delta, str):
|
|
encoded = delta.encode("utf-8")
|
|
output_bytes += len(encoded)
|
|
output_digest.update(encoded)
|
|
if output_bytes > max_output_bytes:
|
|
return _result(
|
|
context_id=context_id, packet_id=packet_id, profile=profile,
|
|
status="output_budget_exceeded", decoder=decoder,
|
|
event_count=event_count, output_bytes=output_bytes,
|
|
output_digest=output_digest, response_id=response_id,
|
|
last_event_id=last_event_id, reason_code="output_overrun",
|
|
checkpoint_path=checkpoint_path,
|
|
)
|
|
if kind in {"response.failed", "error"}:
|
|
code, message = _error_info(payload)
|
|
status = "context_budget_exceeded" if _context_error(code, message) else "provider_failed"
|
|
return _result(
|
|
context_id=context_id, packet_id=packet_id, profile=profile,
|
|
status=status, decoder=decoder, event_count=event_count,
|
|
output_bytes=output_bytes, output_digest=output_digest,
|
|
response_id=response_id, last_event_id=last_event_id,
|
|
reason_code=code, error=message, checkpoint_path=checkpoint_path,
|
|
)
|
|
if kind in {"response.completed", "response.incomplete"}:
|
|
try:
|
|
validated = validate_provider_response(
|
|
payload, max_output_tokens=max_output_tokens
|
|
)
|
|
except ModelBudgetError as exc:
|
|
message = str(exc)
|
|
overrun = "exceeded output token budget" in message
|
|
return _result(
|
|
context_id=context_id, packet_id=packet_id, profile=profile,
|
|
status="output_budget_exceeded" if overrun else "protocol_error",
|
|
decoder=decoder,
|
|
event_count=event_count, output_bytes=output_bytes,
|
|
output_digest=output_digest, response_id=response_id,
|
|
last_event_id=last_event_id,
|
|
reason_code="output_overrun" if overrun else "invalid_usage",
|
|
error=message, checkpoint_path=checkpoint_path,
|
|
)
|
|
status = "truncated" if bool(validated.get("truncated")) else "completed"
|
|
return _result(
|
|
context_id=context_id, packet_id=packet_id, profile=profile,
|
|
status=status, decoder=decoder, event_count=event_count,
|
|
output_bytes=output_bytes, output_digest=output_digest,
|
|
response_id=response_id, last_event_id=last_event_id,
|
|
reason_code="length" if status == "truncated" else None,
|
|
usage={key: int(validated[key]) for key in ("input_tokens", "output_tokens", "total_tokens")},
|
|
finish_reason=validated.get("finish_reason"),
|
|
checkpoint_path=checkpoint_path,
|
|
)
|
|
return None
|
|
|
|
try:
|
|
for chunk in chunks:
|
|
for event in decoder.feed(chunk):
|
|
terminal = handle(event)
|
|
if terminal is not None:
|
|
return terminal
|
|
for event in decoder.finish():
|
|
terminal = handle(event)
|
|
if terminal is not None:
|
|
return terminal
|
|
except (ConnectionError, TimeoutError, OSError) as exc:
|
|
return _result(
|
|
context_id=context_id, packet_id=packet_id, profile=profile,
|
|
status="interrupted", decoder=decoder, event_count=event_count,
|
|
output_bytes=output_bytes, output_digest=output_digest,
|
|
response_id=response_id, last_event_id=last_event_id,
|
|
reason_code="stream_disconnected", error=type(exc).__name__,
|
|
checkpoint_path=checkpoint_path,
|
|
)
|
|
except StreamProtocolError as exc:
|
|
return _result(
|
|
context_id=context_id, packet_id=packet_id, profile=profile,
|
|
status="protocol_error", decoder=decoder, event_count=event_count,
|
|
output_bytes=output_bytes, output_digest=output_digest,
|
|
response_id=response_id, last_event_id=last_event_id,
|
|
reason_code="stream_decode_error", error=str(exc),
|
|
checkpoint_path=checkpoint_path,
|
|
)
|
|
except Exception as exc:
|
|
# SDKs use different exception classes for a body that ends before
|
|
# ``response.completed``. Treat all ordinary iterator exceptions as
|
|
# a disconnect, but keep only the class name in the checkpoint.
|
|
return _result(
|
|
context_id=context_id, packet_id=packet_id, profile=profile,
|
|
status="interrupted", decoder=decoder, event_count=event_count,
|
|
output_bytes=output_bytes, output_digest=output_digest,
|
|
response_id=response_id, last_event_id=last_event_id,
|
|
reason_code="stream_disconnected", error=type(exc).__name__,
|
|
checkpoint_path=checkpoint_path,
|
|
)
|
|
|
|
return _result(
|
|
context_id=context_id, packet_id=packet_id, profile=profile,
|
|
status="interrupted", decoder=decoder, event_count=event_count,
|
|
output_bytes=output_bytes, output_digest=output_digest,
|
|
response_id=response_id, last_event_id=last_event_id,
|
|
reason_code="stream_ended_without_completion",
|
|
checkpoint_path=checkpoint_path,
|
|
)
|
|
|
|
|
|
__all__ = ["SSEDecoder", "SSEEvent", "StreamProtocolError", "consume_sse"]
|