Implement WebAssembly runtime and kinematics coverage
This commit is contained in:
426
scripts/model_stream.py
Normal file
426
scripts/model_stream.py
Normal file
@@ -0,0 +1,426 @@
|
||||
#!/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"]
|
||||
Reference in New Issue
Block a user