tokenable 1.0.0__py3-none-any.whl

tokenable/providers/__init__.py

@@ -0,0 +1,485 @@
+"""Pricing database — model lookup, cost calculation, capability inference."""
+
+from __future__ import annotations
+
+import json
+import re
+from importlib.resources import files
+from typing import TYPE_CHECKING
+
+from tokenable.models import (
+    Capability,
+    ModelPricing,
+    ModelStatus,
+    ModelTier,
+    Provider,
+)
+
+# ── Tier thresholds ──────────────────────────────────────────────────
+
+_TIER_BUDGET_MAX = 5.0
+_TIER_PREMIUM_MIN = 20.0
+
+
+def compute_tier(output_cost_per_million: float) -> ModelTier:
+    """Derive cost tier from output pricing."""
+    if output_cost_per_million <= _TIER_BUDGET_MAX:
+        return ModelTier.BUDGET
+    if output_cost_per_million < _TIER_PREMIUM_MIN:
+        return ModelTier.MID
+    return ModelTier.PREMIUM
+
+
+# ── Data loading ─────────────────────────────────────────────────────
+
+_PROVIDER_DATA: dict[Provider, list[ModelPricing]] = {}
+_BENCHMARKS: dict[str, dict] = {}
+
+
+def _load_provider_json(provider: Provider) -> list[ModelPricing]:
+    """Load and parse a provider's pricing JSON."""
+    data_dir = files("tokenable.providers") / "data"
+    raw = (data_dir / f"{provider.value}.json").read_text("utf-8")
+    data = json.loads(raw)
+    models: list[ModelPricing] = []
+    for m in data.get("models", []):
+        caps = [Capability(c) for c in m.get("capabilities", [])]
+        status = ModelStatus(m.get("status", "current"))
+        tier = compute_tier(m.get("output_cost_per_million", 0))
+        models.append(
+            ModelPricing(
+                id=m["id"],
+                name=m["name"],
+                provider=provider,
+                aliases=m.get("aliases", []),
+                status=status,
+                input_cost_per_million=m["input_cost_per_million"],
+                output_cost_per_million=m["output_cost_per_million"],
+                cache_read_input_cost_per_million=m.get("cache_read_input_cost_per_million"),
+                cache_write_input_cost_per_million=m.get("cache_write_input_cost_per_million"),
+                batch_input_cost_per_million=m.get("batch_input_cost_per_million"),
+                batch_output_cost_per_million=m.get("batch_output_cost_per_million"),
+                fast_input_cost_per_million=m.get("fast_input_cost_per_million"),
+                fast_output_cost_per_million=m.get("fast_output_cost_per_million"),
+                input_cost_above_200k_per_million=m.get("input_cost_above_200k_per_million"),
+                output_cost_above_200k_per_million=m.get("output_cost_above_200k_per_million"),
+                context_window=m["context_window"],
+                max_output_tokens=m["max_output_tokens"],
+                supports_vision=m.get("supports_vision", False),
+                supports_tools=m.get("supports_tools", False),
+                supports_prompt_caching=m.get("supports_prompt_caching", False),
+                supports_reasoning=m.get("supports_reasoning", False),
+                supports_computer_use=m.get("supports_computer_use", False),
+                tier=tier,
+                capabilities=caps,
+                knowledge_cutoff=m.get("knowledge_cutoff"),
+            )
+        )
+    return models
+
+
+def _ensure_loaded() -> None:
+    """Lazy-load all provider data."""
+    if _PROVIDER_DATA:
+        return
+    for p in Provider:
+        _PROVIDER_DATA[p] = _load_provider_json(p)
+    # Load benchmarks
+    data_dir = files("tokenable.providers") / "data"
+    raw = (data_dir / "benchmarks.json").read_text("utf-8")
+    data = json.loads(raw)
+    _BENCHMARKS.update(data.get("models", {}))
+
+
+# ── Public API ───────────────────────────────────────────────────────
+
+
+def get_all_providers() -> list[Provider]:
+    """Get all supported providers."""
+    return list(Provider)
+
+
+def get_provider_models(provider: Provider) -> list[ModelPricing]:
+    """Get all models for a provider."""
+    _ensure_loaded()
+    return _PROVIDER_DATA.get(provider, [])
+
+
+def normalize_model_id(model_id: str) -> str:
+    """Strip platform/routing prefixes and version suffixes from model IDs."""
+    mid = model_id
+    mid = re.sub(r"^openrouter/", "", mid)
+    mid = re.sub(r"^(aws|gcp|azure)/(anthropic|openai|google|meta|mistralai|cohere|ai21)/", "", mid)
+    mid = re.sub(r"^(bedrock/|azure/|vertex_ai/|azure_ai/)", "", mid)
+    mid = re.sub(
+        r"^(models/|gemini/|xai/|openai/|perplexity/|anthropic/|google/|meta-llama/|meta/|mistralai/)",
+        "",
+        mid,
+    )
+    mid = re.sub(r"^bedrock-", "", mid)
+    mid = re.sub(r"^(anthropic|amazon|meta|cohere|ai21|mistral|stability)\.", "", mid)
+    mid = re.sub(r"-v\d+:\d+$", "", mid)
+    return mid
+
+
+def get_model(provider: Provider, model_id: str) -> ModelPricing | None:
+    """Look up a model by canonical ID or alias."""
+    _ensure_loaded()
+    models = _PROVIDER_DATA.get(provider, [])
+
+    # Exact match on ID
+    for m in models:
+        if m.id == model_id:
+            return m
+
+    # Exact match on alias
+    for m in models:
+        if model_id in m.aliases:
+            return m
+
+    # Normalized match
+    normalized = normalize_model_id(model_id)
+    if normalized != model_id:
+        for m in models:
+            if m.id == normalized:
+                return m
+        for m in models:
+            if normalized in m.aliases:
+                return m
+
+    return None
+
+
+def get_all_models() -> list[ModelPricing]:
+    """Get all models across all providers."""
+    _ensure_loaded()
+    result: list[ModelPricing] = []
+    for models in _PROVIDER_DATA.values():
+        result.extend(models)
+    return result
+
+
+def calculate_cost(
+    model: ModelPricing,
+    input_tokens: int,
+    output_tokens: int,
+    cached_input_tokens: int = 0,
+    use_batch: bool = False,
+    use_fast: bool = False,
+    is_long_context: bool | None = None,
+) -> float:
+    """Calculate total USD cost for a request."""
+    if is_long_context is None:
+        is_long_context = input_tokens > 200_000
+
+    effective_cached = min(cached_input_tokens, input_tokens)
+    uncached_input = max(0, input_tokens - effective_cached)
+
+    # Input rate
+    if use_fast and model.fast_input_cost_per_million is not None:
+        input_rate = model.fast_input_cost_per_million
+    elif use_batch and model.batch_input_cost_per_million is not None:
+        input_rate = model.batch_input_cost_per_million
+    elif is_long_context and model.input_cost_above_200k_per_million is not None:
+        input_rate = model.input_cost_above_200k_per_million
+    else:
+        input_rate = model.input_cost_per_million
+
+    # Output rate
+    if use_fast and model.fast_output_cost_per_million is not None:
+        output_rate = model.fast_output_cost_per_million
+    elif use_batch and model.batch_output_cost_per_million is not None:
+        output_rate = model.batch_output_cost_per_million
+    elif is_long_context and model.output_cost_above_200k_per_million is not None:
+        output_rate = model.output_cost_above_200k_per_million
+    else:
+        output_rate = model.output_cost_per_million
+
+    input_cost = (uncached_input / 1_000_000) * input_rate
+    cached_cost = (
+        (effective_cached / 1_000_000) * model.cache_read_input_cost_per_million
+        if effective_cached > 0 and model.cache_read_input_cost_per_million is not None
+        else 0.0
+    )
+    output_cost = (output_tokens / 1_000_000) * output_rate
+
+    return input_cost + cached_cost + output_cost
+
+
+# ── Capability inference ─────────────────────────────────────────────
+
+_CAPABILITY_KEYWORDS: list[tuple[Capability, re.Pattern[str]]] = [
+    (
+        Capability.CODE,
+        re.compile(
+            r"\b(code|coding|debug|refactor|function|typescript|python|javascript|programming|syntax|compile|lint|regex|sql|api endpoint)\b",
+            re.IGNORECASE,
+        ),
+    ),
+    (
+        Capability.REASONING,
+        re.compile(
+            r"\b(step[- ]by[- ]step|analyze|reason|think carefully|evaluate|logic|math|proof|calculate|deduce|chain of thought)\b",
+            re.IGNORECASE,
+        ),
+    ),
+    (
+        Capability.CREATIVE,
+        re.compile(
+            r"\b(story|creative|poem|narrative|fiction|blog post|essay|copywriting)\b",
+            re.IGNORECASE,
+        ),
+    ),
+    (
+        Capability.VISION,
+        re.compile(
+            r"\b(image|screenshot|photo|picture|diagram|chart|visual|ocr|describe the image)\b",
+            re.IGNORECASE,
+        ),
+    ),
+    (
+        Capability.SEARCH,
+        re.compile(
+            r"\b(search|find information|look up|latest news|real[- ]time)\b", re.IGNORECASE
+        ),
+    ),
+    (
+        Capability.AUDIO,
+        re.compile(r"\b(audio|transcribe|speech|voice|listen)\b", re.IGNORECASE),
+    ),
+]
+
+
+def infer_required_capabilities(text: str) -> list[Capability]:
+    """Infer required capabilities from free-form text."""
+    if not text.strip():
+        return [Capability.GENERAL]
+    matched: set[Capability] = set()
+    for cap, pattern in _CAPABILITY_KEYWORDS:
+        if pattern.search(text):
+            matched.add(cap)
+    return list(matched) if matched else [Capability.GENERAL]
+
+
+def get_models_by_capabilities(
+    required: list[Capability],
+    provider: Provider | None = None,
+    status: ModelStatus = ModelStatus.CURRENT,
+    max_output_cost_per_million: float | None = None,
+) -> list[ModelPricing]:
+    """Get all models that support ALL specified capabilities."""
+    _ensure_loaded()
+    source = get_provider_models(provider) if provider else get_all_models()
+    result = []
+    for m in source:
+        if m.status != status:
+            continue
+        if (
+            max_output_cost_per_million is not None
+            and m.output_cost_per_million > max_output_cost_per_million
+        ):
+            continue
+        if all(cap in m.capabilities for cap in required):
+            result.append(m)
+    result.sort(key=lambda m: m.output_cost_per_million)
+    return result
+
+
+# ── Benchmarks ───────────────────────────────────────────────────────
+
+_CAP_TO_BENCHMARK: dict[Capability, str] = {
+    Capability.CODE: "coding",
+    Capability.REASONING: "reasoning",
+    Capability.GENERAL: "overall",
+    Capability.CREATIVE: "creative_writing",
+    Capability.VISION: "overall",
+    Capability.SEARCH: "overall",
+    Capability.AUDIO: "overall",
+}
+
+
+def get_model_benchmarks(provider: Provider, model_id: str) -> dict | None:
+    """Get benchmark scores for a model."""
+    _ensure_loaded()
+    key = f"{provider.value}/{model_id}"
+    if key in _BENCHMARKS:
+        return _BENCHMARKS[key]
+    normalized = normalize_model_id(model_id)
+    if normalized != model_id:
+        norm_key = f"{provider.value}/{normalized}"
+        if norm_key in _BENCHMARKS:
+            return _BENCHMARKS[norm_key]
+    model = get_model(provider, model_id)
+    if model and model.id != model_id:
+        canon_key = f"{provider.value}/{model.id}"
+        if canon_key in _BENCHMARKS:
+            return _BENCHMARKS[canon_key]
+    return None
+
+
+def get_quality_score(provider: Provider, model_id: str, capability: Capability) -> float | None:
+    """Get quality score for a model on a specific capability."""
+    bench = get_model_benchmarks(provider, model_id)
+    if not bench:
+        return None
+    field = _CAP_TO_BENCHMARK.get(capability, "overall")
+    score = bench.get(field)
+    return float(score) if isinstance(score, (int, float)) else bench.get("overall")
+
+
+def get_min_quality_score(
+    provider: Provider, model_id: str, capabilities: list[Capability]
+) -> float | None:
+    """Get minimum quality score across multiple capabilities."""
+    scores: list[float] = []
+    for cap in capabilities:
+        score = get_quality_score(provider, model_id, cap)
+        if score is None:
+            return None
+        scores.append(score)
+    return min(scores) if scores else None
+
+
+# ── Model suggestion ─────────────────────────────────────────────────
+
+_TIER_ORDER: dict[ModelTier, int] = {ModelTier.BUDGET: 0, ModelTier.MID: 1, ModelTier.PREMIUM: 2}
+_DEFAULT_MIN_QUALITY_RATIO = 0.7
+
+
+class AlternativeSuggestion:
+    """A suggested alternative model."""
+
+    def __init__(
+        self,
+        model: ModelPricing,
+        reasoning: str,
+        savings_percent: int,
+        quality_score: float | None = None,
+        current_quality_score: float | None = None,
+    ):
+        self.model = model
+        self.reasoning = reasoning
+        self.savings_percent = savings_percent
+        self.quality_score = quality_score
+        self.current_quality_score = current_quality_score
+
+
+def suggest_alternatives(
+    current_model_id: str,
+    current_provider: Provider,
+    required_capabilities: list[Capability],
+    confidence: str = "medium",
+    min_quality_ratio: float = _DEFAULT_MIN_QUALITY_RATIO,
+) -> list[AlternativeSuggestion]:
+    """Suggest cheaper alternatives that match required capabilities."""
+    current = get_model(current_provider, current_model_id)
+    if not current:
+        return []
+
+    current_quality = get_min_quality_score(
+        current_provider, current_model_id, required_capabilities
+    )
+
+    candidates = get_models_by_capabilities(required_capabilities)
+    candidates = [
+        m
+        for m in candidates
+        if not (m.id == current.id and m.provider == current.provider)
+        and m.output_cost_per_million < current.output_cost_per_million
+    ]
+
+    # Filter by confidence
+    filtered = []
+    for m in candidates:
+        current_tier = _TIER_ORDER[current.tier]
+        candidate_tier = _TIER_ORDER[m.tier]
+        tier_drop = current_tier - candidate_tier
+
+        if confidence == "low" and (m.provider != current_provider or tier_drop > 0):
+            continue
+        if confidence == "medium" and tier_drop > 1:
+            continue
+        # Quality gate
+        if current_quality is not None:
+            cq = get_min_quality_score(m.provider, m.id, required_capabilities)
+            if cq is not None and cq < current_quality * min_quality_ratio:
+                continue
+        filtered.append(m)
+
+    # Sort by quality-adjusted cost
+    ref_input, ref_output = 4096, 1024
+    current_cost = calculate_cost(current, ref_input, ref_output)
+
+    def sort_key(m: ModelPricing) -> float:
+        cost = calculate_cost(m, ref_input, ref_output)
+        q = get_min_quality_score(m.provider, m.id, required_capabilities)
+        return cost / (q / 100) if q else cost
+
+    filtered.sort(key=sort_key)
+
+    results: list[AlternativeSuggestion] = []
+    for m in filtered[:3]:
+        candidate_cost = calculate_cost(m, ref_input, ref_output)
+        savings = (
+            round(((current_cost - candidate_cost) / current_cost) * 100) if current_cost > 0 else 0
+        )
+        caps = f"[{', '.join(c.value for c in required_capabilities)}]"
+        cq = get_min_quality_score(m.provider, m.id, required_capabilities)
+        quality_suffix = (
+            f" (quality: {cq} vs {current_quality})"
+            if cq is not None and current_quality is not None
+            else ""
+        )
+        if m.provider == current_provider:
+            reasoning = f"Task requires {caps} — {m.name} handles that at {savings}% lower cost{quality_suffix}"
+        else:
+            reasoning = f"Task requires {caps} — {m.provider.value}/{m.name} handles that at {savings}% lower cost{quality_suffix}"
+        results.append(
+            AlternativeSuggestion(
+                model=m,
+                reasoning=reasoning,
+                savings_percent=savings,
+                quality_score=cq,
+                current_quality_score=current_quality,
+            )
+        )
+    return results
+
+
+class TaskSuggestion:
+    """Result of a task-based model suggestion."""
+
+    def __init__(
+        self, model: ModelPricing, inferred_capabilities: list[Capability], reasoning: str
+    ):
+        self.model = model
+        self.inferred_capabilities = inferred_capabilities
+        self.reasoning = reasoning
+
+
+def suggest_model_for_task(
+    text: str,
+    provider: Provider | None = None,
+    max_cost_per_million: float | None = None,
+) -> TaskSuggestion | None:
+    """Suggest the best value model for a free-form task description."""
+    capabilities = infer_required_capabilities(text)
+    candidates = get_models_by_capabilities(
+        capabilities,
+        provider=provider,
+        max_output_cost_per_million=max_cost_per_million,
+    )
+    if not candidates:
+        return None
+
+    def sort_key(m: ModelPricing) -> float:
+        q = get_min_quality_score(m.provider, m.id, capabilities)
+        return m.output_cost_per_million / (q / 100) if q else m.output_cost_per_million
+
+    sorted_candidates = sorted(candidates, key=sort_key)
+    best = sorted_candidates[0]
+    caps = f"[{', '.join(c.value for c in capabilities)}]"
+    quality = get_min_quality_score(best.provider, best.id, capabilities)
+    quality_suffix = f" — quality: {quality}/100" if quality is not None else ""
+    reasoning = f"Inferred capabilities: {caps} — {best.provider.value}/{best.name} (${best.output_cost_per_million}/M output) is the best value model{quality_suffix}"
+    return TaskSuggestion(model=best, inferred_capabilities=capabilities, reasoning=reasoning)