rlmness-cli 0.2.0__py3-none-any.whl

rlmagent/model/oai_compatible.py

@@ -0,0 +1,414 @@
+"""OpenAI-compatible provider adapter for the rlm-agent agent harness."""
+
+from __future__ import annotations
+
+import os
+import uuid
+from collections.abc import AsyncIterator, Sequence
+from contextlib import aclosing
+from dataclasses import replace
+
+import httpx
+
+from rlmagent.harness.contracts.tooling import CancelToken, ToolSpec
+from rlmagent.harness.contracts.transcript import ModelEntry, TranscriptEntry
+from rlmagent.harness.provider.wire import StreamFaultEvent, WireEvent
+from rlmagent.model._oai.helpers import (
+    build_reasoning_extra,
+    is_passing_failure,
+    pick_endpoint,
+    stream_failure,
+    try_parse_json,
+)
+from rlmagent.model._oai.normalize import EventAssembler
+from rlmagent.model._oai.parsers import ChatDecoder, ParseFault, ResponsesDecoder, StreamDecoder
+from rlmagent.model._oai.payloads import build_chat_payload, build_responses_payload
+from rlmagent.model._oai.transport import HttpStreamError, MalformedPayload, open_event_stream
+from rlmagent.model.settings import Credential, OpenAIProfile, ReasoningPolicy
+from rlmagent.model.token_bounds import TokenBoundError, TokenReservation
+from rlmagent.model.transport.backoff import RetryBudget, build_retry_event, compute_delay, pause_for_retry
+from rlmagent.model.transport.client import build_async_client
+from rlmagent.model.transport.faults import format_http_error
+
+_FIRST_PARTY_HOST = "api.openai.com"
+# Routers that front several backends name the one that answered here.
+_SERVED_BY_HEADER = "x-inference-provider"
+_ROUTING_KEY_LIMIT = 64
+
+
+class OpenAIProvider:
+    """Streaming model adapter for OpenAI-compatible API endpoints."""
+
+    def __init__(self, profile: OpenAIProfile) -> None:
+        self._retry_budget = RetryBudget()
+        self._profile = profile
+        self.default_service_tier = profile.service_tier
+        self.supports_service_tier = profile.name in ("openai", "codex", "azure")
+        self._client: httpx.AsyncClient = build_async_client(
+            timeout=httpx.Timeout(profile.timeout_seconds),
+        )
+        #: Backend that served the latest accepted response, when reported.
+        self.served_by: str | None = None
+        #: HTTP status behind the latest request that ended in a fault.
+        self.last_failure_status: int | None = None
+        self._sockets = None
+
+    def set_reasoning(self, policy: ReasoningPolicy) -> None:
+        """Swap the reasoning-effort policy for subsequent requests."""
+        self._profile = replace(self._profile, reasoning=policy)
+
+    async def close(self) -> None:
+        """Release the underlying HTTP client resources."""
+        await self._client.aclose()
+        if self._sockets is not None:
+            await self._sockets.close()
+
+    async def _open_events(self, url, payload, headers, *, signal=None, cache_key=None):
+        if self._profile.transport == "websocket":
+            from rlmagent.model.transport.responses_socket import ResponsesSockets
+
+            if not url.endswith("/responses"):
+                raise httpx.ConnectError("WebSocket transport requires Responses")
+            if self._sockets is None:
+                self._sockets = ResponsesSockets(timeout=self._profile.timeout_seconds)
+            async for event in self._sockets.stream(url, payload, headers, signal=signal,
+                                                   affinity=cache_key):
+                yield event
+        else:
+            async for event in open_event_stream(self._client, url, payload, headers,
+                                                 signal=signal, on_accept=self._note_server):
+                yield event
+
+    def set_service_tier(self, tier: str | None) -> None:
+        if tier is not None and not self.supports_service_tier:
+            raise ValueError("The current provider does not expose service tiers")
+        if tier not in (None, "auto", "default", "flex", "priority"):
+            raise ValueError("Use auto, default, flex, priority or off")
+        self._profile = replace(self._profile, service_tier=tier)
+
+    # --- ModelProvider protocol ---------------------------------------------
+
+    def stream_response(
+        self,
+        *,
+        model: str,
+        system: str,
+        messages: Sequence[TranscriptEntry],
+        tools: Sequence[ToolSpec],
+        signal: CancelToken | None = None,
+        cache_key: str | None = None,
+    ) -> AsyncIterator[WireEvent]:
+        """Stream wire events for one model round-trip."""
+        return self._stream(
+            model=model,
+            system=system,
+            messages=messages,
+            tools=tools,
+            signal=signal,
+            cache_key=cache_key,
+        )
+
+    async def prepare_token_request(self, *, model, system, messages, tools, remaining,
+                                    signal=None, cache_key=None):
+        if self._profile.name == "codex":
+            raise TokenBoundError("Codex subscription does not accept a provider-enforced output cap")
+        if self._profile.name != "openai" or self._profile.transport != "http":
+            raise TokenBoundError("strict token reservations require first-party Responses over HTTP")
+        cred = await self._resolve_credential()
+        base = cred.base_url.rstrip("/")
+        if base != "https://api.openai.com/v1" or (self._profile.endpoint or pick_endpoint(model, base)) != "responses":
+            raise TokenBoundError("strict token reservations require the first-party Responses endpoint")
+        payload = self._responses_payload(model, system, messages, tools)
+        counted = {key: value for key, value in payload.items() if key != "stream"}
+        try:
+            response = await self._client.post(f"{base}/responses/input_tokens", json=counted,
+                                               headers=self._build_headers(cred))
+            response.raise_for_status()
+            result = response.json()
+        except (httpx.HTTPError, ValueError) as exc:
+            raise TokenBoundError("Responses did not supply an input token count") from exc
+        count = result.get("input_tokens") if isinstance(result, dict) else None
+        if type(count) is not int or count < 0:
+            raise TokenBoundError("Responses returned an invalid input token count")
+        output = min(8192, remaining - count)
+        if output < 16:
+            raise TokenBoundError(f"{count} input tokens leave less than 16 tokens for a reply")
+        payload.update(max_output_tokens=output, store=False)
+
+        def start():
+            return self._stream(model=model, system=system, messages=messages, tools=tools,
+                                signal=signal, cache_key=cache_key, prepared=(cred, payload))
+
+        return TokenReservation(count, output, start)
+
+    def _responses_payload(self, model, system, messages, tools):
+        reasoning = build_reasoning_extra(self._profile.reasoning, model=model, endpoint="responses")
+        return build_responses_payload(model=model, system=system, messages=messages,
+                                       tools=tools, extra=reasoning or None)
+
+    # --- internals ----------------------------------------------------------
+
+    async def _resolve_credential(self) -> Credential:
+        """Obtain authentication material, refreshing if a resolver is configured."""
+        if self._profile.resolver is not None:
+            return await self._profile.resolver()
+        return self._profile.credential
+
+    def _build_headers(self, cred: Credential) -> dict[str, str]:
+        """Assemble HTTP headers from the credential and profile."""
+        headers: dict[str, str] = {"Content-Type": "application/json"}
+        if cred.api_key:
+            headers["Authorization"] = f"Bearer {cred.api_key}"
+        if self._profile.organization:
+            headers["OpenAI-Organization"] = self._profile.organization
+        for key, value in cred.extra_headers:
+            headers[key] = value
+        return headers
+
+    async def _stream(
+        self,
+        *,
+        model: str,
+        system: str,
+        messages: Sequence[TranscriptEntry],
+        tools: Sequence[ToolSpec],
+        signal: CancelToken | None,
+        cache_key: str | None = None,
+        prepared=None,
+    ) -> AsyncIterator[WireEvent]:
+        """Execute the streaming request with retry and normalisation."""
+        self.served_by = None
+        self.last_failure_status = None
+        cred = prepared[0] if prepared is not None else await self._resolve_credential()
+        base_url = cred.base_url.rstrip("/")
+        endpoint = "responses" if prepared is not None else self._profile.endpoint or pick_endpoint(model, base_url)
+
+        if endpoint == "responses":
+            url = f"{base_url}/responses"
+            payload = prepared[1] if prepared is not None else self._responses_payload(model, system, messages, tools)
+        else:
+            url = f"{base_url}/chat/completions"
+            reasoning_extra = build_reasoning_extra(self._profile.reasoning, model=model, endpoint=endpoint)
+            payload = build_chat_payload(
+                model=model,
+                system=system,
+                messages=messages,
+                tools=tools,
+                extra=reasoning_extra or None,
+            )
+
+        routing_key = self._routing_key(cache_key, base_url)
+        if self._profile.name == "codex":
+            payload["store"] = False
+        if self._profile.service_tier is not None:
+            payload["service_tier"] = self._profile.service_tier
+        if routing_key is not None:
+            payload["prompt_cache_key"] = routing_key
+
+        retry = self._profile.retry
+        if prepared is not None:
+            retry = replace(retry, max_retries=0)
+        max_attempts = retry.max_retries + 1
+        # One key per logical call, kept across its retries, so a proxy in
+        # front of the provider can drop a duplicate of a request it served.
+        # Opt-in: the providers' own clients do not send it on these
+        # endpoints, so nothing shows the providers accept it.
+        idempotency_key = (
+            uuid.uuid4().hex if (self._profile.idempotency if self._profile.idempotency is not None
+                                else os.environ.get("RLM_AGENT_IDEMPOTENCY") == "1") else None
+        )
+        logical_request = idempotency_key or uuid.uuid4().hex
+
+        for attempt in range(max_attempts):
+            if signal is not None and signal.is_cancelled():
+                yield self._abort_event(model, endpoint)
+                return
+
+            # Refresh credentials on retries
+            if attempt > 0:
+                if self._profile.resolver is not None:
+                    cred = await self._profile.resolver()
+                headers = self._build_headers(cred)
+            else:
+                headers = self._build_headers(cred)
+            if routing_key is not None and endpoint == "responses":
+                headers["session_id"] = routing_key
+            if idempotency_key is not None:
+                headers["Idempotency-Key"] = idempotency_key
+            headers["X-Stainless-Retry-Count"] = str(attempt)
+            headers["X-Rlmagent-Attempt-Id"] = f"{logical_request}:{attempt}"
+            from rlmagent.agents.trace import model_attempt
+
+            model_attempt(attempt=attempt, id=headers["X-Rlmagent-Attempt-Id"],
+                          transport=self._profile.transport, phase="started")
+
+            try:
+                decoder: StreamDecoder
+                if endpoint == "chat":
+                    decoder = ChatDecoder()
+                else:
+                    decoder = ResponsesDecoder()
+
+                assembler = EventAssembler(
+                    model=model,
+                    provider=self._profile.name,
+                    api=endpoint,
+                )
+
+                # Once anything reaches the caller the attempt is committed: a
+                # later failure is surfaced, never replayed, so visible output
+                # and tool calls cannot be duplicated.
+                committed = False
+                passing_failure = False
+                failure_kind = ""
+                try:
+                    async with aclosing(self._open_events(
+                        url, payload, headers, signal=signal,
+                        cache_key=cache_key,
+                    )) as events:
+                        async for sse in events:
+                            data = try_parse_json(sse.data)
+                            if data is None:
+                                raise MalformedPayload(sse.data)
+                            # Some compatible servers omit the `event:` line;
+                            # the payload's own `type` names the event then.
+                            name = sse.event or str(data.get("type") or "")
+                            if not committed and attempt < retry.max_retries:
+                                failure = stream_failure(name, data)
+                                if failure is not None and is_passing_failure(*failure):
+                                    passing_failure = True
+                                    failure_kind = failure[0] or "stream error"
+                                    break
+                            for sig in decoder.decode(name, data):
+                                for wire_event in assembler.accept(sig):
+                                    committed = True
+                                    yield wire_event
+                except MalformedPayload as exc:
+                    for wire_event in assembler.accept(ParseFault(message=str(exc))):
+                        yield wire_event
+
+                if not passing_failure:
+                    if signal is not None and signal.is_cancelled():
+                        yield self._abort_event(model, endpoint)
+                        return
+                    for wire_event in assembler.finalize():
+                        yield wire_event
+                    return  # success
+
+                delay = compute_delay(attempt, max_delay_seconds=retry.max_delay_seconds)
+                yield build_retry_event(
+                    attempt=attempt, max_retries=retry.max_retries,
+                    delay_seconds=delay, reason=f"a transient {failure_kind}",
+                )
+                still_alive = await pause_for_retry(delay, signal=signal)
+                if not still_alive:
+                    yield self._abort_event(model, endpoint)
+                    return
+                continue
+
+            except HttpStreamError as exc:
+                if not exc.retriable or attempt >= retry.max_retries or not self._retry_budget.take():
+                    self.last_failure_status = exc.status
+                    yield self._http_error_event(model, endpoint, exc)
+                    return
+
+                delay = compute_delay(
+                    attempt, max_delay_seconds=retry.max_delay_seconds,
+                )
+                if exc.retry_after_seconds is not None:
+                    delay = max(delay, exc.retry_after_seconds)
+
+                yield build_retry_event(
+                    attempt=attempt, max_retries=retry.max_retries,
+                    delay_seconds=delay, reason=f"HTTP {exc.status}",
+                )
+                still_alive = await pause_for_retry(delay, signal=signal)
+                if not still_alive:
+                    yield self._abort_event(model, endpoint)
+                    return
+
+            except (httpx.TimeoutException, httpx.ConnectError, httpx.ReadError) as exc:
+                if committed or attempt >= retry.max_retries or not self._retry_budget.take():
+                    yield self._network_error_event(model, endpoint, exc)
+                    return
+
+                delay = compute_delay(
+                    attempt, max_delay_seconds=retry.max_delay_seconds,
+                )
+                yield build_retry_event(
+                    attempt=attempt, max_retries=retry.max_retries,
+                    delay_seconds=delay, reason="a network error",
+                )
+                still_alive = await pause_for_retry(delay, signal=signal)
+                if not still_alive:
+                    yield self._abort_event(model, endpoint)
+                    return
+
+    def _note_server(self, response: httpx.Response) -> None:
+        self.served_by = response.headers.get(_SERVED_BY_HEADER) or None
+
+    def _routing_key(self, cache_key: str | None, base_url: str) -> str | None:
+        """The prompt-cache routing value for this request, if it should be sent."""
+        if not cache_key:
+            return None
+        enabled = self._profile.cache_affinity
+        if enabled is None:
+            enabled = httpx.URL(base_url).host == _FIRST_PARTY_HOST
+        # The API caps the key length; the session id is clamped to fit.
+        return cache_key[:_ROUTING_KEY_LIMIT] if enabled else None
+
+    # --- error event builders -----------------------------------------------
+
+    def _abort_event(self, model: str, endpoint: str) -> StreamFaultEvent:
+        """Build an abort event for a cancelled request."""
+        entry = ModelEntry(
+            model=model,
+            provider=self._profile.name,
+            api=endpoint,
+            content=[],
+            stop_reason="aborted",
+        )
+        return StreamFaultEvent(reason="aborted", error=entry)
+
+    def _http_error_event(
+        self,
+        model: str,
+        endpoint: str,
+        exc: HttpStreamError,
+    ) -> StreamFaultEvent:
+        """Build a fault event from an HTTP error response."""
+        msg = format_http_error(
+            provider_name=self._profile.name,
+            status_code=exc.status,
+            body=exc.body,
+            model=model,
+        )
+        entry = ModelEntry(
+            model=model,
+            provider=self._profile.name,
+            api=endpoint,
+            content=[],
+            stop_reason="error",
+            error_message=msg,
+        )
+        return StreamFaultEvent(reason="error", error=entry)
+
+    def _network_error_event(
+        self,
+        model: str,
+        endpoint: str,
+        exc: Exception,
+    ) -> StreamFaultEvent:
+        """Build a fault event from a network-level exception."""
+        entry = ModelEntry(
+            model=model,
+            provider=self._profile.name,
+            api=endpoint,
+            content=[],
+            stop_reason="error",
+            error_message=f"{self._profile.name} network error: {exc}",
+        )
+        return StreamFaultEvent(reason="error", error=entry)
+
+
+__all__ = ["OpenAIProvider"]

rlmagent/model/preferences.py

@@ -0,0 +1,56 @@
+"""User-owned model cycles and per-model request settings, without credentials."""
+
+import json
+from uuid import uuid4
+
+from rlmagent.discovery import _resolve_home
+from rlmagent.refinement.store import _lock
+
+
+def read(provider):
+    try:
+        data = json.loads((_resolve_home() / "model-preferences.json").read_text(encoding="utf-8"))
+    except FileNotFoundError:
+        data = {}
+    return data.get(provider, {"cycle": [], "overrides": {}})
+
+
+def change(provider, operation):
+    path = _resolve_home() / "model-preferences.json"
+    with _lock(path.with_suffix(".lock")):
+        try:
+            data = json.loads(path.read_text(encoding="utf-8"))
+        except FileNotFoundError:
+            data = {}
+        selected = data.setdefault(provider, {"cycle": [], "overrides": {}})
+        operation(selected)
+        temporary = path.with_suffix(f".{uuid4().hex}.tmp")
+        try:
+            temporary.write_text(json.dumps(data), encoding="utf-8")
+            temporary.replace(path)
+        finally:
+            temporary.unlink(missing_ok=True)
+    return selected
+
+
+def tier(provider, model, value):
+    change(
+        provider, lambda selected: selected["overrides"].setdefault(model, {}).update(tier=value)
+    )
+
+
+def cycle(provider, operation, model=None):
+    def update(selected):
+        models = selected["cycle"]
+        if operation == "add":
+            if not model or len(model) > 256:
+                raise ValueError("Supply a model identifier")
+            if model not in models:
+                if len(models) >= 100:
+                    raise ValueError("Model cycle limit reached")
+                models.append(model)
+        elif operation == "remove":
+            if model in models:
+                models.remove(model)
+
+    return change(provider, update)["cycle"]

rlmagent/model/registered_oauth.py

@@ -0,0 +1,195 @@
+"""OAuth for a user's registered xAI public client, without borrowed client IDs."""
+
+import asyncio
+import base64
+import hashlib
+import json
+import math
+import secrets
+import threading
+import time
+from urllib.parse import urlsplit
+
+from rlmagent.discovery import _resolve_home
+from rlmagent.model.settings import ConfigError, Credential
+from rlmagent.model.subscription import _callback, _store
+from rlmagent.model.transport.client import build_async_client
+from rlmagent.refinement.store import _lock
+
+NAME = "xai-subscription"
+BASE = "https://api.x.ai/v1"
+
+
+def registration(value):
+    required = {
+        "client_id",
+        "authorization_endpoint",
+        "token_endpoint",
+        "scope",
+        "resource",
+        "redirect_port",
+    }
+    if not isinstance(value, dict) or value.keys() != required:
+        raise ConfigError("Registered OAuth client settings have missing or unknown fields")
+    for key in ("client_id", "scope"):
+        if not isinstance(value[key], str) or not value[key].strip():
+            raise ConfigError("Registered OAuth client ID and scope must be nonempty")
+    if type(value["redirect_port"]) is not int or not 0 <= value["redirect_port"] <= 65535:
+        raise ConfigError("Invalid OAuth loopback port")
+    for key in ("authorization_endpoint", "token_endpoint"):
+        if not isinstance(value[key], str):
+            raise ConfigError("Registered OAuth endpoints must be strings")
+        uri = urlsplit(value[key])
+        host = uri.hostname or ""
+        allowed = host in ("x.ai", "grok.com") or host.endswith((".x.ai", ".grok.com"))
+        if (
+            not allowed
+            or uri.scheme != "https"
+            or uri.username
+            or uri.password
+            or uri.query
+            or uri.fragment
+        ):
+            raise ConfigError("Registered OAuth endpoints must be pinned HTTPS provider endpoints")
+    if value["resource"] is not None and value["resource"] != BASE:
+        raise ConfigError("OAuth resource must match the public inference endpoint")
+    return dict(value)
+
+
+def _tokens(value, old):
+    access = value.get("access_token")
+    if (
+        value.get("token_type", "").lower() != "bearer"
+        or not isinstance(access, str)
+        or not access
+        or any(char.isspace() for char in access)
+    ):
+        raise ConfigError("Provider returned an invalid access token")
+    lifetime = value.get("expires_in")
+    if lifetime is not None and (
+        type(lifetime) not in (int, float) or not math.isfinite(lifetime) or lifetime <= 0
+    ):
+        raise ConfigError("Provider returned an invalid token lifetime")
+    refresh = value.get("refresh_token", old.get("refresh_token"))
+    if refresh is not None and (not isinstance(refresh, str) or not refresh):
+        raise ConfigError("Provider returned an invalid refresh token")
+    return {
+        "registration": old["registration"],
+        "access_token": access,
+        "refresh_token": refresh,
+        "expires_at": time.time() + lifetime if lifetime is not None else None,
+    }
+
+
+async def login(path=None):
+    def selected():
+        with _lock(_resolve_home() / "xai-oauth.lock"):
+            return _store("load", namespace=NAME)
+
+    old = await asyncio.to_thread(selected)
+    if path is not None:
+        settings = registration(json.loads(path.read_text(encoding="utf-8")))
+    elif old:
+        settings = registration(old["registration"])
+    else:
+        raise ConfigError(
+            "Supply --client-registration with your provider-issued public client settings"
+        )
+    verifier = secrets.token_urlsafe(64)
+    params = {
+        "client_id": settings["client_id"],
+        "response_type": "code",
+        "scope": settings["scope"],
+        "state": secrets.token_urlsafe(32),
+        "code_challenge_method": "S256",
+        "code_challenge": base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest())
+        .rstrip(b"=")
+        .decode(),
+    }
+    if settings["resource"]:
+        params["resource"] = settings["resource"]
+    cancelled = threading.Event()
+    try:
+        callback, redirect = await asyncio.to_thread(
+            _callback,
+            settings["authorization_endpoint"],
+            params,
+            cancelled=cancelled,
+            port=settings["redirect_port"],
+        )
+        if callback.get("client_id", settings["client_id"]) != settings["client_id"]:
+            raise ConfigError("Provider returned a different registered client")
+        fields = {
+            "grant_type": "authorization_code",
+            "client_id": settings["client_id"],
+            "code": callback["code"],
+            "code_verifier": verifier,
+            "redirect_uri": redirect,
+        }
+        if settings["resource"]:
+            fields["resource"] = settings["resource"]
+        async with build_async_client(timeout=20, follow_redirects=False) as client:
+            reply = await client.post(settings["token_endpoint"], data=fields)
+            reply.raise_for_status()
+            record = _tokens(reply.json(), {"registration": settings})
+
+        def save():
+            with _lock(_resolve_home() / "xai-oauth.lock"):
+                _store("save", record, namespace=NAME)
+
+        await asyncio.to_thread(save)
+    except ConfigError:
+        raise
+    except Exception:
+        raise ConfigError(
+            "Registered OAuth sign-in failed; saved credentials were retained"
+        ) from None
+    finally:
+        cancelled.set()
+
+
+def _resolve():
+    with _lock(_resolve_home() / "xai-oauth.lock"):
+        record = _store("load", namespace=NAME)
+        if not record or not record.get("access_token"):
+            raise ConfigError(
+                "Run provider login xai-subscription with your registered client settings"
+            )
+        settings = registration(record["registration"])
+        if record["expires_at"] is not None and record["expires_at"] < time.time() + 60:
+            if not record["refresh_token"]:
+                raise ConfigError("Registered OAuth login expired; sign in again")
+
+            async def refresh():
+                fields = {
+                    "grant_type": "refresh_token",
+                    "client_id": settings["client_id"],
+                    "refresh_token": record["refresh_token"],
+                }
+                if settings["resource"]:
+                    fields["resource"] = settings["resource"]
+                async with build_async_client(timeout=20, follow_redirects=False) as client:
+                    reply = await client.post(settings["token_endpoint"], data=fields)
+                    reply.raise_for_status()
+                    return _tokens(reply.json(), record)
+
+            try:
+                record = asyncio.run(refresh())
+                _store("save", record, namespace=NAME)
+            except Exception:
+                raise ConfigError("Registered OAuth refresh failed; sign in again") from None
+        return Credential(record["access_token"], BASE)
+
+
+async def resolve():
+    return await asyncio.to_thread(_resolve)
+
+
+async def logout():
+    def clear():
+        with _lock(_resolve_home() / "xai-oauth.lock"):
+            record = _store("load", namespace=NAME)
+            if record:
+                _store("save", {"registration": record["registration"]}, namespace=NAME)
+
+    await asyncio.to_thread(clear)