tokenable 1.0.0__py3-none-any.whl

tokenable/mcp/__init__.py

@@ -0,0 +1,249 @@
+"""MCP Server — Model Context Protocol integration for AI agents."""
+
+from __future__ import annotations
+
+import json
+import sys
+from typing import Any
+
+from tokenable.models import Provider
+from tokenable.providers import (
+    calculate_cost,
+    get_all_models,
+    get_all_providers,
+    get_min_quality_score,
+    get_model,
+    infer_required_capabilities,
+    suggest_alternatives,
+    suggest_model_for_task,
+)
+
+# ── Tool handlers ────────────────────────────────────────────────────
+
+
+def handle_suggest_model(
+    task: str, provider: str | None = None, max_cost_per_million_tokens: float | None = None
+) -> dict[str, Any]:
+    """Suggest the cheapest capable model for a task."""
+    prov = Provider(provider) if provider else None
+    suggestion = suggest_model_for_task(
+        task, provider=prov, max_cost_per_million=max_cost_per_million_tokens
+    )
+
+    if not suggestion:
+        return {"error": "No model found matching the inferred capabilities for this task."}
+
+    alts = suggest_alternatives(
+        suggestion.model.id, suggestion.model.provider, suggestion.inferred_capabilities
+    )
+
+    rec_quality = get_min_quality_score(
+        suggestion.model.provider, suggestion.model.id, suggestion.inferred_capabilities
+    )
+
+    return {
+        "recommended": {
+            "provider": suggestion.model.provider.value,
+            "model": suggestion.model.id,
+            "name": suggestion.model.name,
+            "inputCostPerMillion": suggestion.model.input_cost_per_million,
+            "outputCostPerMillion": suggestion.model.output_cost_per_million,
+            "capabilities": [c.value for c in suggestion.model.capabilities],
+            **({"qualityScore": rec_quality} if rec_quality is not None else {}),
+        },
+        "alternatives": [
+            {
+                "provider": a.model.provider.value,
+                "model": a.model.id,
+                "name": a.model.name,
+                "outputCostPerMillion": a.model.output_cost_per_million,
+                "savingsPercent": a.savings_percent,
+                **({"qualityScore": a.quality_score} if a.quality_score is not None else {}),
+            }
+            for a in alts
+        ],
+        "inferredCapabilities": [c.value for c in suggestion.inferred_capabilities],
+        "reasoning": suggestion.reasoning,
+    }
+
+
+def handle_estimate_cost(
+    provider: str,
+    model: str,
+    input_tokens: int,
+    output_tokens: int,
+    requests_per_day: int | None = None,
+    use_batch: bool = False,
+    use_cache: bool = False,
+) -> dict[str, Any]:
+    """Estimate cost for an LLM API call."""
+    try:
+        prov = Provider(provider)
+    except ValueError:
+        return {
+            "error": f"Unknown provider '{provider}'. Valid: {[p.value for p in get_all_providers()]}"
+        }
+
+    pricing = get_model(prov, model)
+    if not pricing:
+        return {"error": f"Unknown model '{model}' for provider '{provider}'."}
+
+    cost = calculate_cost(
+        pricing,
+        input_tokens,
+        output_tokens,
+        cached_input_tokens=input_tokens if use_cache else 0,
+        use_batch=use_batch,
+    )
+    monthly = cost * requests_per_day * 30 if requests_per_day else None
+
+    return {
+        "provider": provider,
+        "model": pricing.id,
+        "modelName": pricing.name,
+        "inputTokens": input_tokens,
+        "outputTokens": output_tokens,
+        "costPerCall": cost,
+        "monthlyCost": monthly,
+        "requestsPerDay": requests_per_day,
+        "pricing": {
+            "inputCostPerMillion": pricing.input_cost_per_million,
+            "outputCostPerMillion": pricing.output_cost_per_million,
+        },
+    }
+
+
+def handle_audit(directory: str, volume: int = 1000) -> dict[str, Any]:
+    """Scan directory and return cost estimates with recommendations."""
+    from tokenable.config import load_config
+    from tokenable.estimator import build_estimate_rows
+    from tokenable.recommender import detect_smart_alternatives
+    from tokenable.scanner import scan_directory
+
+    config = load_config()
+    results = scan_directory(directory, config.ignore)
+    rows, unknown_models = build_estimate_rows(results, config, volume, False, None, None)
+    total = sum(r.monthly_cost for r in rows)
+
+    # Get recommendations
+    findings = detect_smart_alternatives(results, volume)
+    recommendations = [
+        {
+            "file": f.file,
+            "line": f.line,
+            "currentProvider": f.current_provider.value,
+            "currentModel": f.current_model,
+            "suggestedProvider": f.suggested_provider.value,
+            "suggestedModel": f.suggested_model,
+            "reasoning": f.reasoning,
+            "savingsPercent": round((f.monthly_savings / f.current_monthly_cost) * 100)
+            if f.current_monthly_cost > 0
+            else 0,
+        }
+        for f in findings
+    ]
+
+    return {
+        "directory": directory,
+        "volume": volume,
+        "totalMonthlyCost": total,
+        "callSites": len(rows),
+        "unknownModels": list(unknown_models),
+        "recommendations": recommendations,
+    }
+
+
+def handle_apply_recommendations(
+    directory: str,
+    volume: int = 1000,
+    dry_run: bool = False,
+    recommendations: list[dict] | None = None,
+) -> dict[str, Any]:
+    """Apply model swap recommendations to source files."""
+    from tokenable.fixer import apply_recommendations
+    from tokenable.models import ModelSwap
+
+    if recommendations:
+        swaps = [
+            ModelSwap(
+                file=r["file"],
+                line=r["line"],
+                current_model=r["currentModel"],
+                suggested_model=r["suggestedModel"],
+            )
+            for r in recommendations
+        ]
+    else:
+        # Auto-detect from audit
+        audit_result = handle_audit(directory, volume)
+        swaps = [
+            ModelSwap(
+                file=r["file"],
+                line=r["line"],
+                current_model=r["currentModel"],
+                suggested_model=r["suggestedModel"],
+            )
+            for r in audit_result.get("recommendations", [])
+        ]
+
+    result = apply_recommendations(swaps, directory, dry_run)
+    return result.model_dump()
+
+
+# ── MCP Server (stdio JSON-RPC) ─────────────────────────────────────
+
+TOOLS = {
+    "suggest_model": {
+        "description": "Suggest the cheapest LLM model capable of handling a given task.",
+        "handler": handle_suggest_model,
+    },
+    "estimate_cost": {
+        "description": "Estimate the cost of an LLM API call.",
+        "handler": handle_estimate_cost,
+    },
+    "audit": {
+        "description": "Scan a directory for LLM API calls and suggest cheaper models.",
+        "handler": handle_audit,
+    },
+    "apply_recommendations": {
+        "description": "Apply model swap recommendations to source files.",
+        "handler": handle_apply_recommendations,
+    },
+}
+
+
+def run_server() -> None:
+    """Run MCP server on stdio (JSON-RPC)."""
+    sys.stderr.write("TokEnable MCP server running on stdio\n")
+
+    for line in sys.stdin:
+        line = line.strip()
+        if not line:
+            continue
+        try:
+            request = json.loads(line)
+            method = request.get("method", "")
+            params = request.get("params", {})
+            req_id = request.get("id")
+
+            if method == "tools/list":
+                result = [
+                    {"name": name, "description": info["description"]}
+                    for name, info in TOOLS.items()
+                ]
+            elif method == "tools/call":
+                tool_name = params.get("name", "")
+                tool_args = params.get("arguments", {})
+                if tool_name in TOOLS:
+                    result = TOOLS[tool_name]["handler"](**tool_args)
+                else:
+                    result = {"error": f"Unknown tool: {tool_name}"}
+            else:
+                result = {"error": f"Unknown method: {method}"}
+
+            response = {"jsonrpc": "2.0", "id": req_id, "result": result}
+        except Exception as e:
+            response = {"jsonrpc": "2.0", "id": None, "error": {"message": str(e)}}
+
+        sys.stdout.write(json.dumps(response) + "\n")
+        sys.stdout.flush()

tokenable/models/__init__.py

@@ -0,0 +1,308 @@
+"""Pydantic data models for TokEnable."""
+
+from __future__ import annotations
+
+from enum import Enum, StrEnum
+from typing import Literal
+
+from pydantic import BaseModel, Field
+
+
+class Provider(StrEnum):
+    """Supported LLM providers."""
+
+    ANTHROPIC = "anthropic"
+    OPENAI = "openai"
+    GOOGLE = "google"
+    XAI = "xai"
+    PERPLEXITY = "perplexity"
+
+
+class ModelTier(StrEnum):
+    """Cost tier derived from output pricing."""
+
+    BUDGET = "budget"
+    MID = "mid"
+    PREMIUM = "premium"
+
+
+class ModelStatus(StrEnum):
+    """Model availability status."""
+
+    CURRENT = "current"
+    LEGACY = "legacy"
+    DEPRECATED = "deprecated"
+
+
+class Capability(StrEnum):
+    """Model capabilities."""
+
+    CODE = "code"
+    REASONING = "reasoning"
+    GENERAL = "general"
+    CREATIVE = "creative"
+    VISION = "vision"
+    SEARCH = "search"
+    AUDIO = "audio"
+
+
+class ModelPricing(BaseModel):
+    """Pricing and metadata for a single model."""
+
+    id: str
+    name: str
+    provider: Provider
+    aliases: list[str] = Field(default_factory=list)
+    status: ModelStatus = ModelStatus.CURRENT
+    input_cost_per_million: float
+    output_cost_per_million: float
+    cache_read_input_cost_per_million: float | None = None
+    cache_write_input_cost_per_million: float | None = None
+    batch_input_cost_per_million: float | None = None
+    batch_output_cost_per_million: float | None = None
+    fast_input_cost_per_million: float | None = None
+    fast_output_cost_per_million: float | None = None
+    input_cost_above_200k_per_million: float | None = None
+    output_cost_above_200k_per_million: float | None = None
+    context_window: int
+    max_output_tokens: int
+    supports_vision: bool = False
+    supports_tools: bool = False
+    supports_prompt_caching: bool = False
+    supports_reasoning: bool = False
+    supports_computer_use: bool = False
+    tier: ModelTier = ModelTier.MID
+    capabilities: list[Capability] = Field(default_factory=list)
+    knowledge_cutoff: str | None = None
+
+
+class ScanResult(BaseModel):
+    """A single detected LLM API call site."""
+
+    file_path: str
+    line_number: int
+    provider: Provider
+    model: str | None = None
+    system_prompt: str | None = None
+    user_prompt: str | None = None
+    max_output_tokens: int | None = None
+    is_dynamic: bool = False
+    framework: str = ""
+
+
+class TokenSource(StrEnum):
+    """How token counts were determined."""
+
+    CODE = "code"
+    MODEL_LIMIT = "model_limit"
+    TYPICAL = "typical"
+    PRODUCTION = "production"
+    CALIBRATED = "calibrated"
+
+
+class EstimateRow(BaseModel):
+    """Cost estimate for a single call site."""
+
+    file: str
+    line: int
+    provider: str
+    model: str
+    input_tokens: int
+    input_token_source: TokenSource
+    output_tokens: int
+    output_token_source: TokenSource
+    cost_per_call: float
+    monthly_cost: float
+    system_prompt: str | None = None
+    user_prompt: str | None = None
+
+
+class EstimateSummary(BaseModel):
+    """Summary of all cost estimates."""
+
+    rows: list[EstimateRow]
+    total_monthly_cost: float
+    volume: int
+
+
+class ModelCalibration(BaseModel):
+    """Calibration factors for a single model."""
+
+    input_ratio: float = Field(ge=0.001, le=100)
+    output_ratio: float = Field(ge=0.001, le=100)
+    sample_size: int = Field(ge=0)
+    confidence: Literal["low", "medium", "high"]
+    actual_avg_input: float = Field(ge=0)
+    actual_avg_output: float = Field(ge=0)
+    estimated_avg_input: float = Field(ge=0)
+    estimated_avg_output: float = Field(ge=0)
+
+
+class CalibrationData(BaseModel):
+    """Stored calibration data."""
+
+    version: Literal[1] = 1
+    calibrated_at: str
+    models: dict[str, ModelCalibration] = Field(default_factory=dict)
+
+
+class ModelStats(BaseModel):
+    """Production usage stats for a model."""
+
+    provider: str
+    model: str
+    request_count: int
+    avg_input_tokens: float
+    avg_output_tokens: float
+    p50_input_tokens: float
+    p50_output_tokens: float
+    avg_cost_per_request: float = 0.0
+    total_cost: float = 0.0
+
+
+class ProviderUsageRecord(BaseModel):
+    """Usage data for a single model from a provider API."""
+
+    model: str
+    request_count: int
+    total_input_tokens: int
+    total_output_tokens: int
+    avg_input_tokens: int
+    avg_output_tokens: int
+
+
+class ProviderUsageResult(BaseModel):
+    """Usage data from a provider API."""
+
+    provider: Provider
+    records: list[ProviderUsageRecord]
+    period_start: str
+    period_end: str
+
+
+TelemetryBackend = Literal["otlp", "grafana-tempo", "inferwise-cloud"]
+
+
+class TelemetryConfig(BaseModel):
+    """Telemetry backend configuration."""
+
+    backend: TelemetryBackend
+    endpoint: str
+    headers: dict[str, str] = Field(default_factory=dict)
+    api_key: str | None = None
+
+
+class BudgetConfig(BaseModel):
+    """Budget thresholds for enforcement."""
+
+    warn: float | None = None
+    block: float | None = None
+    require_approval: float | None = None
+    approvers: list[str] = Field(default_factory=list)
+    max_monthly_cost: float | None = None
+
+
+class OverrideConfig(BaseModel):
+    """Volume override for a file pattern."""
+
+    pattern: str
+    volume: int | None = None
+
+
+class TokEnableConfig(BaseModel):
+    """Main configuration file schema."""
+
+    default_volume: int | None = None
+    ignore: list[str] = Field(default_factory=list)
+    overrides: list[OverrideConfig] = Field(default_factory=list)
+    budgets: BudgetConfig | None = None
+    telemetry: TelemetryConfig | None = None
+    # Deprecated fields
+    api_url: str | None = None
+    api_key: str | None = None
+
+
+class ModelSwap(BaseModel):
+    """A model swap recommendation."""
+
+    file: str
+    line: int
+    current_model: str
+    suggested_model: str
+    monthly_savings: float = 0.0
+
+
+class AppliedSwap(BaseModel):
+    """A successfully applied model swap."""
+
+    file: str
+    line: int
+    from_model: str
+    to_model: str
+
+
+class SkippedSwap(BaseModel):
+    """A skipped model swap with reason."""
+
+    file: str
+    line: int
+    from_model: str
+    to_model: str
+    reason: str
+
+
+class ApplyResult(BaseModel):
+    """Result of applying model swaps."""
+
+    applied: list[AppliedSwap] = Field(default_factory=list)
+    skipped: list[SkippedSwap] = Field(default_factory=list)
+    total_applied: int = 0
+    total_skipped: int = 0
+    estimated_monthly_savings: float = 0.0
+
+
+class SmartAlternativeFinding(BaseModel):
+    """A smart model alternative recommendation."""
+
+    type: Literal["smart-alternative"] = "smart-alternative"
+    file: str
+    line: int
+    current_provider: Provider
+    current_model: str
+    current_monthly_cost: float
+    suggested_provider: Provider
+    suggested_model: str
+    suggested_monthly_cost: float
+    monthly_savings: float
+    required_capabilities: list[Capability]
+    confidence: Literal["high", "medium", "low"]
+    reasoning: str
+    quality_score: float | None = None
+    current_quality_score: float | None = None
+
+
+class CachingFinding(BaseModel):
+    """A prompt caching opportunity."""
+
+    type: Literal["caching"] = "caching"
+    system_prompt: str
+    locations: list[dict[str, str | int]]
+    provider: Provider
+    model: str
+    current_monthly_cost: float
+    monthly_savings: float
+
+
+class BatchFinding(BaseModel):
+    """A batch API opportunity."""
+
+    type: Literal["batch"] = "batch"
+    file: str
+    model: str
+    provider: Provider
+    call_count: int
+    current_monthly_cost: float
+    monthly_savings: float
+
+
+AuditFinding = SmartAlternativeFinding | CachingFinding | BatchFinding

tokenable/pricing_sync.py

@@ -0,0 +1,145 @@
+"""Runtime pricing sync from LiteLLM's community database."""
+
+from __future__ import annotations
+
+import json
+from datetime import date
+from importlib.resources import files
+from pathlib import Path
+from typing import Any
+
+import httpx
+
+LITELLM_URL = "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"
+
+# Models we track per provider — LiteLLM key → our model ID
+TRACKED_MODELS: dict[str, dict[str, str]] = {
+    "anthropic": {
+        "claude-opus-4-6": "claude-opus-4-6",
+        "claude-sonnet-4-6": "claude-sonnet-4-6",
+        "claude-haiku-4-5-20251001": "claude-haiku-4-5-20251001",
+        "claude-3-5-haiku-20241022": "claude-3-5-haiku-20241022",
+    },
+    "openai": {
+        "gpt-4o": "gpt-4o",
+        "gpt-4o-mini": "gpt-4o-mini",
+        "gpt-4.1": "gpt-4.1",
+        "gpt-4.1-mini": "gpt-4.1-mini",
+        "gpt-4.1-nano": "gpt-4.1-nano",
+        "o3": "o3",
+        "o3-mini": "o3-mini",
+        "o4-mini": "o4-mini",
+    },
+    "google": {
+        "gemini/gemini-2.5-pro": "gemini-2.5-pro",
+        "gemini/gemini-2.5-flash": "gemini-2.5-flash",
+        "gemini/gemini-2.0-flash": "gemini-2.0-flash",
+        "gemini/gemini-2.0-flash-lite": "gemini-2.0-flash-lite",
+    },
+    "xai": {
+        "xai/grok-3": "grok-3",
+        "xai/grok-3-mini": "grok-3-mini",
+    },
+    "perplexity": {
+        "perplexity/sonar-pro": "sonar-pro",
+        "perplexity/sonar": "sonar",
+    },
+}
+
+
+def _get_providers_dir() -> Path:
+    """Get the path to the installed providers/data directory."""
+    data_dir = files("tokenable.providers") / "data"
+    # importlib.resources may return a Traversable; get the actual path
+    return Path(str(data_dir))
+
+
+def _find_entry(prices: dict, key: str) -> dict | None:
+    """Find a model in LiteLLM data."""
+    if key in prices:
+        return prices[key]
+    for full_key in prices:
+        if full_key.endswith(f"/{key}") or full_key == key:
+            return prices[full_key]
+    return None
+
+
+def sync_pricing(
+    provider: str | None = None,
+    dry_run: bool = False,
+) -> dict[str, Any]:
+    """Fetch latest pricing from LiteLLM and update local JSON files.
+
+    Returns dict with 'total_changes' and 'changes' list.
+    """
+    resp = httpx.get(LITELLM_URL, timeout=30, follow_redirects=True)
+    resp.raise_for_status()
+    prices = resp.json()
+
+    providers_dir = _get_providers_dir()
+    providers_to_sync = [provider] if provider else list(TRACKED_MODELS.keys())
+    all_changes: list[dict[str, Any]] = []
+
+    for prov in providers_to_sync:
+        tracked = TRACKED_MODELS.get(prov, {})
+        file_path = providers_dir / f"{prov}.json"
+        if not file_path.exists():
+            continue
+
+        data = json.loads(file_path.read_text())
+        models = data.get("models", [])
+        file_modified = False
+
+        for litellm_key, our_id in tracked.items():
+            entry = _find_entry(prices, litellm_key)
+            if not entry:
+                continue
+
+            our_model = next((m for m in models if m["id"] == our_id), None)
+            if not our_model:
+                continue
+
+            # Convert per-token → per-million
+            new_input = round((entry.get("input_cost_per_token", 0) or 0) * 1_000_000, 2)
+            new_output = round((entry.get("output_cost_per_token", 0) or 0) * 1_000_000, 2)
+
+            if new_input == 0 and new_output == 0:
+                continue
+
+            old_input = our_model.get("input_cost_per_million", 0)
+            old_output = our_model.get("output_cost_per_million", 0)
+
+            if new_input != old_input:
+                all_changes.append({"model": f"{prov}/{our_id}", "field": "input/M", "old": old_input, "new": new_input})
+                our_model["input_cost_per_million"] = new_input
+                file_modified = True
+
+            if new_output != old_output:
+                all_changes.append({"model": f"{prov}/{our_id}", "field": "output/M", "old": old_output, "new": new_output})
+                our_model["output_cost_per_million"] = new_output
+                file_modified = True
+
+            # Update context/output limits
+            if entry.get("max_input_tokens") and entry["max_input_tokens"] != our_model.get("context_window"):
+                our_model["context_window"] = entry["max_input_tokens"]
+                file_modified = True
+
+            max_out = entry.get("max_output_tokens") or entry.get("max_tokens")
+            if max_out and max_out != our_model.get("max_output_tokens"):
+                our_model["max_output_tokens"] = max_out
+                file_modified = True
+
+            # Cache pricing
+            if entry.get("cache_read_input_token_cost"):
+                our_model["cache_read_input_cost_per_million"] = round(entry["cache_read_input_token_cost"] * 1_000_000, 2)
+                file_modified = True
+            if entry.get("cache_creation_input_token_cost"):
+                our_model["cache_write_input_cost_per_million"] = round(entry["cache_creation_input_token_cost"] * 1_000_000, 2)
+                file_modified = True
+
+        if file_modified and not dry_run:
+            data["last_updated"] = date.today().isoformat()
+            data["last_verified"] = date.today().isoformat()
+            file_path.write_text(json.dumps(data, indent=2) + "\n")
+
+    return {"total_changes": len(all_changes), "changes": all_changes}