polymath-agent 0.4.0__py3-none-any.whl

polymath/detector.py

@@ -0,0 +1,442 @@
+"""
+Auto-detect all available models on the device.
+Checks: CLI tools, Ollama API, LM Studio API, env vars, polymath config.
+"""
+from __future__ import annotations
+
+import shutil
+import subprocess
+import time
+from typing import TYPE_CHECKING
+
+import httpx
+
+from polymath.config import MODEL_REGISTRY, ModelInfo, get_api_key, load_config, atomic_write_text
+
+if TYPE_CHECKING:
+    pass
+
+
+_registry_cache: list[ModelInfo] | None = None
+_cache_ts: float = 0
+CACHE_TTL = 300  # seconds
+
+
+def invalidate_cache() -> None:
+    """Force next detect() call to re-probe all providers."""
+    global _registry_cache, _cache_ts
+    _registry_cache = None
+    _cache_ts = 0
+
+
+def detect(force: bool = False) -> list[ModelInfo]:
+    """Return MODEL_REGISTRY with `available` flags set. Results cached for 5 minutes."""
+    global _registry_cache, _cache_ts
+    if not force and _registry_cache and (time.time() - _cache_ts) < CACHE_TTL:
+        return _registry_cache
+    _registry_cache = _run_detect()
+    _cache_ts = time.time()
+    return _registry_cache
+
+
+def _run_detect() -> list[ModelInfo]:
+    """Internal: actually probe all providers and return registry."""
+    cfg = load_config()
+    registry = [_clone(m) for m in MODEL_REGISTRY]
+
+    _detect_claude(registry, cfg)
+    _detect_gemini(registry, cfg)
+    _detect_openai(registry, cfg)
+    _detect_ollama(registry, cfg)
+    _detect_lmstudio(registry, cfg)
+
+    return registry
+
+
+def available_models(registry: list[ModelInfo] | None = None) -> list[ModelInfo]:
+    r = registry or detect()
+    return [m for m in r if m.available]
+
+
+# ── Provider detectors ────────────────────────────────────────────────────────
+
+def _detect_claude(registry: list[ModelInfo], cfg: dict) -> None:
+    key = get_api_key("claude", cfg)
+    cli = shutil.which("claude")
+    if not key and not cli:
+        return
+    # If CLI exists but no explicit key, attempt to read from claude's own config
+    if not key and cli:
+        key = _read_claude_cli_key()
+    if key:
+        cfg.setdefault("api_keys", {})["claude"] = key
+        for m in registry:
+            if m.provider == "claude":
+                m.available = True
+
+
+def _detect_gemini(registry: list[ModelInfo], cfg: dict) -> None:
+    key = get_api_key("gemini", cfg)
+    cli = shutil.which("gemini")
+    if not key and not cli:
+        return
+    if not key and cli:
+        key = _read_gemini_cli_key()
+    if key:
+        cfg.setdefault("api_keys", {})["gemini"] = key
+        # Mark whether this is an OAuth token (from gemini CLI) vs API key
+        # OAuth tokens start with "ya29." — flag it so the adapter knows
+        cfg.setdefault("gemini_auth_type", "oauth" if key.startswith("ya29.") else "api_key")
+        for m in registry:
+            if m.provider == "gemini":
+                m.available = True
+
+
+def _detect_openai(registry: list[ModelInfo], cfg: dict) -> None:
+    key = get_api_key("openai", cfg)
+    cli = shutil.which("codex")
+    if not key and not cli:
+        return
+    if not key and cli:
+        key = _read_codex_cli_key()
+    if key:
+        cfg.setdefault("api_keys", {})["openai"] = key
+        for m in registry:
+            if m.provider == "openai":
+                m.available = True
+
+
+def _detect_ollama(registry: list[ModelInfo], cfg: dict) -> None:
+    base = cfg.get("ollama_url", "http://localhost:11434")
+    try:
+        resp = httpx.get(f"{base}/api/tags", timeout=3)
+        if resp.status_code != 200:
+            return
+        remote_models = {m["name"] for m in resp.json().get("models", [])}
+    except Exception:
+        # Ollama not running — still mark as available if CLI exists
+        if not shutil.which("ollama"):
+            return
+        remote_models = set()
+
+    # Mark known registry entries
+    for m in registry:
+        if m.provider == "ollama" and m.id in remote_models:
+            m.available = True
+
+    # Add any extra Ollama models not in registry
+    known = {m.id for m in registry if m.provider == "ollama"}
+    for name in remote_models - known:
+        registry.append(ModelInfo(
+            id=name, provider="ollama", display_name=name,
+            cost_tier=__import__("polymath.config", fromlist=["CostTier"]).CostTier.FREE,
+            context_window=32_000, speed=3, quality=2,
+            capabilities=["general", "code"],
+            input_cost_per_1k=0.0, output_cost_per_1k=0.0,
+            available=True,
+        ))
+
+
+def _detect_lmstudio(registry: list[ModelInfo], cfg: dict) -> None:
+    base = cfg.get("lmstudio_url", "http://localhost:1234")
+    try:
+        resp = httpx.get(f"{base}/v1/models", timeout=3)
+        if resp.status_code != 200:
+            return
+        for item in resp.json().get("data", []):
+            model_id = item["id"]
+            known = {m.id for m in registry}
+            if model_id not in known:
+                from polymath.config import CostTier
+                registry.append(ModelInfo(
+                    id=model_id, provider="lmstudio", display_name=model_id,
+                    cost_tier=CostTier.FREE,
+                    context_window=32_000, speed=3, quality=3,
+                    capabilities=["general", "code"],
+                    input_cost_per_1k=0.0, output_cost_per_1k=0.0,
+                    available=True,
+                ))
+            else:
+                for m in registry:
+                    if m.id == model_id:
+                        m.available = True
+    except Exception:
+        pass
+
+
+# ── Key readers ───────────────────────────────────────────────────────────────
+
+def _read_claude_cli_key() -> str | None:
+    """
+    Claude Code stores its key in the macOS Keychain, not a plain file.
+    Read it via the `security` CLI, which is the only reliable cross-version method.
+    Falls back to scanning known JSON locations for older/alternate installs.
+    """
+    import json, subprocess
+    from pathlib import Path
+
+    # 1. macOS Keychain — Claude Code stores credentials here
+    # "Claude Code-credentials" holds a JSON blob; extract apiKey if present
+    for service in ("Claude API Key", "Claude Code-credentials", "Claude Safe Storage"):
+        try:
+            result = subprocess.run(
+                ["security", "find-generic-password", "-s", service, "-w"],
+                capture_output=True, text=True, timeout=5,
+            )
+            if result.returncode == 0 and result.stdout.strip():
+                raw = result.stdout.strip()
+                # Try to parse as JSON (Claude Code-credentials format)
+                try:
+                    import json as _json
+                    data = _json.loads(raw)
+                    # "claudeAiOauth" holds a nested dict with accessToken
+                    oauth = data.get("claudeAiOauth")
+                    if isinstance(oauth, dict) and oauth.get("accessToken"):
+                        return oauth["accessToken"]
+                    if isinstance(oauth, str) and "accessToken" in oauth:
+                        # Sometimes stored as a JSON string inside JSON
+                        inner = _json.loads(oauth)
+                        if inner.get("accessToken"):
+                            return inner["accessToken"]
+                    for field in ("apiKey", "api_key", "ANTHROPIC_API_KEY", "claudeApiKey"):
+                        if data.get(field):
+                            return data[field]
+                except Exception:
+                    # Raw string — check it looks like an Anthropic key/token
+                    if raw.startswith(("sk-ant-", "claude-")):
+                        return raw
+        except Exception:
+            pass
+
+    # 2. JSON file fallbacks — Linux / Windows / older Claude CLI versions
+    candidates = [
+        Path.home() / ".claude" / "config.json",
+        Path.home() / ".config" / "claude" / "config.json",
+    ]
+    # Windows: %APPDATA%\Claude\config.json
+    import os as _os
+    _appdata = _os.environ.get("APPDATA")
+    if _appdata:
+        candidates.append(Path(_appdata) / "Claude" / "config.json")
+    # Also check ANTHROPIC_API_KEY env var here as a final fallback
+    _env_key = _os.environ.get("ANTHROPIC_API_KEY")
+    if _env_key:
+        return _env_key
+    for p in candidates:
+        if p.exists():
+            try:
+                data = json.loads(p.read_text())
+                for field in ("apiKey", "api_key", "ANTHROPIC_API_KEY"):
+                    if data.get(field):
+                        return data[field]
+            except Exception:
+                pass
+    return None
+
+
+def _read_gemini_cli_key() -> str | None:
+    """
+    Gemini CLI uses OAuth, not a raw API key.
+    Reads ~/.gemini/oauth_creds.json. If the token is expired (expiry_date < now),
+    attempts a silent refresh via the refresh_token before returning.
+    Also checks standard JSON config locations for API-key-based setups.
+    """
+    import json
+    import time as _time
+    from pathlib import Path
+
+    # 1. OAuth access token (google-gemini CLI)
+    oauth_path = Path.home() / ".gemini" / "oauth_creds.json"
+    if oauth_path.exists():
+        try:
+            data = json.loads(oauth_path.read_text())
+            token = data.get("access_token")
+            expiry_ms = data.get("expiry_date", 0)
+            refresh_token = data.get("refresh_token")
+
+            # expiry_date is in milliseconds
+            is_expired = expiry_ms and (_time.time() * 1000) > expiry_ms
+
+            if is_expired and refresh_token:
+                token = _refresh_gemini_token(oauth_path, data, refresh_token)
+            elif token:
+                return token
+        except Exception:
+            pass
+
+    # 2. JSON config API key (older / API-key based setups)
+    candidates = [
+        Path.home() / ".gemini" / "config.json",
+        Path.home() / ".gemini" / "settings.json",
+        Path.home() / ".config" / "gemini" / "config.json",
+        Path.home() / ".config" / "google" / "gemini.json",
+    ]
+    for p in candidates:
+        if p.exists():
+            try:
+                data = json.loads(p.read_text())
+                for field_name in ("apiKey", "api_key", "GOOGLE_API_KEY", "key"):
+                    if data.get(field_name):
+                        return data[field_name]
+            except Exception:
+                pass
+    return None
+
+
+def _refresh_gemini_token(oauth_path, data: dict, refresh_token: str) -> str | None:
+    """
+    Attempt to refresh an expired Gemini OAuth access token using the refresh_token.
+    Updates oauth_creds.json in-place. Returns new access_token or None on failure.
+    """
+    import json
+    import time as _time
+    try:
+        # Google OAuth2 token endpoint — client_id/secret are public Gemini CLI values
+        resp = httpx.post(
+            "https://oauth2.googleapis.com/token",
+            data={
+                "grant_type": "refresh_token",
+                "refresh_token": refresh_token,
+                # Gemini CLI's public OAuth client credentials
+                "client_id": "681254290090-ob3mcu4n8j2rdjvo8c32pnfenrh3hjqm.apps.googleusercontent.com",
+                "client_secret": "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",
+            },
+            timeout=10,
+        )
+        if resp.status_code == 200:
+            new_data = resp.json()
+            new_token = new_data.get("access_token")
+            if new_token:
+                # Merge new fields back and persist
+                data.update(new_data)
+                import time as _time2
+                if "expires_in" in new_data and "expiry_date" not in new_data:
+                    data["expiry_date"] = int((_time2.time() + new_data["expires_in"]) * 1000)
+                # Another tool's credential file. A truncated write here
+                # logs the user out of the Gemini CLI, not just polymath.
+                atomic_write_text(oauth_path, json.dumps(data, indent=2))
+                return new_token
+    except Exception:
+        pass
+    return None
+
+
+# ── Gemini account management ─────────────────────────────────────────────────
+
+def list_gemini_accounts() -> dict:
+    """
+    Return {"active": email_str, "old": [email, ...]} from google_accounts.json.
+    Returns {"active": None, "old": []} if not found.
+    """
+    import json
+    from pathlib import Path
+    accounts_path = Path.home() / ".gemini" / "google_accounts.json"
+    if not accounts_path.exists():
+        return {"active": None, "old": []}
+    try:
+        data = json.loads(accounts_path.read_text())
+        active = data.get("active")
+        old = data.get("old", [])
+        if isinstance(old, str):
+            old = [old] if old else []
+        return {"active": active, "old": list(old)}
+    except Exception:
+        return {"active": None, "old": []}
+
+
+def switch_gemini_account(email: str) -> bool:
+    """
+    Switch active Gemini account by updating google_accounts.json and clearing
+    oauth_creds.json (forcing re-login on next use). Returns True on success.
+    Note: actually switching accounts requires `gemini auth login` for the new account.
+    This just updates the active pointer and wipes cached tokens.
+    """
+    import json
+    from pathlib import Path
+    accounts_path = Path.home() / ".gemini" / "google_accounts.json"
+    if not accounts_path.exists():
+        return False
+    try:
+        data = json.loads(accounts_path.read_text())
+        current_active = data.get("active")
+        old = data.get("old", [])
+        if isinstance(old, str):
+            old = [old] if old else []
+        old = list(old)
+
+        if email == current_active:
+            return True  # already active
+
+        all_known = ([current_active] if current_active else []) + old
+        if email not in all_known:
+            return False  # unknown account
+
+        # Rotate: current active → old, requested → active
+        new_old = [a for a in all_known if a != email]
+        data["active"] = email
+        data["old"] = new_old
+        atomic_write_text(accounts_path, json.dumps(data, indent=2))
+
+        # Wipe token cache so next detect() re-reads/re-logs
+        oauth_path = Path.home() / ".gemini" / "oauth_creds.json"
+        if oauth_path.exists():
+            oauth_path.unlink()
+        invalidate_cache()
+        return True
+    except Exception:
+        return False
+
+
+def _read_codex_cli_key() -> str | None:
+    """
+    Codex CLI stores the OpenAI key in ~/.codex/auth.json under OPENAI_API_KEY.
+    """
+    import json
+    from pathlib import Path
+
+    # 1. Codex auth.json (primary location)
+    auth = Path.home() / ".codex" / "auth.json"
+    if auth.exists():
+        try:
+            data = json.loads(auth.read_text())
+            # Direct API key (api_key auth mode)
+            key = data.get("OPENAI_API_KEY")
+            if key:
+                return key
+            # ChatGPT OAuth mode — access_token works as bearer for OpenAI API
+            if data.get("auth_mode") == "chatgpt":
+                token = (data.get("tokens") or {}).get("access_token")
+                if token:
+                    return f"chatgpt_oauth:{token}"  # flagged for adapter
+        except Exception:
+            pass
+
+    # 2. Fallback JSON locations
+    candidates = [
+        Path.home() / ".codex" / "config.json",
+        Path.home() / ".config" / "codex" / "config.json",
+        Path.home() / ".openai" / "config.json",
+    ]
+    for p in candidates:
+        if p.exists():
+            try:
+                data = json.loads(p.read_text())
+                for field in ("apiKey", "api_key", "OPENAI_API_KEY", "key"):
+                    if data.get(field):
+                        return data[field]
+            except Exception:
+                pass
+    return None
+
+
+# ── Helpers ───────────────────────────────────────────────────────────────────
+
+def _clone(m: ModelInfo) -> ModelInfo:
+    return ModelInfo(
+        id=m.id, provider=m.provider, display_name=m.display_name,
+        cost_tier=m.cost_tier, context_window=m.context_window,
+        speed=m.speed, quality=m.quality, capabilities=list(m.capabilities),
+        input_cost_per_1k=m.input_cost_per_1k,
+        output_cost_per_1k=m.output_cost_per_1k,
+        available=False,
+    )

polymath/domain.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Callable
+
+from polymath.config import ModelInfo, TaskType
+from polymath.project_config import LocalProject
+
+
+@dataclass
+class Project:
+    name: str = ""
+    source: str = "local"  # local | repo | none
+    shared_root: Path | None = None
+    local_project: LocalProject | None = None
+    repo_root: Path | None = None
+    repo_branch: str = ""
+    repo_dirty: bool = False
+
+    @property
+    def shared(self) -> bool:
+        return self.source == "repo" and self.shared_root is not None
+
+
+@dataclass
+class SessionState:
+    session_id: str
+    title: str = ""
+    active_project: str = ""
+    queued_inputs: list[str] = field(default_factory=list)
+    last_output: str = ""
+
+
+@dataclass
+class ModelSelection:
+    primary: ModelInfo | None = None
+    reviewer: ModelInfo | None = None
+    summarizer: ModelInfo | None = None
+    estimated_prompt_tokens: int = 0
+    fallback_reason: str = ""
+
+
+@dataclass
+class PermissionRequest:
+    """One request to run something that needs the user's consent.
+
+    `tool_name` and `arguments` are the tool call itself, and they are what
+    the permission layer decides on. `action` is a rendering of the same
+    call for the user to read — never parse it back, that is the bug this
+    dataclass exists to prevent.
+    """
+    action: str
+    tool_name: str = ""
+    shell_command: str = ""
+    arguments: dict = field(default_factory=dict)
+    #: Non-empty when this particular call reaches beyond the tool's ordinary
+    #: scope, e.g. a read outside the project. An allowlist entry for the tool
+    #: does not cover it — the user is asked every time.
+    escalation_reason: str = ""
+    requires_prompt: bool = True
+
+
+@dataclass
+class CommandInvocation:
+    raw: str
+    name: str
+    args: str = ""
+    kind: str = "user"  # user | builtin | custom
+
+
+@dataclass
+class RunContext:
+    session: SessionState
+    project: Project
+    task_type: TaskType
+    model_selection: ModelSelection
+    permission_handler: Callable[[PermissionRequest], bool]