prompture 0.0.33.dev1__py3-none-any.whl → 0.0.34__py3-none-any.whl

prompture/model_rates.py

@@ -0,0 +1,217 @@
+"""Live model rates from models.dev API with local caching.
+
+Fetches pricing and metadata for LLM models from https://models.dev/api.json,
+caches locally with TTL-based auto-refresh, and provides lookup functions
+used by drivers for cost calculations.
+"""
+
+import contextlib
+import json
+import logging
+import threading
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Optional
+
+logger = logging.getLogger(__name__)
+
+# Maps prompture provider names to models.dev provider names
+PROVIDER_MAP: dict[str, str] = {
+    "openai": "openai",
+    "claude": "anthropic",
+    "google": "google",
+    "groq": "groq",
+    "grok": "xai",
+    "azure": "azure",
+    "openrouter": "openrouter",
+}
+
+_API_URL = "https://models.dev/api.json"
+_CACHE_DIR = Path.home() / ".prompture" / "cache"
+_CACHE_FILE = _CACHE_DIR / "models_dev.json"
+_META_FILE = _CACHE_DIR / "models_dev_meta.json"
+
+_lock = threading.Lock()
+_data: Optional[dict[str, Any]] = None
+_loaded = False
+
+
+def _get_ttl_days() -> int:
+    """Get TTL from settings if available, otherwise default to 7."""
+    try:
+        from .settings import settings
+
+        return getattr(settings, "model_rates_ttl_days", 7)
+    except Exception:
+        return 7
+
+
+def _cache_is_valid() -> bool:
+    """Check whether the local cache exists and is within TTL."""
+    if not _CACHE_FILE.exists() or not _META_FILE.exists():
+        return False
+    try:
+        meta = json.loads(_META_FILE.read_text(encoding="utf-8"))
+        fetched_at = datetime.fromisoformat(meta["fetched_at"])
+        ttl_days = meta.get("ttl_days", _get_ttl_days())
+        age = datetime.now(timezone.utc) - fetched_at
+        return age.total_seconds() < ttl_days * 86400
+    except Exception:
+        return False
+
+
+def _write_cache(data: dict[str, Any]) -> None:
+    """Write API data and metadata to local cache."""
+    try:
+        _CACHE_DIR.mkdir(parents=True, exist_ok=True)
+        _CACHE_FILE.write_text(json.dumps(data), encoding="utf-8")
+        meta = {
+            "fetched_at": datetime.now(timezone.utc).isoformat(),
+            "ttl_days": _get_ttl_days(),
+        }
+        _META_FILE.write_text(json.dumps(meta), encoding="utf-8")
+    except Exception as exc:
+        logger.debug("Failed to write model rates cache: %s", exc)
+
+
+def _read_cache() -> Optional[dict[str, Any]]:
+    """Read cached API data from disk."""
+    try:
+        return json.loads(_CACHE_FILE.read_text(encoding="utf-8"))
+    except Exception:
+        return None
+
+
+def _fetch_from_api() -> Optional[dict[str, Any]]:
+    """Fetch fresh data from models.dev API."""
+    try:
+        import requests
+
+        resp = requests.get(_API_URL, timeout=15)
+        resp.raise_for_status()
+        return resp.json()
+    except Exception as exc:
+        logger.debug("Failed to fetch model rates from %s: %s", _API_URL, exc)
+        return None
+
+
+def _ensure_loaded() -> Optional[dict[str, Any]]:
+    """Lazy-load data: use cache if valid, otherwise fetch from API."""
+    global _data, _loaded
+    if _loaded:
+        return _data
+
+    with _lock:
+        # Double-check after acquiring lock
+        if _loaded:
+            return _data
+
+        if _cache_is_valid():
+            _data = _read_cache()
+            if _data is not None:
+                _loaded = True
+                return _data
+
+        # Cache missing or expired — fetch fresh
+        fresh = _fetch_from_api()
+        if fresh is not None:
+            _data = fresh
+            _write_cache(fresh)
+        else:
+            # Fetch failed — try stale cache as last resort
+            _data = _read_cache()
+
+        _loaded = True
+        return _data
+
+
+def _lookup_model(provider: str, model_id: str) -> Optional[dict[str, Any]]:
+    """Find a model entry in the cached data.
+
+    The API structure is ``{provider: {model_id: {...}, ...}, ...}``.
+    """
+    data = _ensure_loaded()
+    if data is None:
+        return None
+
+    api_provider = PROVIDER_MAP.get(provider, provider)
+    provider_data = data.get(api_provider)
+    if not isinstance(provider_data, dict):
+        return None
+
+    return provider_data.get(model_id)
+
+
+# ── Public API ──────────────────────────────────────────────────────────────
+
+
+def get_model_rates(provider: str, model_id: str) -> Optional[dict[str, float]]:
+    """Return pricing dict for a model, or ``None`` if unavailable.
+
+    Returned keys mirror models.dev cost fields (per 1M tokens):
+    ``input``, ``output``, and optionally ``cache_read``, ``cache_write``,
+    ``reasoning``.
+    """
+    entry = _lookup_model(provider, model_id)
+    if entry is None:
+        return None
+
+    cost = entry.get("cost")
+    if not isinstance(cost, dict):
+        return None
+
+    rates: dict[str, float] = {}
+    for key in ("input", "output", "cache_read", "cache_write", "reasoning"):
+        val = cost.get(key)
+        if val is not None:
+            with contextlib.suppress(TypeError, ValueError):
+                rates[key] = float(val)
+
+    # Must have at least input and output to be useful
+    if "input" in rates and "output" in rates:
+        return rates
+    return None
+
+
+def get_model_info(provider: str, model_id: str) -> Optional[dict[str, Any]]:
+    """Return full model metadata (cost, limits, capabilities), or ``None``."""
+    return _lookup_model(provider, model_id)
+
+
+def get_all_provider_models(provider: str) -> list[str]:
+    """Return list of model IDs available for a provider."""
+    data = _ensure_loaded()
+    if data is None:
+        return []
+
+    api_provider = PROVIDER_MAP.get(provider, provider)
+    provider_data = data.get(api_provider)
+    if not isinstance(provider_data, dict):
+        return []
+
+    return list(provider_data.keys())
+
+
+def refresh_rates_cache(force: bool = False) -> bool:
+    """Fetch fresh data from models.dev.
+
+    Args:
+        force: If ``True``, fetch even when the cache is still within TTL.
+
+    Returns:
+        ``True`` if fresh data was fetched and cached successfully.
+    """
+    global _data, _loaded
+
+    with _lock:
+        if not force and _cache_is_valid():
+            return False
+
+        fresh = _fetch_from_api()
+        if fresh is not None:
+            _data = fresh
+            _write_cache(fresh)
+            _loaded = True
+            return True
+
+        return False

prompture/runner.py

@@ -1,13 +1,15 @@
 """Test suite runner for executing JSON validation tests across multiple models."""
-from typing import Dict, Any, List
 
-from .core import ask_for_json, Driver
+from typing import Any
+
 from prompture.validator import validate_against_schema
 
+from .core import Driver, ask_for_json
+
 
-def run_suite_from_spec(spec: Dict[str, Any], drivers: Dict[str, Driver]) -> Dict[str, Any]:
+def run_suite_from_spec(spec: dict[str, Any], drivers: dict[str, Driver]) -> dict[str, Any]:
     """Run a test suite specified by a spec dictionary across multiple models.
-    
+
     Args:
         spec: A dictionary containing the test suite specification with the structure:
             {
@@ -21,7 +23,7 @@ def run_suite_from_spec(spec: Dict[str, Any], drivers: Dict[str, Driver]) -> Dic
                 }, ...]
             }
         drivers: A dictionary mapping driver names to driver instances
-        
+
     Returns:
         A dictionary containing test results with the structure:
             {
@@ -42,67 +44,67 @@ def run_suite_from_spec(spec: Dict[str, Any], drivers: Dict[str, Driver]) -> Dic
             }
     """
     results = []
-    
+
     for test in spec["tests"]:
         for model in spec["models"]:
             driver = drivers.get(model["driver"])
             if not driver:
                 continue
-                
+
             # Run test for each input
             for input_data in test["inputs"]:
                 # Format prompt template with input data
                 try:
                     prompt = test["prompt_template"].format(**input_data)
                 except KeyError as e:
-                    results.append({
-                        "test_id": test["id"],
-                        "model_id": model["id"],
-                        "input": input_data,
-                        "prompt": test["prompt_template"],
-                        "error": f"Template formatting error: missing key {e}",
-                        "validation": {"ok": False, "error": "Prompt formatting failed", "data": None},
-                        "usage": {"total_tokens": 0, "cost": 0}
-                    })
+                    results.append(
+                        {
+                            "test_id": test["id"],
+                            "model_id": model["id"],
+                            "input": input_data,
+                            "prompt": test["prompt_template"],
+                            "error": f"Template formatting error: missing key {e}",
+                            "validation": {"ok": False, "error": "Prompt formatting failed", "data": None},
+                            "usage": {"total_tokens": 0, "cost": 0},
+                        }
+                    )
                     continue
-                    
+
                 # Get JSON response from model
                 try:
                     response = ask_for_json(
                         driver=driver,
                         content_prompt=prompt,
                         json_schema=test["schema"],
-                        options=model.get("options", {})
+                        options=model.get("options", {}),
                     )
-                    
+
                     # Validate response against schema
-                    validation = validate_against_schema(
-                        response["json_string"],
-                        test["schema"]
+                    validation = validate_against_schema(response["json_string"], test["schema"])
+
+                    results.append(
+                        {
+                            "test_id": test["id"],
+                            "model_id": model["id"],
+                            "input": input_data,
+                            "prompt": prompt,
+                            "response": response["json_object"],
+                            "validation": validation,
+                            "usage": response["usage"],
+                        }
                     )
-                    
-                    results.append({
-                        "test_id": test["id"],
-                        "model_id": model["id"],
-                        "input": input_data,
-                        "prompt": prompt,
-                        "response": response["json_object"],
-                        "validation": validation,
-                        "usage": response["usage"]
-                    })
-                    
+
                 except Exception as e:
-                    results.append({
-                        "test_id": test["id"],
-                        "model_id": model["id"],
-                        "input": input_data,
-                        "prompt": prompt,
-                        "error": str(e),
-                        "validation": {"ok": False, "error": "Model response error", "data": None},
-                        "usage": {"total_tokens": 0, "cost": 0}
-                    })
-    
-    return {
-        "meta": spec.get("meta", {}),
-        "results": results
-    }
+                    results.append(
+                        {
+                            "test_id": test["id"],
+                            "model_id": model["id"],
+                            "input": input_data,
+                            "prompt": prompt,
+                            "error": str(e),
+                            "validation": {"ok": False, "error": "Model response error", "data": None},
+                            "usage": {"total_tokens": 0, "cost": 0},
+                        }
+                    )
+
+    return {"meta": spec.get("meta", {}), "results": results}

prompture/session.py

@@ -0,0 +1,117 @@
+"""Usage session tracking for Prompture.
+
+Provides :class:`UsageSession` which accumulates token counts, costs, and
+errors across multiple driver calls.  A session instance is compatible as
+both an ``on_response`` and ``on_error`` callback, so you can wire it
+directly into :class:`~prompture.callbacks.DriverCallbacks`.
+
+Usage::
+
+    from prompture import UsageSession, DriverCallbacks
+
+    session = UsageSession()
+    callbacks = DriverCallbacks(
+        on_response=session.record,
+        on_error=session.record_error,
+    )
+
+    # ... pass *callbacks* to your driver / conversation ...
+
+    print(session.summary()["formatted"])
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass
+class UsageSession:
+    """Accumulates usage statistics across multiple driver calls."""
+
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+    total_tokens: int = 0
+    total_cost: float = 0.0
+    call_count: int = 0
+    errors: int = 0
+    _per_model: dict[str, dict[str, Any]] = field(default_factory=dict, repr=False)
+
+    # ------------------------------------------------------------------ #
+    # Recording
+    # ------------------------------------------------------------------ #
+
+    def record(self, response_info: dict[str, Any]) -> None:
+        """Record a successful driver response.
+
+        Compatible as an ``on_response`` callback for
+        :class:`~prompture.callbacks.DriverCallbacks`.
+
+        Args:
+            response_info: Payload dict with at least ``meta`` and
+                optionally ``driver`` keys.
+        """
+        meta = response_info.get("meta", {})
+        pt = meta.get("prompt_tokens", 0)
+        ct = meta.get("completion_tokens", 0)
+        tt = meta.get("total_tokens", 0)
+        cost = meta.get("cost", 0.0)
+
+        self.prompt_tokens += pt
+        self.completion_tokens += ct
+        self.total_tokens += tt
+        self.total_cost += cost
+        self.call_count += 1
+
+        model = response_info.get("driver", "unknown")
+        bucket = self._per_model.setdefault(
+            model,
+            {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0, "cost": 0.0, "calls": 0},
+        )
+        bucket["prompt_tokens"] += pt
+        bucket["completion_tokens"] += ct
+        bucket["total_tokens"] += tt
+        bucket["cost"] += cost
+        bucket["calls"] += 1
+
+    def record_error(self, error_info: dict[str, Any]) -> None:
+        """Record a driver error.
+
+        Compatible as an ``on_error`` callback for
+        :class:`~prompture.callbacks.DriverCallbacks`.
+        """
+        self.errors += 1
+
+    # ------------------------------------------------------------------ #
+    # Reporting
+    # ------------------------------------------------------------------ #
+
+    def summary(self) -> dict[str, Any]:
+        """Return a machine-readable summary with a ``formatted`` string."""
+        formatted = (
+            f"Session: {self.total_tokens:,} tokens across {self.call_count} call(s) costing ${self.total_cost:.4f}"
+        )
+        if self.errors:
+            formatted += f" ({self.errors} error(s))"
+
+        return {
+            "prompt_tokens": self.prompt_tokens,
+            "completion_tokens": self.completion_tokens,
+            "total_tokens": self.total_tokens,
+            "total_cost": self.total_cost,
+            "call_count": self.call_count,
+            "errors": self.errors,
+            "per_model": dict(self._per_model),
+            "formatted": formatted,
+        }
+
+    def reset(self) -> None:
+        """Clear all accumulated counters."""
+        self.prompt_tokens = 0
+        self.completion_tokens = 0
+        self.total_tokens = 0
+        self.total_cost = 0.0
+        self.call_count = 0
+        self.errors = 0
+        self._per_model.clear()

prompture/settings.py

@@ -1,6 +1,8 @@
-from pydantic_settings import BaseSettings, SettingsConfigDict
 from typing import Optional
 
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+
 class Settings(BaseSettings):
     """Application settings loaded from environment variables or .env file."""
 
@@ -52,6 +54,17 @@ class Settings(BaseSettings):
     airllm_model: str = "meta-llama/Llama-2-7b-hf"
     airllm_compression: Optional[str] = None  # "4bit" or "8bit"
 
+    # Model rates cache
+    model_rates_ttl_days: int = 7  # How often to refresh models.dev cache
+
+    # Response cache
+    cache_enabled: bool = False
+    cache_backend: str = "memory"
+    cache_ttl_seconds: int = 3600
+    cache_memory_maxsize: int = 256
+    cache_sqlite_path: Optional[str] = None
+    cache_redis_url: Optional[str] = None
+
     model_config = SettingsConfigDict(
         env_file=".env",
         extra="ignore",