rlmness-cli 0.2.0__py3-none-any.whl

rlmagent/model/_oai/parsers.py

@@ -0,0 +1,317 @@
+"""Streaming response decoders for Chat Completions and Responses APIs."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Protocol
+
+from rlmagent.model._oai.helpers import stream_failure
+
+# ---------------------------------------------------------------------------
+# Parsed signal types
+# ---------------------------------------------------------------------------
+
+
+@dataclass(frozen=True, slots=True)
+class TextChunk:
+    """An incremental fragment of assistant text content."""
+
+    text: str
+
+
+@dataclass(frozen=True, slots=True)
+class ReasoningChunk:
+    """An incremental fragment of reasoning / chain-of-thought output."""
+
+    text: str
+
+
+@dataclass(frozen=True, slots=True)
+class CallBegin:
+    """Marks the start of a new tool call at the given parser index."""
+
+    index: int
+    call_id: str
+    name: str
+
+
+@dataclass(frozen=True, slots=True)
+class CallArgFragment:
+    """An incremental fragment of a tool call's JSON argument string."""
+
+    index: int
+    fragment: str
+
+
+@dataclass(frozen=True, slots=True)
+class CallResolved:
+    """A fully reconstructed tool call ready for execution."""
+
+    index: int
+    call_id: str
+    name: str
+    arguments_json: str
+
+
+@dataclass(frozen=True, slots=True)
+class Metadata:
+    """Response-level metadata received at the start of a stream."""
+
+    response_id: str | None = None
+    model: str | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class Finished:
+    """Marks the end of the model response."""
+
+    finish_reason: str | None = None
+    usage: dict[str, Any] | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class ParseFault:
+    """Signals a parsing error in the response stream."""
+
+    message: str
+
+
+type ParsedSignal = (
+    TextChunk
+    | ReasoningChunk
+    | CallBegin
+    | CallArgFragment
+    | CallResolved
+    | Metadata
+    | Finished
+    | ParseFault
+)
+
+
+# ---------------------------------------------------------------------------
+# Tool-call argument accumulator
+# ---------------------------------------------------------------------------
+
+
+@dataclass(slots=True)
+class _ToolBuffer:
+    """Accumulates incremental argument fragments for one tool call."""
+
+    index: int
+    call_id: str
+    name: str
+    fragments: list[str] = field(default_factory=list)
+
+    @property
+    def arguments_json(self) -> str:
+        return "".join(self.fragments)
+
+
+# ---------------------------------------------------------------------------
+# Decoder protocol
+# ---------------------------------------------------------------------------
+
+
+class StreamDecoder(Protocol):
+    """Common interface implemented by both API-specific decoders."""
+
+    def decode(self, event_name: str, data: dict[str, Any]) -> list[ParsedSignal]:
+        """Process one parsed SSE payload and return any resulting signals."""
+        ...
+
+
+# ---------------------------------------------------------------------------
+# Chat Completions decoder
+# ---------------------------------------------------------------------------
+
+
+def _failure_text(code: str | None, message: str | None) -> str:
+    """Human-readable text for a provider-reported stream failure."""
+    if message:
+        return message
+    if code:
+        return f"Response failed: {code}"
+    return "Provider returned an error"
+
+
+class ChatDecoder:
+    """Decodes Chat Completions streaming chunks into parsed signals."""
+
+    def __init__(self) -> None:
+        self._tools: dict[int, _ToolBuffer] = {}
+
+    def decode(self, event_name: str, data: dict[str, Any]) -> list[ParsedSignal]:
+        # Top-level error object (non-streaming error response)
+        failure = stream_failure(event_name, data)
+        if failure is not None:
+            return [ParseFault(message=_failure_text(*failure))]
+
+        signals: list[ParsedSignal] = []
+
+        # Response metadata
+        resp_id = data.get("id")
+        resp_model = data.get("model")
+        if resp_id or resp_model:
+            signals.append(Metadata(response_id=resp_id, model=resp_model))
+
+        choices = data.get("choices")
+        if not choices:
+            # Usage-only chunk (sent with stream_options.include_usage)
+            usage = data.get("usage")
+            if usage:
+                signals.append(Finished(usage=usage))
+            return signals
+
+        choice = choices[0]
+        delta = choice.get("delta", {})
+        finish_reason = choice.get("finish_reason")
+
+        # Reasoning content (extended-thinking models). When one chunk carries
+        # both fields, the reasoning precedes the answer it leads to.
+        reasoning = delta.get("reasoning_content")
+        if reasoning:
+            signals.append(ReasoningChunk(text=reasoning))
+
+        # Text content
+        content = delta.get("content")
+        if content:
+            signals.append(TextChunk(text=content))
+
+        # Tool call deltas
+        for tc in delta.get("tool_calls", []):
+            idx = tc.get("index", 0)
+            tc_id = tc.get("id")
+            func = tc.get("function", {})
+            name = func.get("name")
+            args_fragment = func.get("arguments", "")
+
+            if tc_id and name:
+                # First chunk — open a new tool buffer
+                buf = _ToolBuffer(index=idx, call_id=tc_id, name=name)
+                self._tools[idx] = buf
+                signals.append(CallBegin(index=idx, call_id=tc_id, name=name))
+                if args_fragment:
+                    buf.fragments.append(args_fragment)
+                    signals.append(CallArgFragment(index=idx, fragment=args_fragment))
+            elif args_fragment and idx in self._tools:
+                # Continuation — accumulate arguments
+                self._tools[idx].fragments.append(args_fragment)
+                signals.append(CallArgFragment(index=idx, fragment=args_fragment))
+
+        # Stream completion
+        if finish_reason is not None:
+            # Resolve any pending tool calls
+            for buf in self._tools.values():
+                signals.append(CallResolved(
+                    index=buf.index,
+                    call_id=buf.call_id,
+                    name=buf.name,
+                    arguments_json=buf.arguments_json,
+                ))
+            self._tools.clear()
+
+            usage = data.get("usage")
+            signals.append(Finished(finish_reason=finish_reason, usage=usage))
+
+        return signals
+
+
+# ---------------------------------------------------------------------------
+# Responses API decoder
+# ---------------------------------------------------------------------------
+
+
+class ResponsesDecoder:
+    """Decodes Responses API streaming events into parsed signals."""
+
+    def __init__(self) -> None:
+        self._tools: dict[int, _ToolBuffer] = {}
+        self._reasoned = False
+
+    def _resolve(self, idx: int, final: object) -> list[ParsedSignal]:
+        """Close the call at *idx*. A non-empty final string is authoritative;"""
+        buf = self._tools.pop(idx, None)
+        if buf is None:
+            return []
+        arguments = final if isinstance(final, str) and final else buf.arguments_json
+        return [CallResolved(
+            index=buf.index,
+            call_id=buf.call_id,
+            name=buf.name,
+            arguments_json=arguments,
+        )]
+
+    def decode(self, event_name: str, data: dict[str, Any]) -> list[ParsedSignal]:
+        signals: list[ParsedSignal] = []
+
+        match event_name:
+            # --- lifecycle ----------------------------------------------------
+            case "response.created":
+                resp = data.get("response", {})
+                signals.append(Metadata(
+                    response_id=resp.get("id"),
+                    model=resp.get("model"),
+                ))
+
+            # --- text content -------------------------------------------------
+            case "response.output_text.delta":
+                delta = data.get("delta", "")
+                if delta:
+                    signals.append(TextChunk(text=delta))
+
+            # --- reasoning content -------------------------------------------
+            case "response.reasoning_summary_text.delta":
+                delta = data.get("delta", "")
+                if delta:
+                    self._reasoned = True
+                    signals.append(ReasoningChunk(text=delta))
+
+            case "response.reasoning_summary_part.done":
+                # Summary parts are separate paragraphs.
+                if self._reasoned:
+                    signals.append(ReasoningChunk(text="\n\n"))
+
+            # --- tool calls ---------------------------------------------------
+            case "response.output_item.added":
+                item = data.get("item", {})
+                idx = data.get("output_index", 0)
+                if item.get("type") == "function_call":
+                    call_id = item.get("call_id", "")
+                    fn_name = item.get("name", "")
+                    self._tools[idx] = _ToolBuffer(
+                        index=idx, call_id=call_id, name=fn_name,
+                    )
+                    signals.append(CallBegin(index=idx, call_id=call_id, name=fn_name))
+
+            case "response.function_call_arguments.delta":
+                idx = data.get("output_index", 0)
+                delta = data.get("delta", "")
+                if delta and idx in self._tools:
+                    self._tools[idx].fragments.append(delta)
+                    signals.append(CallArgFragment(index=idx, fragment=delta))
+
+            case "response.function_call_arguments.done":
+                signals += self._resolve(data.get("output_index", 0), data.get("arguments"))
+
+            case "response.output_item.done":
+                # Some servers finish a call here without an arguments `done`.
+                item = data.get("item", {})
+                if item.get("type") == "function_call":
+                    signals += self._resolve(data.get("output_index", 0), item.get("arguments"))
+
+            # --- completion ---------------------------------------------------
+            case "response.completed":
+                for idx in sorted(self._tools):
+                    signals += self._resolve(idx, None)
+                resp = data.get("response", {})
+                usage = resp.get("usage")
+                status = resp.get("status", "completed")
+                reason = "stop" if status == "completed" else status
+                signals.append(Finished(finish_reason=reason, usage=usage))
+
+            # --- errors -------------------------------------------------------
+            case "error" | "response.failed":
+                failure = stream_failure(event_name, data) or (None, None)
+                signals.append(ParseFault(message=_failure_text(*failure)))
+
+        return signals

rlmagent/model/_oai/payloads.py

@@ -0,0 +1,226 @@
+"""Request payload construction for Chat Completions and Responses APIs."""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Sequence
+from typing import Any
+
+from rlmagent.harness.contracts.tooling import ToolSpec
+from rlmagent.harness.contracts.transcript import (
+    CallBlock,
+    ExtensionEntry,
+    ForkSummaryEntry,
+    HumanEntry,
+    ImageSegment,
+    ModelEntry,
+    PruneSummaryEntry,
+    ShellResultEntry,
+    TextSegment,
+    ToolOutcomeEntry,
+    TranscriptEntry,
+    entry_to_human,
+)
+from rlmagent.model.correlation import wire_call_id
+
+# ---------------------------------------------------------------------------
+# Tool schema conversion
+# ---------------------------------------------------------------------------
+
+
+def _chat_function_schema(tool: ToolSpec) -> dict[str, Any]:
+    """Convert a ``ToolSpec`` to a Chat Completions function tool definition."""
+    return {
+        "type": "function",
+        "function": {
+            "name": tool.name,
+            "description": tool.description,
+            "parameters": dict(tool.parameters),
+        },
+    }
+
+
+def _responses_function_schema(tool: ToolSpec) -> dict[str, Any]:
+    """Convert a ``ToolSpec`` to a Responses API function tool definition."""
+    return {
+        "type": "function",
+        "name": tool.name,
+        "description": tool.description,
+        "parameters": dict(tool.parameters),
+    }
+
+
+# ---------------------------------------------------------------------------
+# Content helpers
+# ---------------------------------------------------------------------------
+
+
+def _user_blocks(entry: HumanEntry) -> str | list[dict[str, Any]]:
+    """Convert user content to the provider content format."""
+    if isinstance(entry.content, str):
+        return entry.content
+    parts: list[dict[str, Any]] = []
+    for seg in entry.content:
+        if isinstance(seg, TextSegment):
+            parts.append({"type": "text", "text": seg.text})
+        elif isinstance(seg, ImageSegment):
+            parts.append({
+                "type": "image_url",
+                "image_url": {"url": f"data:{seg.mime_type};base64,{seg.data}"},
+            })
+    return parts or ""
+
+
+def _serialize_call(call: CallBlock) -> dict[str, Any]:
+    """Serialise a ``CallBlock`` to the Chat Completions tool_calls format."""
+    return {
+        "id": wire_call_id(call.id),
+        "type": "function",
+        "function": {
+            "name": call.name,
+            "arguments": json.dumps(call.arguments),
+        },
+    }
+
+
+# ---------------------------------------------------------------------------
+# Chat Completions message conversion
+# ---------------------------------------------------------------------------
+
+
+def _chat_assistant_msg(entry: ModelEntry) -> dict[str, Any]:
+    """Convert a ``ModelEntry`` to a Chat Completions assistant message."""
+    msg: dict[str, Any] = {"role": "assistant"}
+    text_parts = [b.text for b in entry.content if isinstance(b, TextSegment)]
+    msg["content"] = "".join(text_parts) or None
+    calls = [_serialize_call(b) for b in entry.content if isinstance(b, CallBlock)]
+    if calls:
+        msg["tool_calls"] = calls
+    return msg
+
+
+def _chat_tool_result_msg(entry: ToolOutcomeEntry) -> dict[str, Any]:
+    """Convert a ``ToolOutcomeEntry`` to a Chat Completions tool message."""
+    return {
+        "role": "tool",
+        "tool_call_id": wire_call_id(entry.tool_call_id),
+        "content": entry.text or "",
+    }
+
+
+def _to_chat_message(entry: TranscriptEntry) -> dict[str, Any] | None:
+    """Convert a single transcript entry to a Chat Completions message dict."""
+    if isinstance(entry, HumanEntry):
+        return {"role": "user", "content": _user_blocks(entry)}
+    if isinstance(entry, ModelEntry):
+        return _chat_assistant_msg(entry)
+    if isinstance(entry, ToolOutcomeEntry):
+        return _chat_tool_result_msg(entry)
+    if isinstance(entry, (ExtensionEntry, ShellResultEntry, PruneSummaryEntry, ForkSummaryEntry)):
+        return {"role": "user", "content": _user_blocks(entry_to_human(entry))}
+    return None
+
+
+# ---------------------------------------------------------------------------
+# Responses API message conversion
+# ---------------------------------------------------------------------------
+
+
+def _to_responses_items(entry: TranscriptEntry) -> list[dict[str, Any]]:
+    """Convert a single transcript entry to Responses API input items."""
+    if isinstance(entry, HumanEntry):
+        return [{"role": "user", "content": _user_blocks(entry)}]
+
+    if isinstance(entry, ModelEntry):
+        items: list[dict[str, Any]] = []
+        text = entry.text
+        if text:
+            items.append({
+                "role": "assistant",
+                "content": [{"type": "output_text", "text": text}],
+            })
+        for call in entry.tool_calls:
+            items.append({
+                "type": "function_call",
+                "call_id": wire_call_id(call.id),
+                "name": call.name,
+                "arguments": json.dumps(call.arguments),
+            })
+        return items
+
+    if isinstance(entry, ToolOutcomeEntry):
+        return [{
+            "type": "function_call_output",
+            "call_id": wire_call_id(entry.tool_call_id),
+            "output": entry.text or "",
+        }]
+
+    if isinstance(entry, (ExtensionEntry, ShellResultEntry, PruneSummaryEntry, ForkSummaryEntry)):
+        return [{"role": "user", "content": _user_blocks(entry_to_human(entry))}]
+
+    return []
+
+
+# ---------------------------------------------------------------------------
+# Payload assembly
+# ---------------------------------------------------------------------------
+
+
+def build_chat_payload(
+    *,
+    model: str,
+    system: str,
+    messages: Sequence[TranscriptEntry],
+    tools: Sequence[ToolSpec],
+    max_tokens: int | None = None,
+    extra: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    """Assemble a complete Chat Completions request body."""
+    chat_msgs: list[dict[str, Any]] = [{"role": "system", "content": system}]
+    for entry in messages:
+        msg = _to_chat_message(entry)
+        if msg is not None:
+            chat_msgs.append(msg)
+
+    body: dict[str, Any] = {
+        "model": model,
+        "stream": True,
+        "stream_options": {"include_usage": True},
+        "messages": chat_msgs,
+    }
+    if tools:
+        body["tools"] = [_chat_function_schema(t) for t in tools]
+    if max_tokens is not None:
+        body["max_tokens"] = max_tokens
+    if extra:
+        body.update(extra)
+    return body
+
+
+def build_responses_payload(
+    *,
+    model: str,
+    system: str,
+    messages: Sequence[TranscriptEntry],
+    tools: Sequence[ToolSpec],
+    max_tokens: int | None = None,
+    extra: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    """Assemble a complete Responses API request body."""
+    items: list[dict[str, Any]] = []
+    for entry in messages:
+        items.extend(_to_responses_items(entry))
+
+    body: dict[str, Any] = {
+        "model": model,
+        "stream": True,
+        "instructions": system,
+        "input": items,
+    }
+    if tools:
+        body["tools"] = [_responses_function_schema(t) for t in tools]
+    if max_tokens is not None:
+        body["max_output_tokens"] = max_tokens
+    if extra:
+        body.update(extra)
+    return body

rlmagent/model/_oai/transport.py

@@ -0,0 +1,127 @@
+"""Streaming HTTP transport for OpenAI-compatible API endpoints."""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Callable
+
+import httpx
+
+from rlmagent.harness.contracts.tooling import CancelToken
+from rlmagent.model._oai.helpers import ServerSentEvent, parse_retry_after
+
+# ---------------------------------------------------------------------------
+# Retriable status codes
+# ---------------------------------------------------------------------------
+
+_RETRIABLE_CODES = frozenset({408, 409, 425, 429})
+
+
+# ---------------------------------------------------------------------------
+# Error types
+# ---------------------------------------------------------------------------
+
+
+
+class MalformedPayload(ValueError):
+    """Raised when an SSE data line cannot be decoded as a valid JSON object."""
+
+    _TRUNCATE = 200
+
+    def __init__(self, raw: str) -> None:
+        self.raw = raw
+        preview = raw[:self._TRUNCATE] + "…" if len(raw) > self._TRUNCATE else raw
+        super().__init__(f"Invalid JSON payload: {preview}")
+
+
+class HttpStreamError(Exception):
+    """Raised when a streaming HTTP request receives a non-200 status."""
+
+    def __init__(
+        self,
+        status: int,
+        body: str,
+        *,
+        retry_after: float | None = None,
+    ) -> None:
+        self.status = status
+        self.body = body
+        self.retry_after_seconds = retry_after
+        super().__init__(f"HTTP {status}")
+
+    @property
+    def retriable(self) -> bool:
+        """Return whether this status code warrants an automatic reattempt."""
+        return self.status in _RETRIABLE_CODES or self.status >= 500
+
+
+# ---------------------------------------------------------------------------
+# SSE stream reader
+# ---------------------------------------------------------------------------
+
+
+async def _read_sse(response: httpx.Response) -> AsyncIterator[ServerSentEvent]:
+    """Parse an HTTP response body as a Server-Sent Events stream."""
+    event_type = ""
+    data_lines: list[str] = []
+
+    async for raw_line in response.aiter_lines():
+        line = raw_line.rstrip("\r\n")
+
+        if not line:
+            # Empty line — dispatch the accumulated event
+            if data_lines:
+                payload = "\n".join(data_lines)
+                if payload == "[DONE]":
+                    return
+                yield ServerSentEvent(event=event_type, data=payload)
+                event_type = ""
+                data_lines = []
+            continue
+
+        if line.startswith(":"):
+            continue  # comment
+
+        if line.startswith("event:"):
+            event_type = line[6:].strip()
+        elif line.startswith("data:"):
+            data_lines.append(line[5:].lstrip())
+
+
+# ---------------------------------------------------------------------------
+# Public streaming request
+# ---------------------------------------------------------------------------
+
+
+async def open_event_stream(
+    client: httpx.AsyncClient,
+    url: str,
+    payload: dict,
+    headers: dict[str, str],
+    *,
+    signal: CancelToken | None = None,
+    on_accept: Callable[[httpx.Response], None] | None = None,
+) -> AsyncIterator[ServerSentEvent]:
+    """Execute a streaming POST and yield parsed SSE events."""
+    from rlmagent.agents.trace import correlate_headers
+    async with client.stream(
+        "POST",
+        url,
+        json=payload,
+        headers=correlate_headers(headers),
+    ) as response:
+        if response.status_code != 200:
+            body = (await response.aread()).decode("utf-8", errors="replace")
+            raise HttpStreamError(
+                response.status_code,
+                body,
+                retry_after=parse_retry_after(
+                    response.headers.get("retry-after")
+                ),
+            )
+        if on_accept is not None:
+            on_accept(response)
+
+        async for sse in _read_sse(response):
+            if signal is not None and signal.is_cancelled():
+                return
+            yield sse