polymath-agent 0.4.0__py3-none-any.whl

polymath/config.py

@@ -0,0 +1,293 @@
+from __future__ import annotations
+
+import json
+import os
+import tempfile
+from dataclasses import dataclass, field
+from enum import Enum
+from pathlib import Path
+from typing import Callable
+
+CONFIG_DIR   = Path.home() / ".polymath"
+CONFIG_FILE  = CONFIG_DIR / "config.json"
+DB_FILE      = CONFIG_DIR / "context.db"
+SESSIONS_DIR = CONFIG_DIR / "sessions"
+PROJECTS_DIR = CONFIG_DIR / "projects"
+
+
+class TaskType(Enum):
+    CODE = "code"
+    ANALYSIS = "analysis"
+    CREATIVE = "creative"
+    MATH = "math"
+    RESEARCH = "research"
+    GENERAL = "general"
+
+
+class TaskComplexity(Enum):
+    SIMPLE = "simple"      # direct answer, skip plan/clarify
+    MEDIUM = "medium"      # one clarify round
+    COMPLEX = "complex"    # full pipeline
+
+
+class CostTier(Enum):
+    FREE = 0
+    CHEAP = 1
+    PREMIUM = 2
+    EXPENSIVE = 3
+
+
+@dataclass
+class ModelInfo:
+    id: str
+    provider: str
+    display_name: str
+    cost_tier: CostTier
+    context_window: int
+    speed: int            # 1–5
+    quality: int          # 1–5
+    capabilities: list[str]
+    input_cost_per_1k: float
+    output_cost_per_1k: float
+    available: bool = False
+    #: Whether the model accepts a sampling temperature. Claude removed
+    #: temperature/top_p/top_k with the 4.7 generation and rejects them with
+    #: a 400, so this is per-model and not per-provider.
+    supports_temperature: bool = True
+
+    def supports(self, task_type: TaskType) -> bool:
+        return task_type.value in self.capabilities or "general" in self.capabilities
+
+
+# Full model registry — availability set at runtime by detector
+MODEL_REGISTRY: list[ModelInfo] = [
+    # ── Claude ────────────────────────────────────────────────────────────────
+    # IDs are complete as written — no date suffix. These models reject
+    # temperature, hence supports_temperature=False; see ClaudeAdapter.
+    ModelInfo("claude-opus-5", "claude", "Claude Opus 5",
+              CostTier.EXPENSIVE, 1_000_000, 3, 5,
+              ["code", "analysis", "creative", "math", "research", "general"],
+              0.005, 0.025, supports_temperature=False),
+    ModelInfo("claude-sonnet-5", "claude", "Claude Sonnet 5",
+              CostTier.PREMIUM, 1_000_000, 4, 4,
+              ["code", "analysis", "creative", "math", "research", "general"],
+              0.002, 0.010, supports_temperature=False),
+    ModelInfo("claude-haiku-4-5", "claude", "Claude Haiku 4.5",
+              CostTier.CHEAP, 200_000, 5, 3,
+              ["code", "analysis", "general"],
+              0.001, 0.005),
+
+    # ── Gemini ────────────────────────────────────────────────────────────────
+    ModelInfo("gemini-2.0-flash", "gemini", "Gemini 2.0 Flash",
+              CostTier.FREE, 1_000_000, 5, 3,
+              ["code", "analysis", "general"],
+              0.0, 0.0),
+    ModelInfo("gemini-1.5-pro", "gemini", "Gemini 1.5 Pro",
+              CostTier.PREMIUM, 2_000_000, 3, 4,
+              ["code", "analysis", "creative", "math", "research", "general"],
+              0.00125, 0.005),
+    ModelInfo("gemini-1.5-flash", "gemini", "Gemini 1.5 Flash",
+              CostTier.CHEAP, 1_000_000, 5, 3,
+              ["code", "analysis", "general"],
+              0.000075, 0.0003),
+
+    # ── OpenAI ────────────────────────────────────────────────────────────────
+    ModelInfo("gpt-4o", "openai", "GPT-4o",
+              CostTier.PREMIUM, 128_000, 4, 5,
+              ["code", "analysis", "creative", "math", "vision", "general"],
+              0.0025, 0.01),
+    ModelInfo("gpt-4o-mini", "openai", "GPT-4o Mini",
+              CostTier.CHEAP, 128_000, 5, 3,
+              ["code", "analysis", "general"],
+              0.00015, 0.0006),
+    ModelInfo("o1-mini", "openai", "o1 Mini",
+              CostTier.CHEAP, 128_000, 3, 4,
+              ["code", "math", "analysis"],
+              0.003, 0.012),
+
+    # ── Ollama (local) ────────────────────────────────────────────────────────
+    ModelInfo("llama3.2:3b", "ollama", "Llama 3.2 3B",
+              CostTier.FREE, 128_000, 4, 2,
+              ["general", "code"],
+              0.0, 0.0),
+    ModelInfo("mistral:latest", "ollama", "Mistral 7B",
+              CostTier.FREE, 32_000, 4, 3,
+              ["code", "analysis", "general"],
+              0.0, 0.0),
+]
+
+def get_model(model_id: str) -> ModelInfo | None:
+    """Registry row for a model id, or None if we have never heard of it."""
+    for model in MODEL_REGISTRY:
+        if model.id == model_id:
+            return model
+    return None
+
+
+# ── Priority sort keys ────────────────────────────────────────────────────────
+
+PRIORITY_PROFILES: dict[str, Callable[[ModelInfo], tuple]] = {
+    "quality-first":  lambda m: (-m.quality, -m.context_window, -m.speed, m.cost_tier.value),
+    "cost-first":     lambda m: (m.cost_tier.value, -m.quality, -m.speed),
+    "speed-first":    lambda m: (-m.speed, m.cost_tier.value, -m.quality),
+    "balanced":       lambda m: (-(m.quality + m.speed) / 2, m.cost_tier.value),
+}
+
+
+# ── Persistent user config ────────────────────────────────────────────────────
+
+DEFAULT_CONFIG = {
+    "priority_profile": "quality-first",
+    "review_rounds": 2,
+    "context_window_messages": 8,
+    "api_keys": {},          # provider → key
+    "ollama_url": "http://localhost:11434",
+    "lmstudio_url": "http://localhost:1234",
+    # Plan-first is Polymath's headline UX: every non-`--ask` session pauses
+    # after the plan step for explicit approval. Boris-pattern: a good plan
+    # is the difference between a great session and a wasted one.
+    "plan_mode": "always",   # always | complexity | off
+    "memory": {
+        "embedding_provider": "gemini",      # gemini | openai
+        "embedding_model": "text-embedding-004",
+        "retrieval_k": 12,
+        "dedup_threshold": 0.92,
+        # Writeback OFF by default. Use /learn-from for deliberate corrections
+        # instead — see slash_commands._run_learn_from.
+        "auto_writeback": False,
+        "writeback_min_chars": 200,
+        "stale_candidate_days": 30,
+        # Head start the primary model gets before the failover model is
+        # also fired. See orchestrator/race.py.
+        "race_delay_seconds": 5.0,
+        # Run the reviewer concurrently against partial output. Off until
+        # it has proved itself. See orchestrator/speculative_review.py.
+        "speculative_review": False,
+    },
+}
+
+
+def merge_defaults(defaults: dict, saved: dict) -> dict:
+    """Overlay saved settings on the defaults, nested blocks included.
+
+    A one-level merge replaced a whole nested block whenever the saved file
+    had one, so saving a single "memory" setting silently dropped every
+    other memory default. Values the user has set win; everything else falls
+    through to the default.
+    """
+    merged = dict(defaults)
+    for key, value in saved.items():
+        current = merged.get(key)
+        if isinstance(current, dict) and isinstance(value, dict):
+            merged[key] = merge_defaults(current, value)
+        else:
+            merged[key] = value
+    return merged
+
+
+#: The config file holds every provider's API key, so it is ours alone to
+#: read. Both modes are enforced on each load, which also repairs a file
+#: written by an older version under the default umask (0644).
+CONFIG_DIR_MODE = 0o700
+CONFIG_FILE_MODE = 0o600
+
+
+def atomic_write_text(path: Path, text: str, mode: int | None = None) -> None:
+    """Replace a file's contents in one step.
+
+    A plain write truncates first, so an interruption leaves a partial
+    file. That matters most for files we do not own: polymath rewrites the
+    Gemini CLI's oauth_creds.json when refreshing a token, and a truncated
+    write there destroys the user's login to a different tool.
+
+    An existing file keeps its own permissions unless `mode` says otherwise.
+    """
+    path.parent.mkdir(parents=True, exist_ok=True)
+    if mode is None and path.exists():
+        try:
+            mode = path.stat().st_mode & 0o777
+        except OSError:
+            mode = None
+
+    fd, tmp_name = tempfile.mkstemp(dir=str(path.parent), prefix=path.name + ".", suffix=".tmp")
+    tmp = Path(tmp_name)
+    try:
+        with os.fdopen(fd, "w") as f:
+            f.write(text)
+        if mode is not None:
+            _secure(tmp, mode)
+        os.replace(tmp, path)
+    except Exception:
+        tmp.unlink(missing_ok=True)
+        raise
+
+
+def _secure(path: Path, mode: int) -> None:
+    """Best-effort chmod. A filesystem that cannot represent the mode (a
+    mounted share, say) must not stop polymath from starting."""
+    try:
+        if path.exists() and (path.stat().st_mode & 0o777) != mode:
+            path.chmod(mode)
+    except OSError:
+        pass
+
+
+def load_config() -> dict:
+    CONFIG_DIR.mkdir(parents=True, exist_ok=True)
+    _secure(CONFIG_DIR, CONFIG_DIR_MODE)
+    _secure(CONFIG_FILE, CONFIG_FILE_MODE)
+    if CONFIG_FILE.exists():
+        with open(CONFIG_FILE) as f:
+            saved = json.load(f)
+        return merge_defaults(DEFAULT_CONFIG, saved)
+    return dict(DEFAULT_CONFIG)
+
+
+def save_config(cfg: dict) -> None:
+    """Write the config atomically and readable only by its owner.
+
+    Atomic because a crash partway through a plain write leaves a truncated
+    file, and the next start would then fail to parse its own config and
+    silently fall back to defaults, losing the user's API keys.
+    """
+    CONFIG_DIR.mkdir(parents=True, exist_ok=True)
+    _secure(CONFIG_DIR, CONFIG_DIR_MODE)
+
+    # Session state borrows this dictionary — the ensemble cost warning
+    # records that it has fired as cfg["_ensemble_warned"]. The underscore
+    # marks it as not-for-disk, and this is what enforces that: persisting
+    # it would mean the warning never appears again in any future session.
+    cfg = {key: value for key, value in cfg.items() if not key.startswith("_")}
+
+    atomic_write_text(
+        CONFIG_FILE, json.dumps(cfg, indent=2), mode=CONFIG_FILE_MODE,
+    )
+
+
+def memory_setting(cfg: dict, key: str):
+    """Read one cfg["memory"] setting.
+
+    Falls back to DEFAULT_CONFIG rather than to a literal written at the call
+    site, so each default has exactly one home and cannot drift. Tolerates a
+    partial cfg, which tests and older config files both produce.
+    """
+    block = cfg.get("memory") if isinstance(cfg, dict) else None
+    if isinstance(block, dict) and key in block:
+        return block[key]
+    return DEFAULT_CONFIG["memory"][key]
+
+
+def get_api_key(provider: str, cfg: dict) -> str | None:
+    # 1. user config file
+    key = cfg.get("api_keys", {}).get(provider)
+    if key:
+        return key
+    # 2. environment variable
+    env_map = {
+        "claude":  "ANTHROPIC_API_KEY",
+        "gemini":  "GOOGLE_API_KEY",
+        "openai":  "OPENAI_API_KEY",
+        "mistral": "MISTRAL_API_KEY",
+        "groq":    "GROQ_API_KEY",
+    }
+    return os.environ.get(env_map.get(provider, ""))

polymath/context_manager.py

@@ -0,0 +1,76 @@
+"""
+Legacy facade — historical import location for the brain index.
+
+The actual implementation now lives in `nexus.index`. This file
+re-exports the public API so any existing callers
+(`from polymath.context_manager import X`) keep working unchanged.
+
+New code should import from `nexus` directly:
+
+    from nexus import append_context, read_context, TaskType
+
+This file will continue to exist for backwards compatibility — old brains
+on disk, old session exports, and external scripts may reference it for
+years to come. Long-term contract: never delete a facade.
+"""
+from __future__ import annotations
+
+from nexus.config import BRAIN_PROJECTS_DIR as PROJECTS_DIR
+from nexus.discovery import current_owner_slug, find_brain_root
+from nexus.index import (
+    ALWAYS_INJECT,
+    CONTEXT_DESCRIPTIONS,
+    CONTEXT_FILES,
+    TASK_CONTEXT_MAP,
+    _resolve_context_dir,
+    append_context,
+    artifacts_dir,
+    brain_grep,
+    brain_usage,
+    build_context_injection,
+    context_dir,
+    create_project,
+    create_session_dirs,
+    export_session_to_project,
+    get_project_meta,
+    init_brain_root,
+    list_context_files,
+    list_projects,
+    project_dir,
+    read_context,
+    save_artifact,
+    session_dir,
+    sessions_dir,
+    write_context,
+)
+from nexus.types import TaskType
+
+__all__ = [
+    "ALWAYS_INJECT",
+    "CONTEXT_DESCRIPTIONS",
+    "CONTEXT_FILES",
+    "PROJECTS_DIR",
+    "TASK_CONTEXT_MAP",
+    "TaskType",
+    "append_context",
+    "artifacts_dir",
+    "brain_grep",
+    "brain_usage",
+    "build_context_injection",
+    "context_dir",
+    "create_project",
+    "create_session_dirs",
+    "current_owner_slug",
+    "export_session_to_project",
+    "find_brain_root",
+    "get_project_meta",
+    "init_brain_root",
+    "list_context_files",
+    "list_projects",
+    "project_dir",
+    "read_context",
+    "save_artifact",
+    "session_dir",
+    "sessions_dir",
+    "write_context",
+]