code-muse 0.0.1__py3-none-any.whl

code_muse/model_switching.py

@@ -0,0 +1,59 @@
+"""Shared helpers for switching models and reloading agents safely."""
+
+from code_muse.config import set_model_name
+
+
+def _get_effective_agent_model(agent) -> str | None:
+    """Safely fetch the effective model name for an agent."""
+    try:
+        return agent.get_model_name()
+    except Exception:
+        return None
+
+
+def set_model_and_reload_agent(
+    model_name: str,
+    *,
+    warn_on_pinned_mismatch: bool = True,
+) -> None:
+    """Set the global model and reload the active agent.
+
+    This keeps model switching consistent across commands while avoiding
+    direct imports that can trigger circular dependencies.
+    """
+    from code_muse.messaging import emit_info, emit_warning
+
+    set_model_name(model_name)
+
+    try:
+        from code_muse.agents import get_current_agent
+
+        current_agent = get_current_agent()
+        if current_agent is None:
+            emit_warning("Model changed but no active agent was found to reload")
+            return
+
+        # JSON agents may need to refresh their config before reload
+        if hasattr(current_agent, "refresh_config"):
+            try:
+                current_agent.refresh_config()
+            except Exception:
+                # Non-fatal, continue to reload
+                ...
+
+        if warn_on_pinned_mismatch:
+            effective_model = _get_effective_agent_model(current_agent)
+            if effective_model and effective_model != model_name:
+                display_name = getattr(
+                    current_agent, "display_name", current_agent.name
+                )
+                emit_warning(
+                    "Active agent "
+                    f"'{display_name}' is pinned to '{effective_model}', "
+                    f"so '{model_name}' will not take effect until unpinned."
+                )
+
+        current_agent.reload_code_generation_agent()
+        emit_info("Active agent reloaded")
+    except Exception as exc:
+        emit_warning(f"Model changed but agent reload failed: {exc}")

code_muse/model_utils.py

@@ -0,0 +1,156 @@
+"""Model-related utilities shared across agents and tools.
+
+This module is intentionally model-agnostic. Anything model-family-specific
+(e.g. claude-code OAuth prompt handling) lives in its own plugin and hooks
+into the ``prepare_model_prompt`` or ``get_model_system_prompt`` callbacks.
+
+Plugins can register:
+
+- ``prepare_model_prompt``: fully take over prompt prep for a model family.
+- ``get_model_system_prompt``: augment/override the system prompt for a model.
+"""
+
+from dataclasses import dataclass
+
+
+@dataclass
+class PreparedPrompt:
+    """Result of preparing a prompt for a specific model.
+
+    Attributes:
+        instructions: The system instructions to use for the agent
+        user_prompt: The user prompt (possibly modified)
+        is_claude_code: Whether this is a claude-code model (set by the
+            claude_code_oauth plugin via the ``prepare_model_prompt`` hook).
+    """
+
+    instructions: str
+    user_prompt: str
+    is_claude_code: bool
+
+
+def prepare_prompt_for_model(
+    model_name: str,
+    system_prompt: str,
+    user_prompt: str,
+    prepend_system_to_user: bool = True,
+) -> PreparedPrompt:
+    """Prepare instructions and prompt for a specific model.
+
+    Core fires two hooks to let plugins customize prompt prep:
+
+    1. ``prepare_model_prompt`` — first winner with ``handled=True`` takes
+       over entirely (used by the claude_code_oauth plugin).
+    2. ``get_model_system_prompt`` — legacy per-model system-prompt hook;
+       still fired for compatibility with plugins (e.g. agent_skills) that
+       rely on it.
+
+    If no plugin handles the model, we return the original system/user prompt
+    unchanged.
+
+    Args:
+        model_name: The name of the model being used.
+        system_prompt: The default system prompt from the agent.
+        user_prompt: The user's prompt/message.
+        prepend_system_to_user: Whether to prepend the system prompt to the
+            user prompt (only meaningful for plugins that opt into it).
+
+    Returns:
+        PreparedPrompt ready for the model.
+    """
+    from code_muse import callbacks
+
+    # 1) Give the dedicated prepare_model_prompt hook first crack. First
+    #    plugin to claim ``handled=True`` wins.
+    for result in callbacks.on_prepare_model_prompt(
+        model_name, system_prompt, user_prompt, prepend_system_to_user
+    ):
+        if result and isinstance(result, dict) and result.get("handled"):
+            return PreparedPrompt(
+                instructions=result.get("instructions", system_prompt),
+                user_prompt=result.get("user_prompt", user_prompt),
+                is_claude_code=bool(result.get("is_claude_code", False)),
+            )
+
+    # 2) Fall back to the legacy per-model system-prompt hook for plugins
+    #    that still register there.
+    for result in callbacks.on_get_model_system_prompt(
+        model_name, system_prompt, user_prompt
+    ):
+        if result and isinstance(result, dict) and result.get("handled"):
+            return PreparedPrompt(
+                instructions=result.get("instructions", system_prompt),
+                user_prompt=result.get("user_prompt", user_prompt),
+                is_claude_code=bool(result.get("is_claude_code", False)),
+            )
+
+    # 3) No plugin handled it — return the caller's prompts unchanged.
+    return PreparedPrompt(
+        instructions=system_prompt,
+        user_prompt=user_prompt,
+        is_claude_code=False,
+    )
+
+
+def supports_adaptive_thinking(
+    model_name: str, actual_model_id: str | None = None
+) -> bool:
+    """Return whether a model should default to adaptive thinking.
+
+    Opus 4-6, Opus 4-7, and Sonnet 4-6 models support adaptive thinking.
+    Checks both the alias/key and the real model ID to handle Bedrock-style
+    names like ``us.anthropic.claude-opus-4-7``.
+
+    Args:
+        model_name: The model alias/key (e.g. ``"bedrock-opus-4-7"``).
+        actual_model_id: The real model ID from config (e.g.
+            ``"us.anthropic.claude-opus-4-7"``).
+    """
+    candidates = [model_name.lower()]
+    if actual_model_id:
+        candidates.append(actual_model_id.lower())
+
+    _ADAPTIVE_TAGS = (
+        "opus-4-6",
+        "4-6-opus",
+        "opus-4-7",
+        "4-7-opus",
+        "sonnet-4-6",
+        "4-6-sonnet",
+    )
+    return any(tag in c for c in candidates for tag in _ADAPTIVE_TAGS)
+
+
+def get_default_extended_thinking(
+    model_name: str, actual_model_id: str | None = None
+) -> str:
+    """Return the default extended_thinking mode for an Anthropic model.
+
+    Opus 4-6, Opus 4-7, and Sonnet 4-6 models default to ``"adaptive"``
+    thinking; all other Anthropic models default to ``"enabled"``.
+
+    Args:
+        model_name: The model alias/key (e.g. ``"bedrock-opus-4-7"``).
+        actual_model_id: The real model ID from config (e.g.
+            ``"us.anthropic.claude-opus-4-7"``).
+
+    Returns:
+        ``"adaptive"`` for supported variants, ``"enabled"`` otherwise.
+    """
+    if supports_adaptive_thinking(model_name, actual_model_id):
+        return "adaptive"
+    return "enabled"
+
+
+def should_use_anthropic_thinking_summary(
+    model_name: str, actual_model_id: str | None = None
+) -> bool:
+    """Return whether Anthropic adaptive thinking should request summary display.
+
+    Anthropic's newer Opus 4.7 models require ``display: \"summarized\"`` alongside
+    ``thinking={"type": "adaptive"}``.
+    """
+    candidates = [model_name.lower()]
+    if actual_model_id:
+        candidates.append(actual_model_id.lower())
+    return any("opus-4-7" in c or "4-7-opus" in c for c in candidates)

code_muse/models.json

@@ -0,0 +1,66 @@
+{
+  "wafer.ai-glm-5.1": {
+    "type": "custom_openai",
+    "provider": "wafer",
+    "name": "GLM-5.1",
+    "custom_endpoint": {
+      "url": "https://pass.wafer.ai/v1",
+      "api_key": "$WAFER_API_KEY"
+    },
+    "context_length": 200000,
+    "supported_settings": ["temperature", "seed", "top_p", "max_tokens"]
+  },
+  "wafer.ai-Qwen3.5-397B-A17B": {
+    "type": "custom_openai",
+    "provider": "wafer",
+    "name": "Qwen3.5-397B-A17B",
+    "custom_endpoint": {
+      "url": "https://pass.wafer.ai/v1",
+      "api_key": "$WAFER_API_KEY"
+    },
+    "context_length": 262144,
+    "supported_settings": ["temperature", "seed", "top_p", "max_tokens"]
+  },
+  "wafer-DeepSeek-V4-Pro": {
+    "type": "custom_openai",
+    "provider": "wafer",
+    "name": "DeepSeek-V4-Pro",
+    "custom_endpoint": {
+      "url": "https://pass.wafer.ai/v1",
+      "api_key": "$WAFER_API_KEY"
+    },
+    "context_length": 1000000,
+    "supported_settings": ["temperature", "seed", "top_p", "max_tokens"]
+  },
+  "wafer-MiniMax-M2.7": {
+    "type": "custom_openai",
+    "provider": "wafer",
+    "name": "MiniMax-M2.7",
+    "custom_endpoint": {
+      "url": "https://pass.wafer.ai/v1",
+      "api_key": "$WAFER_API_KEY"
+    },
+    "context_length": 204800,
+    "supported_settings": ["temperature", "seed", "top_p", "max_tokens"]
+  },
+  "crof-kimi-k2.5-lightning": {
+    "type": "custom_openai",
+    "provider": "crof",
+    "name": "kimi-k2.5-lightning",
+    "custom_endpoint": {
+      "url": "https://crof.ai/v1",
+      "api_key": "$CROF_API_KEY"
+    },
+    "context_length": 131072,
+    "streaming": false,
+    "strict_tools": false,
+    "supported_settings": ["temperature", "seed", "top_p"]
+  },
+  "zai-glm-5-coding": {
+    "type": "zai_coding",
+    "provider": "zai",
+    "name": "glm-5",
+    "context_length": 200000,
+    "supported_settings": ["temperature", "top_p"]
+  }
+}

code_muse/models_cache/__init__.py

@@ -0,0 +1,26 @@
+"""Models Cache + LRU module for Muse.
+
+Ports Codex's models_cache.json pattern and Oh-My-Pi's fs_cache LRU
+to Python with thread-safe generics.
+"""
+
+from code_muse.models_cache.blocking_lru_cache import BlockingLruCache
+from code_muse.models_cache.cache_writer import write_models_cache
+from code_muse.models_cache.sha256_hash import sha256_digest, sha256_digest_file
+from code_muse.models_cache.startup_integration import (
+    CACHE_TTL,
+    MODELS_CACHE_PATH,
+    load_cached_models,
+    refresh_models_cache,
+)
+
+__all__ = [
+    "BlockingLruCache",
+    "sha256_digest",
+    "sha256_digest_file",
+    "write_models_cache",
+    "load_cached_models",
+    "refresh_models_cache",
+    "CACHE_TTL",
+    "MODELS_CACHE_PATH",
+]

code_muse/models_cache/blocking_lru_cache.py

@@ -0,0 +1,98 @@
+"""Thread-safe generic LRU cache using OrderedDict and threading.Lock.
+
+FREE-THREADED: This cache uses threading.Lock because it may be accessed
+from both sync and async contexts. If all callers become fully async,
+consider switching to asyncio.Lock.
+"""
+
+import threading
+from collections import OrderedDict
+from collections.abc import Callable
+from typing import TypeVar
+
+K = TypeVar("K")
+V = TypeVar("V")
+
+
+class BlockingLruCache[K, V]:
+    """Thread-safe LRU cache with get_or_insert_with pattern.
+
+    Uses OrderedDict for O(1) LRU eviction and threading.Lock for
+    thread safety. Capacity must be a positive integer (> 0).
+    Gracefully works outside asyncio/tokio runtimes.
+    """
+
+    def __init__(self, capacity: int) -> None:
+        if capacity < 1:
+            raise ValueError("capacity must be positive")
+        self._capacity = capacity
+        # FREE-THREADED: Generic cache — may be accessed from sync or async code.
+        # Keep threading.Lock; migrate to asyncio.Lock only if all callers are async.
+        self._lock = threading.Lock()
+        self._cache: OrderedDict[K, V] = OrderedDict()
+
+    def get(self, key: K) -> V | None:
+        """Return cached value or None. Promotes to MRU on hit."""
+        with self._lock:
+            if key in self._cache:
+                self._cache.move_to_end(key)
+                return self._cache[key]
+            return None
+
+    def get_or_insert_with(self, key: K, factory: Callable[[], V]) -> V:
+        """Return cached value, or call factory on miss and cache it.
+
+        The factory is called OUTSIDE the lock to avoid holding it during
+        expensive computation. A double-check is performed under lock.
+        """
+        # Fast path: check cache under lock
+        with self._lock:
+            if key in self._cache:
+                self._cache.move_to_end(key)
+                return self._cache[key]
+
+        # Slow path: compute outside lock
+        value = factory()
+
+        # Insert under lock with eviction check
+        with self._lock:
+            # Double-check: another thread may have inserted while we computed
+            if key in self._cache:
+                self._cache.move_to_end(key)
+                return self._cache[key]
+
+            # Evict oldest if at capacity
+            while len(self._cache) >= self._capacity:
+                self._cache.popitem(last=False)
+
+            self._cache[key] = value
+            return value
+
+    def insert(self, key: K, value: V) -> None:
+        """Insert a value, evicting oldest if at capacity."""
+        with self._lock:
+            if key in self._cache:
+                self._cache.move_to_end(key)
+                self._cache[key] = value
+                return
+            while len(self._cache) >= self._capacity:
+                self._cache.popitem(last=False)
+            self._cache[key] = value
+
+    def remove(self, key: K) -> V | None:
+        """Remove and return the value for key, or None."""
+        with self._lock:
+            return self._cache.pop(key, None)
+
+    def clear(self) -> None:
+        """Remove all entries."""
+        with self._lock:
+            self._cache.clear()
+
+    def __len__(self) -> int:
+        with self._lock:
+            return len(self._cache)
+
+    @property
+    def capacity(self) -> int:
+        return self._capacity

code_muse/models_cache/cache_writer.py

@@ -0,0 +1,86 @@
+"""Write models_cache.json to the data directory.
+
+Ports Codex's models_cache.json pattern: a snapshot of the full model
+catalog bundled with the client, refreshed on startup.
+"""
+
+import dataclasses
+import json
+import logging
+import os
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import Any
+
+from code_muse.config import DATA_DIR
+from code_muse.models_dev_parser import (
+    BUNDLED_JSON_FILENAME,
+    ModelInfo,
+    ModelsDevRegistry,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _client_version() -> str:
+    """Return the installed code-muse version."""
+    try:
+        from importlib.metadata import version
+
+        return version("code-muse")
+    except Exception:
+        return "unknown"
+
+
+def _model_to_dict(model: ModelInfo, priority: int) -> dict[str, Any]:
+    """Serialize a ModelInfo into the cache entry format."""
+    base = {
+        "slug": model.full_id,
+        "display_name": model.name,
+        "description": "",
+        "visibility": "List",
+        "priority": priority,
+    }
+    # Include all dataclass fields for completeness
+    fields = dataclasses.asdict(model)
+    # Avoid duplicating keys we already set explicitly
+    for key, value in fields.items():
+        if key not in base:
+            base[key] = value
+    return base
+
+
+def write_models_cache(models: list[ModelInfo] | None = None) -> Path | None:
+    """Write the models cache file.
+
+    Args:
+        models: Optional explicit list of ModelInfo objects. If None,
+            loads from the bundled models_dev_api.json.
+
+    Returns:
+        Path to the written file, or None if writing failed.
+    """
+    try:
+        if models is None:
+            bundled_path = Path(__file__).parent.parent / BUNDLED_JSON_FILENAME
+            registry = ModelsDevRegistry(bundled_path)
+            models = list(registry.models.values())
+
+        cache_data: dict[str, Any] = {
+            "fetched_at": datetime.now(UTC).isoformat(),
+            "etag": None,
+            "client_version": _client_version(),
+            "models": [_model_to_dict(m, priority=idx) for idx, m in enumerate(models)],
+        }
+
+        cache_path = Path(DATA_DIR) / "models_cache.json"
+        os.makedirs(cache_path.parent, exist_ok=True)
+        with open(cache_path, "w", encoding="utf-8") as f:
+            json.dump(cache_data, f, indent=2, default=str)
+
+        logger.info(f"Wrote models cache ({len(models)} models) to {cache_path}")
+        return cache_path
+
+    except (OSError, ValueError, TypeError) as exc:
+        logger.warning(f"Failed to write models cache: {exc}")
+        return None

code_muse/models_cache/sha256_hash.pyx

@@ -0,0 +1,34 @@
+# cython: language_level=3
+"""SHA-256 content hash utility for content-addressed cache keys."""
+
+import hashlib
+from pathlib import Path
+
+
+def sha256_digest(data: bytes) -> str:
+    """Return the SHA-256 hex digests of the given bytes.
+
+    Returns a 64-character lowercase hex string.
+    """
+    return hashlib.sha256(data).hexdigest()
+
+
+def sha256_digest_file(path: Path) -> str:
+    """Return the SHA-256 hex digests of a file's contents.
+
+    Streams the file in 64KB chunks to handle large files efficiently.
+    The chunk-reading loop uses typed locals and bound methods to avoid
+    Python-level attribute lookups.
+    """
+    cdef object hasher = hashlib.sha256()
+    cdef bytes chunk
+    cdef object update = hasher.update
+    cdef object read
+    with open(path, "rb") as f:
+        read = f.read
+        while True:
+            chunk = read(65536)
+            if len(chunk) == 0:
+                break
+            update(chunk)
+    return hasher.hexdigest()

code_muse/models_cache/startup_integration.py

@@ -0,0 +1,75 @@
+"""Startup cache integration — load cached models on session start."""
+
+import json
+import logging
+from datetime import UTC, datetime, timedelta
+from pathlib import Path
+from typing import Any
+
+from code_muse.config import MODELS_CACHE_FILE
+from code_muse.models_cache.cache_writer import write_models_cache
+
+logger = logging.getLogger(__name__)
+
+MODELS_CACHE_PATH = Path(MODELS_CACHE_FILE)
+CACHE_TTL = timedelta(hours=24)
+
+
+def load_cached_models() -> list[dict[str, Any]] | None:
+    """Load models from cache if fresh (< 24h old).
+
+    Returns:
+        List of model dicts if cache is fresh, None if stale or missing.
+        Never raises — logs errors and returns None.
+    """
+    try:
+        if not MODELS_CACHE_PATH.exists():
+            logger.debug("models_cache.json not found")
+            return None
+
+        with open(MODELS_CACHE_PATH, encoding="utf-8") as f:
+            cache = json.load(f)
+
+        fetched_at_str = cache.get("fetched_at")
+        if not fetched_at_str:
+            logger.warning("models_cache.json missing fetched_at")
+            return None
+
+        fetched_at = datetime.fromisoformat(fetched_at_str)
+        age = datetime.now(UTC) - fetched_at
+
+        if age > CACHE_TTL:
+            logger.debug(f"models_cache.json is {age} old, needs refresh")
+            return None
+
+        models = cache.get("models")
+        if not isinstance(models, list):
+            logger.warning("models_cache.json models is not a list")
+            return None
+
+        logger.info(f"Using cached models ({len(models)} models, {age} old)")
+        return models
+
+    except (OSError, ValueError) as exc:
+        logger.warning(f"Failed to load models_cache.json: {exc}")
+        return None
+
+
+def refresh_models_cache() -> list[dict[str, Any]]:
+    """Fetch fresh models and update cache.
+
+    Returns:
+        Fresh list of model dicts.
+    """
+    write_models_cache()
+    cached = load_cached_models()
+    if cached is not None:
+        return cached
+    # If cache write failed, load from bundled as ultimate fallback
+    from code_muse.models_dev_parser import ModelsDevRegistry
+
+    bundled_path = Path(__file__).parent.parent / "models_dev_api.json"
+    registry = ModelsDevRegistry(bundled_path)
+    return [
+        {"slug": m.full_id, "display_name": m.name} for m in registry.models.values()
+    ]