rlmness-cli 0.2.0__py3-none-any.whl

rlmagent/model/transport/responses_socket.py

@@ -0,0 +1,223 @@
+"""Persistent, independently locked Responses connections per agent affinity."""
+
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import json
+import time
+import traceback
+from contextlib import suppress
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import httpx
+
+from rlmagent.model._oai.helpers import ServerSentEvent
+from rlmagent.model._oai.transport import HttpStreamError
+from rlmagent.model.transport.client import certificate_context, patched_proxy_environment
+
+
+def canonical(items):
+    result = []
+    for item in items:
+        kind = item.get("type")
+        if kind == "reasoning":
+            continue
+        if item.get("role"):
+            content = item.get("content", [])
+            if isinstance(content, str):
+                content = [
+                    {
+                        "type": "output_text" if item["role"] == "assistant" else "input_text",
+                        "text": content,
+                    }
+                ]
+            result.append(
+                {
+                    "role": item["role"],
+                    "content": [
+                        {
+                            key: value
+                            for key, value in part.items()
+                            if key in ("type", "text", "image_url", "detail")
+                        }
+                        for part in content
+                    ],
+                }
+            )
+        elif kind == "function_call":
+            arguments = item.get("arguments", "{}")
+            try:
+                arguments = json.dumps(json.loads(arguments), sort_keys=True)
+            except (ValueError, TypeError):
+                pass
+            result.append(
+                {
+                    "type": kind,
+                    "call_id": item.get("call_id"),
+                    "name": item.get("name"),
+                    "arguments": arguments,
+                }
+            )
+        else:
+            result.append(item)
+    return result
+
+
+@dataclass
+class Lane:
+    lock: asyncio.Lock = field(default_factory=asyncio.Lock)
+    socket: object = None
+    auth: str | None = None
+    created: float = 0
+    previous: str | None = None
+    prefix: list = field(default_factory=list)
+    settings: dict = field(default_factory=dict)
+
+
+class ResponsesSockets:
+    def __init__(self, *, timeout=120, connect=None):
+        self.timeout = timeout
+        self.connect = connect
+        self.lanes = {}
+        self.last_error = None
+        self.last_error_location = []
+
+    async def _discard(self, lane):
+        if lane.socket is not None:
+            with suppress(Exception):
+                await lane.socket.close()
+        lane.socket, lane.previous, lane.prefix = None, None, []
+
+    async def close(self):
+        for lane in self.lanes.values():
+            await self._discard(lane)
+        self.lanes.clear()
+
+    async def stream(self, url, payload, headers, *, signal=None, affinity=None):
+        try:
+            from websockets.asyncio.client import connect
+            from websockets.exceptions import InvalidStatus
+        except ImportError:
+            raise httpx.ConnectError(
+                "Install rlmagent[subscription] for WebSocket transport"
+            ) from None
+        key = affinity or "root"
+        temporary = key not in self.lanes and len(self.lanes) >= 32
+        lane = Lane() if temporary else self.lanes.setdefault(key, Lane())
+        async with lane.lock:
+            socket_url = str(httpx.URL(url).copy_with(scheme="wss"))
+            auth = {
+                name: value
+                for name, value in headers.items()
+                if name.lower() in ("authorization", "openai-organization", "openai-project")
+            }
+            fingerprint = hashlib.sha256(
+                (socket_url + json.dumps(auth, sort_keys=True)).encode()
+            ).hexdigest()
+            try:
+                if lane.auth != fingerprint or time.monotonic() - lane.created > 2700:
+                    await self._discard(lane)
+                if lane.socket is None:
+                    with patched_proxy_environment():
+                        lane.socket = await (self.connect or connect)(
+                            socket_url,
+                            additional_headers=auth,
+                            ssl=certificate_context(),
+                            open_timeout=20,
+                            max_size=16 * 1024 * 1024,
+                        )
+                    lane.auth, lane.created = fingerprint, time.monotonic()
+                request = {
+                    key: value
+                    for key, value in payload.items()
+                    if key not in ("stream", "background")
+                }
+                original = request.get("input", [])
+                settings = {key: value for key, value in request.items() if key != "input"}
+                same_prefix = canonical(original[: len(lane.prefix)]) == lane.prefix
+                if lane.previous and lane.prefix and same_prefix and settings == lane.settings:
+                    request["previous_response_id"] = lane.previous
+                    request["input"] = original[len(lane.prefix) :]
+                for recovery in range(2):
+                    await lane.socket.send(json.dumps({"type": "response.create", **request}))
+                    visible = False
+                    deadline = time.monotonic() + self.timeout
+                    while True:
+                        receiver = asyncio.create_task(lane.socket.recv())
+                        try:
+                            while not receiver.done():
+                                if signal is not None and signal.is_cancelled():
+                                    await self._discard(lane)
+                                    return
+                                if time.monotonic() >= deadline:
+                                    raise TimeoutError("Responses socket timed out")
+                                await asyncio.wait({receiver}, timeout=0.05)
+                            raw = receiver.result()
+                        finally:
+                            if not receiver.done():
+                                receiver.cancel()
+                                with suppress(asyncio.CancelledError):
+                                    await receiver
+                        event = json.loads(raw)
+                        kind = event.get("type", "")
+                        failure = event.get("error") or event.get("response", {}).get("error") or {}
+                        if (
+                            failure.get("code") == "previous_response_not_found"
+                            and recovery == 0
+                            and not visible
+                            and "previous_response_id" in request
+                        ):
+                            request.pop("previous_response_id")
+                            request["input"] = original
+                            lane.previous, lane.prefix = None, []
+                            from rlmagent.agents.trace import model_attempt
+
+                            model_attempt(
+                                attempt=recovery + 1, transport="websocket", phase="cache_recovery"
+                            )
+                            break
+                        visible = visible or kind in (
+                            "response.output_text.delta",
+                            "response.output_item.added",
+                        )
+                        if kind == "response.incomplete":
+                            event = {
+                                "type": "error",
+                                "error": {"message": "Response was incomplete"},
+                            }
+                            kind = "error"
+                        yield ServerSentEvent(event=kind, data=json.dumps(event))
+                        if kind == "response.completed":
+                            response = event["response"]
+                            lane.previous = response.get("id")
+                            lane.prefix = canonical([*original, *response.get("output", [])])
+                            lane.settings = settings
+                            return
+                        if kind in ("response.failed", "error"):
+                            lane.previous, lane.prefix = None, []
+                            return
+            except InvalidStatus as error:
+                await self._discard(lane)
+                raise HttpStreamError(
+                    error.response.status_code, "WebSocket upgrade rejected"
+                ) from None
+            except asyncio.CancelledError:
+                await self._discard(lane)
+                raise
+            except HttpStreamError:
+                raise
+            except Exception as error:
+                self.last_error = type(error).__name__
+                self.last_error_location = [
+                    f"{Path(frame.filename).name}:{frame.lineno}:{frame.name}"
+                    for frame in traceback.extract_tb(error.__traceback__)
+                ]
+                await self._discard(lane)
+                raise httpx.ReadError(
+                    f"Responses WebSocket disconnected ({self.last_error})"
+                ) from None
+            finally:
+                if temporary:
+                    await self._discard(lane)

rlmagent/model/vertex.py

@@ -0,0 +1,93 @@
+"""Google cloud authentication for the compatible chat endpoint."""
+
+from __future__ import annotations
+
+import asyncio
+import os
+import shutil
+import subprocess
+import time
+from urllib.parse import quote
+
+from rlmagent.model.settings import ConfigError, Credential, OpenAIProfile
+
+
+class ApplicationCredentials:
+    def __init__(self):
+        self.credentials = None
+        self.project = (
+            os.environ.get("GOOGLE_CLOUD_PROJECT")
+            or os.environ.get("CLOUDSDK_CORE_PROJECT")
+            or os.environ.get("GCLOUD_PROJECT")
+        )
+        self.location = os.environ.get("GOOGLE_CLOUD_LOCATION", "us-central1")
+        self._lock = asyncio.Lock()
+        self._token = None
+        self._expires = 0
+
+    def _gcloud(self):
+        if not self.project:
+            raise ConfigError("Set GOOGLE_CLOUD_PROJECT for the cloud endpoint")
+        if self._token is None or time.monotonic() >= self._expires:
+            executable = shutil.which("gcloud")
+            if executable is None:
+                raise ConfigError("Configure application credentials or install gcloud")
+            result = subprocess.run(
+                [executable, "auth", "application-default", "print-access-token"],
+                stdin=subprocess.DEVNULL,
+                capture_output=True,
+                text=True,
+                encoding="utf-8",
+                errors="replace",
+                timeout=20,
+            )
+            if result.returncode != 0 or not result.stdout.strip():
+                raise ConfigError("Cloud login failed; run gcloud auth application-default login")
+            self._token = result.stdout.strip()
+            self._expires = time.monotonic() + 3500
+        return self._endpoint(self._token)
+
+    def _load(self):
+        try:
+            import google.auth
+            from google.auth.transport.requests import Request
+        except ImportError:
+            return self._gcloud()
+        if self.credentials is None:
+            try:
+                self.credentials, detected = google.auth.default(
+                    scopes=["https://www.googleapis.com/auth/cloud-platform"]
+                )
+            except google.auth.exceptions.DefaultCredentialsError:
+                return self._gcloud()
+            self.project = self.project or detected
+        if not self.project:
+            raise ConfigError("Set GOOGLE_CLOUD_PROJECT for the cloud endpoint")
+        if not self.credentials.valid:
+            self.credentials.refresh(Request())
+        return self._endpoint(self.credentials.token)
+
+    def _endpoint(self, token):
+        project, location = quote(self.project, safe=""), quote(self.location, safe="")
+        host = (
+            "aiplatform.googleapis.com"
+            if self.location == "global"
+            else f"{location}-aiplatform.googleapis.com"
+        )
+        return Credential(
+            token,
+            f"https://{host}/v1beta1/projects/{project}/locations/{location}/endpoints/openapi",
+        )
+
+    async def __call__(self):
+        async with self._lock:
+            return await asyncio.to_thread(self._load)
+
+
+def build_vertex():
+    from rlmagent.model.oai_compatible import OpenAIProvider
+
+    resolver = ApplicationCredentials()
+    return OpenAIProvider(
+        OpenAIProfile(name="vertex", credential=Credential(None, ""), resolver=resolver)
+    )

rlmagent/network.py

@@ -0,0 +1,123 @@
+"""Bounded host-side text fetching."""
+
+from __future__ import annotations
+
+import asyncio
+import ipaddress
+import socket
+from html.parser import HTMLParser
+from urllib.parse import urljoin, urlsplit, urlunsplit
+
+
+class _Readable(HTMLParser):
+    """The text of a page, without its scripts, styles and tags."""
+
+    BLOCKS = {"p", "div", "br", "li", "tr", "section", "article", "h1", "h2", "h3", "h4", "pre"}
+
+    def __init__(self) -> None:
+        super().__init__(convert_charrefs=True)
+        self.parts: list[str] = []
+        self.hidden = 0
+
+    def handle_starttag(self, tag, attrs):
+        if tag in ("script", "style", "noscript", "svg"):
+            self.hidden += 1
+        elif tag in self.BLOCKS:
+            self.parts.append("\n")
+
+    def handle_endtag(self, tag):
+        if tag in ("script", "style", "noscript", "svg") and self.hidden:
+            self.hidden -= 1
+
+    def handle_data(self, data):
+        if not self.hidden:
+            self.parts.append(data)
+
+
+def readable(html: str) -> str:
+    page = _Readable()
+    page.feed(html)
+    lines = (" ".join(line.split()) for line in "".join(page.parts).splitlines())
+    return "\n".join(line for line in lines if line)
+
+
+MAX_REDIRECTS = 5
+
+
+def _is_public(address: str) -> bool:
+    ip = ipaddress.ip_address(address.split("%")[0])
+    ip = getattr(ip, "ipv4_mapped", None) or ip
+    return ip.is_global and not ip.is_multicast
+
+
+async def _public_address(host: str) -> str:
+    """An address for `host`, refusing a name where any address is not public."""
+    try:
+        found = await asyncio.get_running_loop().getaddrinfo(host, None, type=socket.SOCK_STREAM)
+    except OSError as error:
+        raise ValueError(f"cannot resolve {host}: {error}") from None
+    addresses = [item[4][0] for item in found]
+    if not addresses or not all(_is_public(a) for a in addresses):
+        raise ValueError(f"{host} is a private or local address; fetch_url only reaches public "
+                         "ones unless rlm-agent runs with --allow-private-fetch")
+    return addresses[0]
+
+
+async def _open(client, url: str, allow_private: bool):
+    """One request. Public hosts are reached at the address just checked, so DNS cannot
+    answer differently for the connection; the name still goes in Host and TLS."""
+    parts = urlsplit(url)
+    if parts.scheme not in {"http", "https"} or not parts.hostname:
+        raise ValueError("fetch_url requires an HTTP or HTTPS URL")
+    if allow_private:
+        return await client.send(client.build_request("GET", url), stream=True)
+    address = await _public_address(parts.hostname)
+    shown = f"[{address}]" if ":" in address else address
+    netloc = f"{shown}:{parts.port}" if parts.port else shown
+    request = client.build_request(
+        "GET", urlunsplit(parts._replace(netloc=netloc)),
+        headers={"Host": parts.netloc.rsplit("@", 1)[-1]},
+        extensions={"sni_hostname": parts.hostname},
+    )
+    return await client.send(request, stream=True)
+
+
+async def fetch_text(url: str, max_chars: int = 50000, raw: bool = False,
+                     allow_private: bool = False) -> dict:
+    import httpx
+
+    limit = max(1, min(int(max_chars), 50000))
+    data = bytearray()
+    truncated = False
+    try:
+        async with httpx.AsyncClient(timeout=20, follow_redirects=False) as client:
+            current = url
+            for _ in range(MAX_REDIRECTS + 1):
+                response = await _open(client, current, allow_private)
+                if not response.is_redirect:
+                    break
+                current = urljoin(current, response.headers["location"])
+                await response.aclose()
+            else:
+                await response.aclose()
+                return {"ok": False, "error": f"more than {MAX_REDIRECTS} redirects"}
+            try:
+                response.raise_for_status()
+                async for chunk in response.aiter_bytes(chunk_size=8192):
+                    remaining = limit * 4 + 1 - len(data)
+                    data.extend(chunk[:remaining])
+                    if len(data) >= limit * 4 + 1:
+                        truncated = True
+                        break
+                text = data.decode(response.encoding or "utf-8", errors="replace")
+                if not raw and "html" in response.headers.get("content-type", ""):
+                    text = readable(text)
+                return {
+                    "url": current,
+                    "text": text[:limit],
+                    "truncated": truncated or len(text) > limit,
+                }
+            finally:
+                await response.aclose()
+    except httpx.HTTPError as error:
+        return {"ok": False, "error": str(error)}

rlmagent/packages.py

@@ -0,0 +1,246 @@
+"""Explicitly installed capability packages and their resource directories."""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import importlib.metadata
+import json
+import re
+import shutil
+import subprocess
+import sys
+import time
+from pathlib import Path
+from uuid import uuid4
+
+from rlmagent.refinement.store import _lock
+
+RETAIN_DAYS = 30
+
+
+def prune(store: Path, records: dict, *, older_than=RETAIN_DAYS, apply=False) -> list[str]:
+    """Collect only expired inactive installs directly owned by this store."""
+    active = {entry["directory"] for entry in records.values()}
+    cutoff = time.time() - older_than * 86400
+    selected = []
+    if not store.exists():
+        return selected
+    for candidate in store.iterdir():
+        if candidate.name in active or not re.fullmatch(r"[a-f0-9]{32}", candidate.name):
+            continue
+        if candidate.is_symlink() or candidate.is_junction() or not candidate.is_dir():
+            continue
+        resolved = _location(candidate.name, store)
+        if str(resolved) in sys.path or candidate.stat().st_mtime >= cutoff:
+            continue
+        selected.append(candidate.name)
+        if apply:
+            shutil.rmtree(resolved)
+    return sorted(selected)
+
+
+def root(*, local: bool = False, paths=None) -> Path:
+    from rlmagent.discovery import _resolve_home
+
+    if local:
+        return (paths.project if paths else Path.cwd()) / ".rlm-agent" / "packages"
+    return (paths.home if paths else _resolve_home()) / "packages"
+
+
+def _records(store: Path | None = None) -> dict:
+    store = store or root()
+    try:
+        entries = json.loads((store / "installed.json").read_text(encoding="utf-8"))
+    except FileNotFoundError:
+        return {}
+    if not isinstance(entries, dict):
+        raise ValueError("Invalid installed-package registry")
+    for name, value in entries.items():
+        if not re.fullmatch(r"[a-z0-9][a-z0-9_-]{0,63}", name) or not isinstance(value, dict):
+            raise ValueError("Invalid installed-package registry")
+        _location(value["directory"], store)
+    return entries
+
+
+def _installer(target: Path, source: str) -> list[str]:
+    """The command that installs `source` into `target`.
+
+    Environments made by uv (including `uv tool install`) carry no pip, so uv
+    installs there for this interpreter; pip is used wherever it exists.
+    """
+    import importlib.util
+
+    if importlib.util.find_spec("pip") is not None:
+        return [sys.executable, "-m", "pip", "install", "--target", str(target), source]
+    uv = shutil.which("uv")
+    if uv is None:
+        raise ValueError("Installing a package needs pip or uv in this environment")
+    return [uv, "pip", "install", "--python", sys.executable, "--target", str(target), source]
+
+
+def _location(directory: str, store: Path | None = None) -> Path:
+    store = store or root()
+    if not isinstance(directory, str) or not re.fullmatch(r"[a-f0-9]{32}", directory):
+        raise ValueError("Invalid package directory")
+    path = (store / directory).resolve()
+    if not path.is_relative_to(store.resolve()):
+        raise ValueError("Package directory is outside the package store")
+    return path
+
+
+def _save(records: dict, store: Path | None = None) -> None:
+    target = (store or root()) / "installed.json"
+    temporary = target.with_suffix(f".{uuid4().hex}.tmp")
+    try:
+        temporary.write_text(json.dumps(records, indent=2), encoding="utf-8")
+        temporary.replace(target)
+    finally:
+        temporary.unlink(missing_ok=True)
+
+
+def activate() -> None:
+    from rlmagent.discovery import default_paths
+
+    paths = default_paths()
+    stores = [root(paths=paths)]
+    if paths.project_enabled:
+        stores.insert(0, root(local=True, paths=paths))
+    for store in stores:
+        for record in _records(store).values():
+            directory = str(_location(record["directory"], store))
+            if directory not in sys.path:
+                sys.path.append(directory)
+
+
+def resource_dirs(
+    kind: str, *, records: dict | None = None, paths=None, store: Path | None = None
+) -> list[Path]:
+    if paths is not None:
+        stores = [root(paths=paths)]
+        if paths.project_enabled:
+            stores.insert(0, root(local=True, paths=paths))
+        return [
+            directory for location in stores for directory in resource_dirs(kind, store=location)
+        ]
+    store = store or root()
+    directories = []
+    for record in (_records(store) if records is None else records).values():
+        package = _location(record["directory"], store)
+        manifest = package / "rlmagent-resources.json"
+        if not manifest.is_file():
+            continue
+        resources = json.loads(manifest.read_text(encoding="utf-8"))
+        if not isinstance(resources, dict) or set(resources) - {"skills", "prompts", "themes"}:
+            raise ValueError("Invalid package resource manifest")
+        paths = resources.get(kind, [])
+        if not isinstance(paths, list) or not all(isinstance(path, str) for path in paths):
+            raise ValueError("Package resources must be lists of relative paths")
+        for path in paths:
+            resolved = (package / path).resolve()
+            if not resolved.is_relative_to(package) or not resolved.is_dir():
+                raise ValueError("Package resource directory is absent or outside its installation")
+            directories.append(resolved)
+    return directories
+
+
+def run(argv: list[str]) -> int:
+    parser = argparse.ArgumentParser(prog="rlm-agent package")
+    sub = parser.add_subparsers(dest="action", required=True)
+    listing = sub.add_parser("list")
+    listing.add_argument("--local", action="store_true", help="Use this project's package store")
+    cleanup = sub.add_parser("prune", help="Review or remove expired inactive installations")
+    cleanup.add_argument("--local", action="store_true")
+    cleanup.add_argument("--older-than", type=float, default=RETAIN_DAYS, metavar="DAYS")
+    cleanup.add_argument("--apply", action="store_true")
+    install = sub.add_parser("install")
+    install.add_argument("name")
+    install.add_argument("source", help="Trusted wheel, directory or package-index requirement")
+    install.add_argument("--local", action="store_true", help="Install for this project only")
+    for action in ("remove", "update"):
+        operation = sub.add_parser(action)
+        operation.add_argument("name")
+        if action == "update":
+            operation.add_argument(
+                "--source", default=None, help="Replacement local artifact or index requirement"
+            )
+        operation.add_argument(
+            "--local", action="store_true", help="Use this project's package store"
+        )
+    args = parser.parse_args(argv)
+    store = root(local=args.local)
+    try:
+        if args.action == "list":
+            print(json.dumps(_records(store), indent=2))
+            return 0
+        if args.action == "prune":
+            import math
+
+            if not math.isfinite(args.older_than) or args.older_than < 0:
+                raise ValueError("Cache retention must be a finite nonnegative number of days")
+            store.mkdir(parents=True, exist_ok=True)
+            with _lock(store / "registry.lock"):
+                selected = prune(
+                    store, _records(store), older_than=args.older_than, apply=args.apply
+                )
+            print(json.dumps({"removed" if args.apply else "candidates": selected}, indent=2))
+            return 0
+        if not re.fullmatch(r"[a-z0-9][a-z0-9_-]{0,63}", args.name):
+            parser.error("invalid package name")
+        store.mkdir(parents=True, exist_ok=True)
+        with _lock(store / "registry.lock"):
+            records = _records(store)
+            if args.action == "remove":
+                if args.name not in records:
+                    raise ValueError("Package is not installed")
+                del records[args.name]
+                _save(records, store)
+            else:
+                source = (
+                    args.source
+                    if args.action == "install"
+                    else args.source or records[args.name]["source"]
+                )
+                if (
+                    not source
+                    or source.startswith("-")
+                    or "://" in source
+                    or any(ch in source for ch in "\r\n")
+                ):
+                    raise ValueError("Use a package-index requirement or a local package path")
+                directory = uuid4().hex
+                target = _location(directory, store)
+                result = subprocess.run(_installer(target, source), check=False)
+                if result.returncode:
+                    return result.returncode
+                # Validate declared resources before making the installation active.
+                packages = [
+                    {"name": distribution.metadata["Name"], "version": distribution.version}
+                    for distribution in importlib.metadata.distributions(path=[str(target)])
+                ]
+                provenance = None
+                artifact = Path(source)
+                if artifact.is_file():
+                    with artifact.open("rb") as stream:
+                        provenance = hashlib.file_digest(stream, "sha256").hexdigest()
+                records[args.name] = {
+                    "source": source,
+                    "directory": directory,
+                    "distributions": packages,
+                    "artifact_sha256": provenance,
+                }
+                for kind in ("skills", "prompts", "themes"):
+                    resource_dirs(kind, records=records, store=store)
+                _save(records, store)
+            try:
+                prune(store, records, apply=True)
+            except OSError as error:
+                print(
+                    f"Package saved; cache cleanup deferred ({type(error).__name__}).",
+                    file=sys.stderr,
+                )
+        print("Package registry updated. Restart the agent to reload Python modules.")
+        return 0
+    except (OSError, ValueError, KeyError) as exc:
+        print(f"Package operation failed: {exc}", file=sys.stderr)
+        return 1