rlmness-cli 0.2.0__py3-none-any.whl

rlmagent/model/_oai/helpers.py

@@ -0,0 +1,256 @@
+"""Shared utilities for the OpenAI-compatible provider layer."""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from typing import Any
+
+from rlmagent.harness.contracts.transcript import CostBreakdown, UsageStats
+from rlmagent.model.settings import ReasoningPolicy
+
+
+@dataclass(frozen=True, slots=True)
+class ServerSentEvent:
+    """A single Server-Sent Event parsed from an HTTP stream."""
+
+    event: str
+    data: str
+
+
+# ---------------------------------------------------------------------------
+# Finish-reason mapping
+# ---------------------------------------------------------------------------
+
+
+def map_halt_reason(reason: str | None) -> str:
+    """Translate an OpenAI finish reason to rlm-agent's HaltReason vocabulary."""
+    match reason:
+        case "stop" | None:
+            return "stop"
+        case "length" | "max_tokens":
+            return "length"
+        case "tool_calls" | "function_call":
+            return "toolUse"
+        case _:
+            return "stop"
+
+
+# ---------------------------------------------------------------------------
+# Usage extraction
+# ---------------------------------------------------------------------------
+
+
+def extract_usage(raw: dict[str, Any] | None) -> UsageStats:
+    """Build a ``UsageStats`` from an OpenAI-style usage dictionary."""
+    if not raw:
+        return UsageStats()
+
+    input_tok = raw.get("prompt_tokens") or raw.get("input_tokens") or 0
+    output_tok = raw.get("completion_tokens") or raw.get("output_tokens") or 0
+    total = raw.get("total_tokens") or (int(input_tok) + int(output_tok))
+
+    # Chat Completions and Responses name the prompt breakdown differently.
+    cache_read = cache_write = 0
+    prompt_detail = raw.get("prompt_tokens_details") or raw.get("input_tokens_details")
+    if isinstance(prompt_detail, dict):
+        cache_read = int(prompt_detail.get("cached_tokens", 0) or 0)
+        cache_write = int(prompt_detail.get("cache_write_tokens", 0) or 0)
+
+    reasoning = None
+    comp_detail = raw.get("completion_tokens_details") or raw.get("output_tokens_details")
+    if isinstance(comp_detail, dict):
+        r = comp_detail.get("reasoning_tokens")
+        if r:
+            reasoning = int(r)
+
+    # Gateways such as OpenRouter report the billed cost in USD with usage.
+    cost = CostBreakdown()
+    reported = raw.get("cost")
+    if isinstance(reported, int | float) and not isinstance(reported, bool):
+        details = raw.get("cost_details")
+        details = details if isinstance(details, dict) else {}
+        cost = CostBreakdown(
+            total=float(reported),
+            input=float(details.get("upstream_inference_prompt_cost") or 0.0),
+            output=float(details.get("upstream_inference_completions_cost") or 0.0),
+        )
+
+    # The reported prompt count includes cached tokens; ``input`` keeps only
+    # the fresh part so every provider reports usage the same way.
+    return UsageStats(
+        cost=cost,
+        total_tokens=int(total),
+        input=max(0, int(input_tok) - cache_read - cache_write),
+        output=int(output_tok),
+        cache_read=cache_read,
+        cache_write=cache_write,
+        reasoning=reasoning,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Endpoint selection
+# ---------------------------------------------------------------------------
+
+_RESPONSES_PREFIXES = ("o1", "o3", "o4")
+_DEFAULT_OPENAI_BASE = "https://api.openai.com/v1"
+
+
+def pick_endpoint(model: str, base_url: str) -> str:
+    """Return ``"responses"`` or ``"chat"`` for the given model and base URL."""
+    if base_url.rstrip("/") != _DEFAULT_OPENAI_BASE:
+        return "chat"
+    lower = model.lower()
+    for prefix in _RESPONSES_PREFIXES:
+        if lower.startswith(prefix):
+            return "responses"
+    return "chat"
+
+
+# ---------------------------------------------------------------------------
+# JSON helpers
+# ---------------------------------------------------------------------------
+
+
+def try_parse_json(text: str) -> dict[str, Any] | None:
+    """Attempt to parse *text* as a JSON object, returning ``None`` on failure."""
+    try:
+        obj = json.loads(text)
+        return obj if isinstance(obj, dict) else None
+    except (json.JSONDecodeError, ValueError):
+        return None
+
+
+def parse_tool_arguments(text: str) -> dict[str, Any]:
+    """Parse tool-call argument JSON, returning an empty dict on failure."""
+    if not text or not text.strip():
+        return {}
+    try:
+        obj = json.loads(text)
+        return obj if isinstance(obj, dict) else {}
+    except (json.JSONDecodeError, ValueError):
+        return {}
+
+
+# ---------------------------------------------------------------------------
+# Reasoning parameters
+# ---------------------------------------------------------------------------
+
+
+# Model families that accept a reasoning-effort setting. Other models reject
+# the parameter outright, so it is only sent to these.
+_REASONING_FAMILIES = ("o1", "o3", "o4", "gpt-5", "gpt-oss")
+# The highest level the API names; stronger local levels are clamped to it.
+_EFFORT_CEILING = "high"
+_EFFORTS = ("minimal", "low", "medium", "high")
+
+
+def reasons(model: str) -> bool:
+    """Whether ``model`` takes a reasoning-effort parameter."""
+    name = model.lower().rsplit("/", 1)[-1]
+    return name.startswith(_REASONING_FAMILIES)
+
+
+def build_reasoning_extra(
+    policy: ReasoningPolicy, *, model: str, endpoint: str,
+) -> dict[str, Any]:
+    """Reasoning parameters for one request, shaped for its endpoint."""
+    if not policy.enabled or not reasons(model):
+        return {}
+    effort = policy.effort or ("high" if policy.budget_tokens else "medium")
+    if effort not in _EFFORTS:
+        effort = _EFFORT_CEILING
+    if endpoint == "responses":
+        return {"reasoning": {"effort": effort}}
+    return {"reasoning_effort": effort}
+
+
+# ---------------------------------------------------------------------------
+# Retry-After parsing
+# ---------------------------------------------------------------------------
+
+
+def parse_retry_after(header: str | None) -> float | None:
+    """Extract a wait duration in seconds from a ``Retry-After`` header."""
+    if header is None:
+        return None
+    try:
+        return max(0.0, float(header))
+    except (ValueError, TypeError):
+        return None
+
+
+# ---------------------------------------------------------------------------
+# In-stream failure classification
+# ---------------------------------------------------------------------------
+
+# Substrings (matched against the lower-cased code + message) that mark a
+# failure as a passing condition worth one more attempt.
+_PASSING_FAILURE_HINTS = (
+    "overloaded",
+    "service_unavailable",
+    "temporarily_unavailable",
+    "rate_limit",
+    "internal_error",
+    "server_error",
+    "timeout",
+)
+
+# Rate-limit wording that actually means the account is out of budget; retrying
+# those only burns the backoff budget before failing anyway.
+_EXHAUSTED_QUOTA_HINTS = (
+    "gousagelimiterror",
+    "freeusagelimiterror",
+    "monthly usage limit reached",
+    "available balance",
+    "insufficient_quota",
+    "out of budget",
+    "quota exceeded",
+    "billing",
+)
+
+
+def stream_failure(event: str, data: dict[str, Any]) -> tuple[str | None, str | None] | None:
+    """Return ``(code, message)`` if this SSE event reports a failure, else ``None``."""
+    nested = data.get("error")
+    response = data.get("response")
+    is_failure = (
+        event in ("error", "response.failed")
+        or data.get("type") in ("error", "response.failed")
+        or isinstance(nested, dict)
+    )
+    if not is_failure:
+        return None
+
+    sources: list[dict[str, Any]] = [data]
+    if isinstance(nested, dict):
+        sources.append(nested)
+    if isinstance(response, dict):
+        for key in ("error", "last_error"):
+            found = response.get(key)
+            if isinstance(found, dict):
+                sources.append(found)
+
+    code: str | None = None
+    message: str | None = None
+    for source in sources:
+        raw_message = source.get("message")
+        if message is None and isinstance(raw_message, str) and raw_message:
+            message = raw_message
+        raw_code = source.get("code")
+        if code is None and isinstance(raw_code, str) and raw_code:
+            code = raw_code
+    if code is None and isinstance(nested, dict):
+        nested_type = nested.get("type")
+        if isinstance(nested_type, str) and nested_type:
+            code = nested_type
+    return code, message
+
+
+def is_passing_failure(code: str | None, message: str | None) -> bool:
+    """Whether an in-stream failure looks transient and safe to reissue."""
+    text = " ".join(part for part in (code, message) if part).lower()
+    if not text or any(hint in text for hint in _EXHAUSTED_QUOTA_HINTS):
+        return False
+    return any(hint in text for hint in _PASSING_FAILURE_HINTS)

rlmagent/model/_oai/normalize.py

@@ -0,0 +1,263 @@
+"""Convert parsed stream signals into rlm-agent wire events."""
+
+from __future__ import annotations
+
+from rlmagent.harness.contracts.transcript import (
+    CallBlock,
+    ModelEntry,
+    ReplyContent,
+    TextSegment,
+    ThoughtSegment,
+    UsageStats,
+)
+from rlmagent.harness.provider.wire import (
+    CallChunkEvent,
+    CallCloseEvent,
+    CallOpenEvent,
+    CompletionCause,
+    ContentChunkEvent,
+    ContentCloseEvent,
+    ContentOpenEvent,
+    FaultCause,
+    ReasoningChunkEvent,
+    ReasoningCloseEvent,
+    ReasoningOpenEvent,
+    StreamCloseEvent,
+    StreamFaultEvent,
+    StreamOpenEvent,
+    WireEvent,
+)
+from rlmagent.model._oai.helpers import extract_usage, map_halt_reason, parse_tool_arguments
+from rlmagent.model._oai.parsers import (
+    CallArgFragment,
+    CallBegin,
+    CallResolved,
+    Finished,
+    Metadata,
+    ParsedSignal,
+    ParseFault,
+    ReasoningChunk,
+    TextChunk,
+)
+
+
+def _segment(kind: str, body: str) -> ReplyContent:
+    return TextSegment(text=body) if kind == "text" else ThoughtSegment(thinking=body)
+
+
+class EventAssembler:
+    """Accumulates parsed signals and emits rlm-agent wire events."""
+
+    def __init__(self, *, model: str, provider: str, api: str) -> None:
+        self._model = model
+        self._provider = provider
+        self._api = api
+        self._response_id: str | None = None
+        self._response_model: str | None = None
+
+        # Ordered sequence of completed content blocks
+        self._blocks: list[ReplyContent] = []
+
+        # Open text/reasoning blocks by channel, in first-seen order. Chat
+        # Completions sends reasoning and answer as independent fields that can
+        # interleave (even within one chunk), so there each channel keeps one
+        # stable block until a boundary. Other APIs stream blocks in sequence,
+        # so opening one channel closes the other.
+        self._open: dict[str, tuple[int, list[str]]] = {}
+        self._independent = api == "chat"
+
+        # Usage and halt
+        self._usage: UsageStats = UsageStats()
+        self._halt: str = "stop"
+
+        # Content-index tracking
+        self._next_idx = 0
+        self._call_idx_map: dict[int, int] = {}  # parser index → content index
+
+        # State flags
+        self._opened = False
+        self._faulted = False
+
+    # --- snapshot -----------------------------------------------------------
+
+    def _snapshot(self) -> ModelEntry:
+        """Build a ``ModelEntry`` reflecting the current accumulated state."""
+        blocks: list[ReplyContent] = list(self._blocks)
+
+        # Append partial content from open blocks, in content-index order
+        for kind, (_idx, chunks) in sorted(self._open.items(), key=lambda kv: kv[1][0]):
+            body = "".join(chunks)
+            if body:
+                blocks.append(_segment(kind, body))
+
+        return ModelEntry(
+            model=self._response_model or self._model,
+            provider=self._provider,
+            api=self._api,
+            response_model=self._response_model,
+            response_id=self._response_id,
+            content=blocks,
+            stop_reason=self._halt,  # type: ignore[arg-type]
+            usage=self._usage,
+        )
+
+    def _ensure_open(self) -> list[WireEvent]:
+        """Emit ``StreamOpenEvent`` if the stream hasn't started yet."""
+        if self._opened:
+            return []
+        self._opened = True
+        return [StreamOpenEvent(partial=self._snapshot())]
+
+    # --- signal dispatch ----------------------------------------------------
+
+    def accept(self, signal: ParsedSignal) -> list[WireEvent]:
+        """Process one parsed signal and return any resulting wire events."""
+        if isinstance(signal, Metadata):
+            return self._on_metadata(signal)
+        if isinstance(signal, TextChunk):
+            return self._on_text(signal)
+        if isinstance(signal, ReasoningChunk):
+            return self._on_reasoning(signal)
+        if isinstance(signal, CallBegin):
+            return self._on_call_begin(signal)
+        if isinstance(signal, CallArgFragment):
+            return self._on_call_arg(signal)
+        if isinstance(signal, CallResolved):
+            return self._on_call_resolved(signal)
+        if isinstance(signal, Finished):
+            return self._on_finished(signal)
+        if isinstance(signal, ParseFault):
+            return self._on_fault(signal)
+        return []
+
+    # --- signal handlers ----------------------------------------------------
+
+    def _on_metadata(self, signal: Metadata) -> list[WireEvent]:
+        if signal.response_id:
+            self._response_id = signal.response_id
+        if signal.model:
+            self._response_model = signal.model
+        return []
+
+    def _on_text(self, signal: TextChunk) -> list[WireEvent]:
+        return self._on_channel("text", signal.text)
+
+    def _on_reasoning(self, signal: ReasoningChunk) -> list[WireEvent]:
+        return self._on_channel("reasoning", signal.text)
+
+    def _on_channel(self, kind: str, fragment: str) -> list[WireEvent]:
+        events = self._ensure_open()
+        if kind not in self._open:
+            if not self._independent:
+                events.extend(self._close_active())
+            idx = self._next_idx
+            self._next_idx += 1
+            self._open[kind] = (idx, [])
+            opener = ContentOpenEvent if kind == "text" else ReasoningOpenEvent
+            events.append(opener(content_index=idx, partial=self._snapshot()))
+        idx, chunks = self._open[kind]
+        chunks.append(fragment)
+        chunk_event = ContentChunkEvent if kind == "text" else ReasoningChunkEvent
+        events.append(chunk_event(content_index=idx, delta=fragment, partial=self._snapshot()))
+        return events
+
+    def _on_call_begin(self, signal: CallBegin) -> list[WireEvent]:
+        events = self._ensure_open()
+        events.extend(self._close_active())
+        idx = self._next_idx
+        self._next_idx += 1
+        self._call_idx_map[signal.index] = idx
+        events.append(CallOpenEvent(
+            content_index=idx,
+            partial=self._snapshot(),
+        ))
+        return events
+
+    def _on_call_arg(self, signal: CallArgFragment) -> list[WireEvent]:
+        idx = self._call_idx_map.get(signal.index)
+        if idx is None:
+            return []
+        return [CallChunkEvent(
+            content_index=idx,
+            delta=signal.fragment,
+            partial=self._snapshot(),
+        )]
+
+    def _on_call_resolved(self, signal: CallResolved) -> list[WireEvent]:
+        args = parse_tool_arguments(signal.arguments_json)
+        call = CallBlock(id=signal.call_id, name=signal.name, arguments=args)
+        self._blocks.append(call)
+        idx = self._call_idx_map.get(signal.index)
+        if idx is None:
+            return []
+        return [CallCloseEvent(
+            content_index=idx,
+            tool_call=call,
+            partial=self._snapshot(),
+        )]
+
+    def _on_finished(self, signal: Finished) -> list[WireEvent]:
+        if signal.finish_reason is not None:
+            self._halt = map_halt_reason(signal.finish_reason)
+        if signal.usage is not None:
+            self._usage = extract_usage(signal.usage)
+        return []
+
+    def _on_fault(self, signal: ParseFault) -> list[WireEvent]:
+        self._faulted = True
+        self._halt = "error"
+        events = self._ensure_open()
+        error_entry = ModelEntry(
+            model=self._response_model or self._model,
+            provider=self._provider,
+            api=self._api,
+            response_model=self._response_model,
+            response_id=self._response_id,
+            content=list(self._snapshot().content),
+            stop_reason="error",
+            error_message=signal.message,
+            usage=self._usage,
+        )
+        events.append(StreamFaultEvent(reason="error", error=error_entry))
+        return events
+
+    # --- block closing ------------------------------------------------------
+
+    def _close_active(self) -> list[WireEvent]:
+        """Close every open text/reasoning block, committing them in index order."""
+        events: list[WireEvent] = []
+        for kind, (idx, chunks) in sorted(self._open.items(), key=lambda kv: kv[1][0]):
+            content = "".join(chunks)
+            self._blocks.append(_segment(kind, content))
+            del self._open[kind]
+            closer = ContentCloseEvent if kind == "text" else ReasoningCloseEvent
+            events.append(closer(content_index=idx, content=content, partial=self._snapshot()))
+        return events
+
+    # --- finalisation -------------------------------------------------------
+
+    def finalize(self) -> list[WireEvent]:
+        """Close all open content blocks and emit the terminal event."""
+        if self._faulted:
+            return []
+
+        events: list[WireEvent] = []
+        events.extend(self._close_active())
+
+        message = self._snapshot()
+
+        if not self._opened:
+            events.append(StreamOpenEvent(partial=message))
+
+        if self._halt in ("error", "aborted"):
+            cause: FaultCause = "aborted" if self._halt == "aborted" else "error"
+            events.append(StreamFaultEvent(reason=cause, error=message))
+        else:
+            reason: CompletionCause = "stop"
+            if self._halt == "length":
+                reason = "length"
+            elif self._halt == "toolUse":
+                reason = "toolUse"
+            events.append(StreamCloseEvent(reason=reason, message=message))
+
+        return events