chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,366 @@
|
||||
"""SSE frame parser — converts raw byte chunks into typed events.
|
||||
|
||||
Handles the ``event:`` / ``data:`` / ``[DONE]`` framing from the
|
||||
server's ``text/event-stream`` responses.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any
|
||||
|
||||
from omnigent.server import schemas as _srv_events
|
||||
|
||||
from ._events import (
|
||||
NATIVE_TOOL_TYPES,
|
||||
ClientTaskCancel,
|
||||
CompactionInProgress,
|
||||
ElicitationRequest,
|
||||
ErrorEvent,
|
||||
MessageDone,
|
||||
NativeToolCall,
|
||||
OutputFileDone,
|
||||
ReasoningDelta,
|
||||
ReasoningStarted,
|
||||
ReasoningSummaryDelta,
|
||||
ResponseCancelled,
|
||||
ResponseCompleted,
|
||||
ResponseCreated,
|
||||
ResponseFailed,
|
||||
ResponseIncomplete,
|
||||
ResponseInProgress,
|
||||
ResponseQueued,
|
||||
RetryEvent,
|
||||
StreamEvent,
|
||||
TextDelta,
|
||||
ToolCall,
|
||||
ToolResult,
|
||||
)
|
||||
from ._types import ErrorInfo, Response
|
||||
|
||||
_log = logging.getLogger("omnigent_client.sse")
|
||||
|
||||
|
||||
def _wire_type(cls: type) -> str:
|
||||
"""
|
||||
Extract the wire ``type`` literal from a server event class.
|
||||
|
||||
Each :mod:`omnigent.server.schemas` event class pins its
|
||||
``type`` field as ``Literal["..."]``; this helper unwraps that
|
||||
to a plain string so the SDK's ``str == str`` dispatch table
|
||||
stays a ``str`` comparison rather than introducing a class-side
|
||||
isinstance check.
|
||||
|
||||
:param cls: A subclass of the server's ``_SSEEventBase``,
|
||||
e.g. :class:`omnigent.server.schemas.OutputTextDeltaEvent`.
|
||||
:returns: The wire ``type`` literal, e.g.
|
||||
``"response.output_text.delta"``.
|
||||
"""
|
||||
return cls.model_fields["type"].annotation.__args__[0] # type: ignore[union-attr]
|
||||
|
||||
|
||||
# Wire type literals — pulled from the server's typed source of
|
||||
# truth so a rename there is a one-edit change here as well.
|
||||
_T_RESPONSE_CREATED = _wire_type(_srv_events.CreatedEvent)
|
||||
_T_RESPONSE_QUEUED = _wire_type(_srv_events.QueuedEvent)
|
||||
_T_RESPONSE_IN_PROGRESS = _wire_type(_srv_events.InProgressEvent)
|
||||
_T_RESPONSE_COMPLETED = _wire_type(_srv_events.CompletedEvent)
|
||||
_T_RESPONSE_FAILED = _wire_type(_srv_events.FailedEvent)
|
||||
_T_RESPONSE_INCOMPLETE = _wire_type(_srv_events.IncompleteEvent)
|
||||
_T_RESPONSE_CANCELLED = _wire_type(_srv_events.CancelledEvent)
|
||||
_T_RESPONSE_OUTPUT_TEXT_DELTA = _wire_type(_srv_events.OutputTextDeltaEvent)
|
||||
_T_RESPONSE_REASONING_STARTED = _wire_type(_srv_events.ReasoningStartedEvent)
|
||||
_T_RESPONSE_REASONING_TEXT_DELTA = _wire_type(_srv_events.ReasoningTextDeltaEvent)
|
||||
_T_RESPONSE_REASONING_SUMMARY_TEXT_DELTA = _wire_type(_srv_events.ReasoningSummaryTextDeltaEvent)
|
||||
_T_RESPONSE_OUTPUT_ITEM_DONE = _wire_type(_srv_events.OutputItemDoneEvent)
|
||||
_T_RESPONSE_OUTPUT_FILE_DONE = _wire_type(_srv_events.OutputFileDoneEvent)
|
||||
_T_RESPONSE_RETRY = _wire_type(_srv_events.RetryEvent)
|
||||
_T_RESPONSE_ERROR = _wire_type(_srv_events.ErrorEvent)
|
||||
_T_RESPONSE_COMPACTION_IN_PROGRESS = _wire_type(_srv_events.CompactionInProgressEvent)
|
||||
_T_RESPONSE_CLIENT_TASK_CANCEL = _wire_type(_srv_events.ClientTaskCancelEvent)
|
||||
_T_RESPONSE_ELICITATION_REQUEST = _wire_type(_srv_events.ElicitationRequestEvent)
|
||||
|
||||
|
||||
async def parse_sse_stream(
|
||||
byte_stream: AsyncIterator[bytes],
|
||||
) -> AsyncIterator[StreamEvent]:
|
||||
"""Parse an SSE byte stream into typed events.
|
||||
|
||||
:param byte_stream: Raw bytes from ``httpx.Response.aiter_bytes()``.
|
||||
:yields: Typed :class:`StreamEvent` instances.
|
||||
"""
|
||||
buf = ""
|
||||
current_event: str | None = None
|
||||
|
||||
async for chunk in byte_stream:
|
||||
buf += chunk.decode("utf-8", errors="replace")
|
||||
while "\n" in buf:
|
||||
line, buf = buf.split("\n", 1)
|
||||
line = line.rstrip("\r")
|
||||
|
||||
if line.startswith("event: "):
|
||||
current_event = line[7:]
|
||||
elif line.startswith("data: "):
|
||||
data_str = line[6:]
|
||||
if data_str.strip() == "[DONE]":
|
||||
# The server's terminal sentinel is a bare ``data: [DONE]``
|
||||
# with no preceding ``event:`` line, so detect it regardless
|
||||
# of ``current_event``.
|
||||
return
|
||||
# A ``data:`` line is only a parseable event when an ``event:``
|
||||
# line preceded it; a lone data line (other than [DONE]) is
|
||||
# ignored.
|
||||
if current_event is not None:
|
||||
try:
|
||||
data = json.loads(data_str)
|
||||
except json.JSONDecodeError:
|
||||
_log.warning("Failed to parse SSE data: %s", data_str[:200])
|
||||
current_event = None
|
||||
continue
|
||||
event = _parse_event(current_event, data)
|
||||
if event is not None:
|
||||
yield event
|
||||
current_event = None
|
||||
elif line == "":
|
||||
current_event = None
|
||||
|
||||
|
||||
def _normalize_event_type(event_type: str) -> str:
|
||||
"""Normalize event type to handle server enum rendering.
|
||||
|
||||
The server builds terminal events as ``f"response.{task.status}"``
|
||||
where ``task.status`` may be a Python enum (rendering as
|
||||
``response.TaskStatus.COMPLETED``) instead of the expected
|
||||
``response.completed``. Normalize by extracting and lowercasing
|
||||
the enum value.
|
||||
"""
|
||||
if ".TaskStatus." in event_type:
|
||||
# "response.TaskStatus.COMPLETED" → "response.completed"
|
||||
parts = event_type.split(".")
|
||||
status = parts[-1].lower()
|
||||
return f"response.{status}"
|
||||
return event_type
|
||||
|
||||
|
||||
def _parse_event(event_type: str, data: dict[str, Any]) -> StreamEvent | None:
|
||||
"""Convert a raw SSE event type + JSON data into a typed event.
|
||||
|
||||
:param event_type: Wire name of the SSE ``event:`` field, e.g.
|
||||
``"response.output_text.delta"``. Compared against the
|
||||
``_T_RESPONSE_*`` wire-type constants (sourced from
|
||||
:mod:`omnigent.server.schemas`) to dispatch.
|
||||
:param data: Decoded JSON payload from the SSE ``data:`` field.
|
||||
:returns: A typed :class:`StreamEvent` for known event names, or
|
||||
``None`` when the payload is missing required fields or the
|
||||
event type is unrecognized (forward-compatible skip).
|
||||
"""
|
||||
event_type = _normalize_event_type(event_type)
|
||||
|
||||
# Response lifecycle
|
||||
if event_type == _T_RESPONSE_CREATED:
|
||||
return ResponseCreated(response=_parse_response(data))
|
||||
if event_type == _T_RESPONSE_QUEUED:
|
||||
return ResponseQueued(response=_parse_response(data))
|
||||
if event_type == _T_RESPONSE_IN_PROGRESS:
|
||||
return ResponseInProgress(response=_parse_response(data))
|
||||
if event_type == _T_RESPONSE_COMPLETED:
|
||||
return ResponseCompleted(response=_parse_response(data))
|
||||
if event_type == _T_RESPONSE_FAILED:
|
||||
return ResponseFailed(response=_parse_response(data))
|
||||
if event_type == _T_RESPONSE_INCOMPLETE:
|
||||
resp = _parse_response(data)
|
||||
reason = ""
|
||||
if resp.incomplete_details is not None:
|
||||
reason = resp.incomplete_details.reason
|
||||
return ResponseIncomplete(response=resp, reason=reason)
|
||||
if event_type == _T_RESPONSE_CANCELLED:
|
||||
return ResponseCancelled(response=_parse_response(data))
|
||||
|
||||
# Text streaming
|
||||
if event_type == _T_RESPONSE_OUTPUT_TEXT_DELTA:
|
||||
delta = data.get("delta")
|
||||
if isinstance(delta, str):
|
||||
return TextDelta(delta=delta)
|
||||
return None
|
||||
|
||||
# Reasoning
|
||||
if event_type == _T_RESPONSE_REASONING_STARTED:
|
||||
return ReasoningStarted()
|
||||
if event_type == _T_RESPONSE_REASONING_TEXT_DELTA:
|
||||
delta = data.get("delta")
|
||||
if isinstance(delta, str):
|
||||
return ReasoningDelta(delta=delta)
|
||||
return None
|
||||
if event_type == _T_RESPONSE_REASONING_SUMMARY_TEXT_DELTA:
|
||||
delta = data.get("delta")
|
||||
if isinstance(delta, str):
|
||||
return ReasoningSummaryDelta(delta=delta)
|
||||
return None
|
||||
|
||||
# Output items
|
||||
if event_type == _T_RESPONSE_OUTPUT_ITEM_DONE:
|
||||
return _parse_output_item(data)
|
||||
|
||||
# File output
|
||||
if event_type == _T_RESPONSE_OUTPUT_FILE_DONE:
|
||||
return OutputFileDone(
|
||||
file_id=str(data.get("file_id", "")),
|
||||
filename=str(data["filename"]) if data.get("filename") is not None else None,
|
||||
content_type=str(data["content_type"])
|
||||
if data.get("content_type") is not None
|
||||
else None,
|
||||
)
|
||||
|
||||
# Retry
|
||||
if event_type == _T_RESPONSE_RETRY:
|
||||
return RetryEvent(
|
||||
source=str(data.get("source", "")),
|
||||
tool_name=str(data["tool_name"]) if data.get("tool_name") is not None else None,
|
||||
attempt=int(data.get("attempt", 0)),
|
||||
max_attempts=int(data.get("max_attempts", 0)),
|
||||
delay_seconds=float(data.get("delay_seconds", 0.0)),
|
||||
error=_parse_error_info(data.get("error", {})),
|
||||
)
|
||||
|
||||
# Error
|
||||
if event_type == _T_RESPONSE_ERROR:
|
||||
return ErrorEvent(
|
||||
source=str(data.get("source", "")),
|
||||
tool_name=str(data["tool_name"]) if data.get("tool_name") is not None else None,
|
||||
error=_parse_error_info(data.get("error", {})),
|
||||
)
|
||||
|
||||
# Compaction
|
||||
if event_type == _T_RESPONSE_COMPACTION_IN_PROGRESS:
|
||||
return CompactionInProgress()
|
||||
|
||||
# Async client-tool cancel notification
|
||||
if event_type == _T_RESPONSE_CLIENT_TASK_CANCEL:
|
||||
task_id = data.get("task_id")
|
||||
if isinstance(task_id, str) and task_id:
|
||||
raw_call_id = data.get("call_id")
|
||||
call_id = raw_call_id if isinstance(raw_call_id, str) and raw_call_id else None
|
||||
return ClientTaskCancel(task_id=task_id, call_id=call_id)
|
||||
_log.warning("response.client_task.cancel missing task_id: %r", data)
|
||||
return None
|
||||
|
||||
# MCP-shape elicitation request (POLICIES.md §7 + the
|
||||
# universal API additions in
|
||||
# designs/SERVER_HARNESS_CONTRACT.md). The ``params`` block
|
||||
# mirrors MCP's ``ElicitRequestFormParams`` field-for-field;
|
||||
# we surface it as a typed event so the stream consumer can
|
||||
# route it to the elicitation hook (not into a ToolHandler).
|
||||
if event_type == "response.elicitation_request":
|
||||
elicitation_id = data.get("elicitation_id")
|
||||
if not isinstance(elicitation_id, str) or not elicitation_id:
|
||||
_log.warning(
|
||||
"response.elicitation_request missing elicitation_id: %r",
|
||||
data,
|
||||
)
|
||||
return None
|
||||
params = data.get("params")
|
||||
if not isinstance(params, dict):
|
||||
_log.warning(
|
||||
"response.elicitation_request missing params dict: %r",
|
||||
data,
|
||||
)
|
||||
return None
|
||||
requested_schema = params.get("requestedSchema")
|
||||
return ElicitationRequest(
|
||||
elicitation_id=elicitation_id,
|
||||
message=str(params.get("message") or ""),
|
||||
# MCP spec restricts requestedSchema to a JSON-Schema
|
||||
# dict — we accept anything dict-shaped and pass it
|
||||
# through; non-dict inputs (defensive against
|
||||
# malformed producers) become an empty schema.
|
||||
requested_schema=requested_schema if isinstance(requested_schema, dict) else {},
|
||||
mode=str(params.get("mode") or "form"),
|
||||
phase=str(params.get("phase") or ""),
|
||||
policy_name=str(params.get("policy_name") or ""),
|
||||
content_preview=str(params.get("content_preview") or ""),
|
||||
url=str(params["url"]) if isinstance(params.get("url"), str) else None,
|
||||
)
|
||||
|
||||
# Unknown event — skip gracefully for forward-compatibility
|
||||
_log.debug("Skipping unknown SSE event type: %s", event_type)
|
||||
return None
|
||||
|
||||
|
||||
def _parse_output_item(data: dict[str, Any]) -> StreamEvent | None:
|
||||
"""Parse a ``response.output_item.done`` event into a typed event."""
|
||||
item = data.get("item")
|
||||
if not isinstance(item, dict):
|
||||
return None
|
||||
|
||||
item_type = item.get("type", "")
|
||||
|
||||
if item_type == "function_call":
|
||||
args_str = str(item.get("arguments", "{}"))
|
||||
try:
|
||||
arguments = json.loads(args_str)
|
||||
except json.JSONDecodeError:
|
||||
arguments = {}
|
||||
name = str(item.get("name", ""))
|
||||
call_id = str(item.get("call_id", ""))
|
||||
return ToolCall(
|
||||
name=name,
|
||||
arguments=arguments,
|
||||
call_id=call_id,
|
||||
status=str(item.get("status", "")),
|
||||
agent_name=str(item.get("model", "")),
|
||||
)
|
||||
|
||||
if item_type == "function_call_output":
|
||||
raw_arguments = item.get("arguments")
|
||||
arguments: dict[str, object] = {}
|
||||
if isinstance(raw_arguments, dict):
|
||||
arguments = raw_arguments
|
||||
elif isinstance(raw_arguments, str) and raw_arguments:
|
||||
try:
|
||||
decoded_arguments = json.loads(raw_arguments)
|
||||
except json.JSONDecodeError:
|
||||
decoded_arguments = {}
|
||||
if isinstance(decoded_arguments, dict):
|
||||
arguments = decoded_arguments
|
||||
return ToolResult(
|
||||
call_id=str(item.get("call_id", "")),
|
||||
output=str(item.get("output", "")),
|
||||
arguments=arguments,
|
||||
)
|
||||
|
||||
if item_type == "message":
|
||||
content = item.get("content", [])
|
||||
return MessageDone(
|
||||
content=content if isinstance(content, list) else [],
|
||||
)
|
||||
|
||||
if item_type in NATIVE_TOOL_TYPES:
|
||||
return NativeToolCall(
|
||||
tool_type=item_type,
|
||||
data=item,
|
||||
)
|
||||
|
||||
# Compaction items, reasoning items, etc. — skip
|
||||
_log.debug("Skipping output item type: %s", item_type)
|
||||
return None
|
||||
|
||||
|
||||
def _parse_response(data: dict[str, Any]) -> Response:
|
||||
"""Extract the Response object from an SSE event payload."""
|
||||
resp_data = data.get("response")
|
||||
if isinstance(resp_data, dict):
|
||||
return Response.from_dict(resp_data)
|
||||
# Some events put fields at the top level
|
||||
return Response.from_dict(data)
|
||||
|
||||
|
||||
def _parse_error_info(raw: Any) -> ErrorInfo:
|
||||
"""Parse an ErrorInfo from a nested dict."""
|
||||
if isinstance(raw, dict):
|
||||
return ErrorInfo(
|
||||
code=str(raw.get("code", "")),
|
||||
message=str(raw.get("message", "")),
|
||||
)
|
||||
return ErrorInfo(code="", message=str(raw))
|
||||
Reference in New Issue
Block a user