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

prompture/model_rates.py

@@ -5,17 +5,18 @@ 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, Dict, List, Optional
+from typing import Any, Optional
 
 logger = logging.getLogger(__name__)
 
 # Maps prompture provider names to models.dev provider names
-PROVIDER_MAP: Dict[str, str] = {
+PROVIDER_MAP: dict[str, str] = {
     "openai": "openai",
     "claude": "anthropic",
     "google": "google",
@@ -31,7 +32,7 @@ _CACHE_FILE = _CACHE_DIR / "models_dev.json"
 _META_FILE = _CACHE_DIR / "models_dev_meta.json"
 
 _lock = threading.Lock()
-_data: Optional[Dict[str, Any]] = None
+_data: Optional[dict[str, Any]] = None
 _loaded = False
 
 
@@ -39,6 +40,7 @@ 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
@@ -58,7 +60,7 @@ def _cache_is_valid() -> bool:
         return False
 
 
-def _write_cache(data: Dict[str, Any]) -> None:
+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)
@@ -72,7 +74,7 @@ def _write_cache(data: Dict[str, Any]) -> None:
         logger.debug("Failed to write model rates cache: %s", exc)
 
 
-def _read_cache() -> Optional[Dict[str, Any]]:
+def _read_cache() -> Optional[dict[str, Any]]:
     """Read cached API data from disk."""
     try:
         return json.loads(_CACHE_FILE.read_text(encoding="utf-8"))
@@ -80,10 +82,11 @@ def _read_cache() -> Optional[Dict[str, Any]]:
         return None
 
 
-def _fetch_from_api() -> Optional[Dict[str, Any]]:
+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()
@@ -92,7 +95,7 @@ def _fetch_from_api() -> Optional[Dict[str, Any]]:
         return None
 
 
-def _ensure_loaded() -> Optional[Dict[str, Any]]:
+def _ensure_loaded() -> Optional[dict[str, Any]]:
     """Lazy-load data: use cache if valid, otherwise fetch from API."""
     global _data, _loaded
     if _loaded:
@@ -122,7 +125,7 @@ def _ensure_loaded() -> Optional[Dict[str, Any]]:
         return _data
 
 
-def _lookup_model(provider: str, model_id: str) -> Optional[Dict[str, Any]]:
+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: {...}, ...}, ...}``.
@@ -142,7 +145,7 @@ def _lookup_model(provider: str, model_id: str) -> Optional[Dict[str, Any]]:
 # ── Public API ──────────────────────────────────────────────────────────────
 
 
-def get_model_rates(provider: str, model_id: str) -> Optional[Dict[str, float]]:
+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):
@@ -157,14 +160,12 @@ def get_model_rates(provider: str, model_id: str) -> Optional[Dict[str, float]]:
     if not isinstance(cost, dict):
         return None
 
-    rates: Dict[str, float] = {}
+    rates: dict[str, float] = {}
     for key in ("input", "output", "cache_read", "cache_write", "reasoning"):
         val = cost.get(key)
         if val is not None:
-            try:
+            with contextlib.suppress(TypeError, ValueError):
                 rates[key] = float(val)
-            except (TypeError, ValueError):
-                pass
 
     # Must have at least input and output to be useful
     if "input" in rates and "output" in rates:
@@ -172,12 +173,12 @@ def get_model_rates(provider: str, model_id: str) -> Optional[Dict[str, float]]:
     return None
 
 
-def get_model_info(provider: str, model_id: str) -> Optional[Dict[str, Any]]:
+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]:
+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:

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."""
 
@@ -55,6 +57,14 @@ class Settings(BaseSettings):
     # 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",