overrides.py

"""User-configured rate overrides for the profiler."""
from __future__ import annotations

import logging
from pathlib import Path
from typing import Dict, Optional

logger = logging.getLogger(__name__)

# Try to import PyYAML, but provide a fallback
try:
    import yaml as _yaml
    HAS_YAML = True
except ImportError:
    _yaml = None
    HAS_YAML = False


class UserOverrides:
    """Load and query user-defined rate overrides from YAML config.

    Supports runtime reloading (bug #25) via reload() — the user can edit
    profiler_overrides.yaml and the next get_rate() call sees the new values
    without requiring a Hermes restart.
    """

    CONFIG_PATH = Path.home() / ".hermes" / "profiler_overrides.yaml"

    def __init__(self):
        self._overrides: Dict[str, Dict[str, float]] = self._load()
        self._config_mtime: Optional[float] = (
            self.CONFIG_PATH.stat().st_mtime if self.CONFIG_PATH.exists() else None
        )

    def reload(self) -> None:
        """Re-read the YAML config from disk."""
        self._overrides = self._load()
        try:
            self._config_mtime = (
                self.CONFIG_PATH.stat().st_mtime if self.CONFIG_PATH.exists() else None
            )
        except OSError:
            self._config_mtime = None

    def _maybe_reload(self) -> None:
        """Auto-reload if the file changed on disk since last load."""
        try:
            if self.CONFIG_PATH.exists():
                mtime = self.CONFIG_PATH.stat().st_mtime
                if self._config_mtime is None or mtime > self._config_mtime:
                    self.reload()
        except OSError:
            pass

    def _load(self) -> Dict[str, Dict[str, float]]:
        """Load overrides from YAML file."""
        if not self.CONFIG_PATH.exists():
            return {}

        if not HAS_YAML:
            logger.warning("PyYAML not installed; user overrides unavailable")
            return {}

        try:
            with open(self.CONFIG_PATH) as f:
                data = _yaml.safe_load(f) or {}
            overrides = data.get("overrides", {}) or {}
            # Strip whitespace from keys (bug #35)
            return {
                (k.strip() if isinstance(k, str) else k): v
                for k, v in overrides.items()
            }
        except Exception as e:
            logger.warning("Failed to load user overrides: %s", e)
            return {}

    def get_rate(self, provider: str, model: str) -> Optional[Dict[str, float]]:
        """Get user override rate for a provider/model combo.

        Auto-reloads the YAML file if it has been modified since last load.
        """
        self._maybe_reload()
        provider_model = f"{provider}/{model}".strip()
        return self._overrides.get(provider_model)