rlmness-cli 0.2.0__py3-none-any.whl

rlmagent/config/store.py

@@ -0,0 +1,275 @@
+"""Persisted user configuration: schema, file location, and TOML read/write."""
+
+from __future__ import annotations
+
+import json
+import os
+import tomllib
+from dataclasses import dataclass, field
+from pathlib import Path
+
+# ---------------------------------------------------------------------------
+# Provider metadata (plain data — the set of backends rlm-agent can talk to)
+# ---------------------------------------------------------------------------
+
+
+@dataclass(frozen=True, slots=True)
+class ProviderInfo:
+    """Static facts about a supported provider backend."""
+
+    key: str
+    label: str
+    key_env: str
+    base_url: str
+    models: tuple[str, ...]
+    needs_key: bool = True
+
+
+PROVIDERS: tuple[ProviderInfo, ...] = (
+    ProviderInfo(key="copilot", label="GitHub Copilot", key_env="COPILOT_GITHUB_TOKEN",
+                 base_url="", models=("auto",), needs_key=False),
+    ProviderInfo(key="claude-subscription", label="Claude Code subscription", key_env="",
+                 base_url="", models=("cli:claude-code",), needs_key=False),
+    ProviderInfo(key="xai", label="xAI", key_env="XAI_API_KEY",
+                 base_url="https://api.x.ai/v1", models=("grok-4",)),
+    ProviderInfo(key="xai-subscription", label="xAI registered OAuth", key_env="",
+                 base_url="https://api.x.ai/v1", models=("grok-4",), needs_key=False),
+    ProviderInfo(key="codex", label="Codex subscription", key_env="",
+                 base_url="https://api.openai.com/v1", models=("gpt-6.1-sol",), needs_key=False),
+    ProviderInfo(key="google", label="Google Gemini", key_env="GOOGLE_API_KEY",
+                 base_url="https://generativelanguage.googleapis.com/v1beta", models=("gemini-2.5-flash",)),
+    ProviderInfo(key="mistral", label="Mistral", key_env="MISTRAL_API_KEY",
+                 base_url="https://api.mistral.ai/v1", models=("mistral-small-latest",)),
+    ProviderInfo(key="azure", label="Azure OpenAI", key_env="AZURE_OPENAI_API_KEY",
+                 base_url="", models=("deployment-name",)),
+    ProviderInfo(
+        key="openai",
+        label="OpenAI",
+        key_env="OPENAI_API_KEY",
+        base_url="https://api.openai.com/v1",
+        models=("gpt-5", "gpt-5-mini", "gpt-4.1"),
+    ),
+    ProviderInfo(
+        key="anthropic",
+        label="Anthropic",
+        key_env="ANTHROPIC_API_KEY",
+        base_url="https://api.anthropic.com",
+        models=("claude-opus-5", "claude-sonnet-5", "claude-haiku-4-5"),
+    ),
+    ProviderInfo(
+        key="openrouter",
+        label="OpenRouter",
+        key_env="OPENROUTER_API_KEY",
+        base_url="https://openrouter.ai/api/v1",
+        models=(
+            "anthropic/claude-opus-5",
+            "openai/gpt-5",
+            "meta-llama/llama-3.3-70b-instruct",
+        ),
+    ),
+    ProviderInfo(
+        key="ollama",
+        label="Ollama",
+        key_env="OLLAMA_API_KEY",
+        base_url="http://localhost:11434/v1",
+        models=("llama3.3", "qwen2.5-coder", "deepseek-r1"),
+        needs_key=False,
+    ),
+    ProviderInfo(
+        key="huggingface",
+        label="Hugging Face",
+        key_env="HF_TOKEN",
+        base_url="https://router.huggingface.co/v1",
+        models=(
+            "openai/gpt-oss-120b",
+            "Qwen/Qwen3-Coder-480B-A35B-Instruct",
+            "moonshotai/Kimi-K2-Instruct",
+        ),
+    ),
+    ProviderInfo(key="vertex", label="Google Vertex", key_env="GOOGLE_CLOUD_PROJECT",
+                 base_url="", models=("google/gemini-3.5-flash",), needs_key=False),
+)
+
+
+def provider_info(key: str) -> ProviderInfo:
+    """Look up provider metadata by key, raising ``KeyError`` when unknown."""
+    for info in PROVIDERS:
+        if info.key == key:
+            return info
+    raise KeyError(key)
+
+
+PROVIDER_KEYS: tuple[str, ...] = tuple(p.key for p in PROVIDERS)
+
+
+# ---------------------------------------------------------------------------
+# Config file schema
+# ---------------------------------------------------------------------------
+
+
+@dataclass(slots=True)
+class RlmAgentConfig:
+    """Contents of the user's ``config.toml``."""
+
+    provider: str | None = None
+    model: str | None = None
+    api_keys: dict[str, str] = field(default_factory=dict)
+    base_urls: dict[str, str] = field(default_factory=dict)
+    # How folders with no saved trust decision are treated: ask, always, never.
+    # Read only from this user-wide file, never from a project.
+    project_trust: str | None = None
+
+
+# ---------------------------------------------------------------------------
+# Location
+# ---------------------------------------------------------------------------
+
+
+MAX_RECENT = 8
+
+
+def recent_models() -> list[dict]:
+    """Models chosen before, newest first: [{"provider", "model"}]. Holds no keys."""
+    try:
+        saved = json.loads((config_dir() / "recent_models.json").read_text(encoding="utf-8"))
+    except (OSError, ValueError):
+        return []
+    return [r for r in saved if isinstance(r, dict) and r.get("provider") and r.get("model")]
+
+
+def remember_model(provider: str, model: str) -> None:
+    entry = {"provider": provider, "model": model}
+    kept = [entry] + [r for r in recent_models() if r != entry]
+    config_dir().mkdir(parents=True, exist_ok=True)
+    (config_dir() / "recent_models.json").write_text(
+        json.dumps(kept[:MAX_RECENT], indent=2), encoding="utf-8")
+
+
+def forget_model(provider: str, model: str) -> None:
+    kept = [r for r in recent_models() if r != {"provider": provider, "model": model}]
+    (config_dir() / "recent_models.json").write_text(json.dumps(kept, indent=2), encoding="utf-8")
+
+
+def config_dir() -> Path:
+    """Directory holding rlm-agent's user configuration."""
+    override = os.environ.get("RLM_AGENT_CONFIG_DIR") or os.environ.get("RLM_AGENT_HOME")
+    if override:
+        return Path(override)
+    return Path.home() / ".rlm-agent"
+
+
+def config_path() -> Path:
+    """Full path to ``config.toml``."""
+    return config_dir() / "config.toml"
+
+
+# ---------------------------------------------------------------------------
+# Read
+# ---------------------------------------------------------------------------
+
+
+def load_config() -> RlmAgentConfig:
+    """Read the config file, returning empty defaults when absent or malformed."""
+    path = config_path()
+    try:
+        raw = path.read_bytes()
+    except OSError:
+        return RlmAgentConfig()
+
+    try:
+        data = tomllib.loads(raw.decode("utf-8"))
+    except (tomllib.TOMLDecodeError, UnicodeDecodeError):
+        return RlmAgentConfig()
+
+    providers = data.get("providers")
+    api_keys: dict[str, str] = {}
+    base_urls: dict[str, str] = {}
+    if isinstance(providers, dict):
+        for name, section in providers.items():
+            if not isinstance(section, dict):
+                continue
+            key = section.get("api_key")
+            if isinstance(key, str) and key:
+                api_keys[name] = key
+            url = section.get("base_url")
+            if isinstance(url, str) and url:
+                base_urls[name] = url
+
+    provider = data.get("provider")
+    model = data.get("model")
+    trust = data.get("project_trust")
+    return RlmAgentConfig(
+        provider=provider if isinstance(provider, str) and provider else None,
+        model=model if isinstance(model, str) and model else None,
+        api_keys=api_keys,
+        base_urls=base_urls,
+        project_trust=trust if trust in ("ask", "always", "never") else None,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Write
+# ---------------------------------------------------------------------------
+
+
+def _quote(value: str) -> str:
+    """Serialize a string as a TOML basic string."""
+    escaped = value.replace("\\", "\\\\").replace('"', '\\"')
+    return f'"{escaped}"'
+
+
+def render_config(config: RlmAgentConfig) -> str:
+    """Render *config* as TOML text."""
+    lines: list[str] = []
+    if config.provider:
+        lines.append(f"provider = {_quote(config.provider)}")
+    if config.model:
+        lines.append(f"model = {_quote(config.model)}")
+    if config.project_trust:
+        lines.append(f"project_trust = {_quote(config.project_trust)}")
+
+    names = sorted(set(config.api_keys) | set(config.base_urls))
+    for name in names:
+        lines.append("")
+        lines.append(f"[providers.{name}]")
+        if name in config.api_keys:
+            lines.append(f"api_key = {_quote(config.api_keys[name])}")
+        if name in config.base_urls:
+            lines.append(f"base_url = {_quote(config.base_urls[name])}")
+    return "\n".join(lines) + "\n"
+
+
+def save_config(config: RlmAgentConfig) -> Path:
+    """Write *config* to disk with owner-only permissions."""
+    path = config_path()
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(render_config(config), encoding="utf-8")
+    try:
+        path.chmod(0o600)
+    except OSError:
+        pass  # best effort — some filesystems (e.g. FAT) reject chmod
+    return path
+
+
+def mask_key(value: str) -> str:
+    """Render an API key safe for display: last four characters only."""
+    if not value:
+        return ""
+    if len(value) <= 4:
+        return "*" * len(value)
+    return "*" * (len(value) - 4) + value[-4:]
+
+
+__all__ = [
+    "RlmAgentConfig",
+    "PROVIDERS",
+    "PROVIDER_KEYS",
+    "ProviderInfo",
+    "config_dir",
+    "config_path",
+    "load_config",
+    "mask_key",
+    "provider_info",
+    "render_config",
+    "save_config",
+]

rlmagent/context/__init__.py

@@ -0,0 +1,47 @@
+"""rlm-agent context: token budgeting, threshold detection, and compaction planning."""
+
+from rlmagent.context.budget import (
+    CHARS_PER_TOKEN,
+    DEFAULT_SUMMARY_SYSTEM,
+    DEFAULT_WINDOW,
+    MODEL_WINDOWS,
+    TOOL_SCHEMA_BASELINE,
+    CompactionPlan,
+    CompactionResult,
+    ContextEstimate,
+    ContextLimits,
+    apply_compaction,
+    build_summary_prompts,
+    estimate_context,
+    estimate_entry_tokens,
+    estimate_system_tokens,
+    estimate_tool_tokens,
+    estimate_transcript_tokens,
+    exceeds_threshold,
+    plan_compaction,
+    resolve_window,
+    serialize_for_summary,
+)
+
+__all__ = [
+    "CHARS_PER_TOKEN",
+    "CompactionPlan",
+    "CompactionResult",
+    "ContextEstimate",
+    "ContextLimits",
+    "DEFAULT_SUMMARY_SYSTEM",
+    "DEFAULT_WINDOW",
+    "MODEL_WINDOWS",
+    "TOOL_SCHEMA_BASELINE",
+    "apply_compaction",
+    "build_summary_prompts",
+    "estimate_context",
+    "estimate_entry_tokens",
+    "estimate_system_tokens",
+    "estimate_tool_tokens",
+    "estimate_transcript_tokens",
+    "exceeds_threshold",
+    "plan_compaction",
+    "resolve_window",
+    "serialize_for_summary",
+]