tokenable 1.0.0__py3-none-any.whl

tokenable/config/__init__.py

@@ -0,0 +1,155 @@
+"""Configuration loading and volume resolution for TokEnable."""
+
+from __future__ import annotations
+
+import json
+import os
+from fnmatch import fnmatch
+from pathlib import Path
+
+from tokenable.models import (
+    BudgetConfig,
+    OverrideConfig,
+    TelemetryConfig,
+    TokEnableConfig,
+)
+
+CONFIG_FILENAME = "tokenable.config.json"
+
+_VOLUME_NAMES: dict[str, int] = {
+    "low": 100,
+    "medium": 1000,
+    "high": 10000,
+}
+
+
+def parseVolume(value: str) -> int | None:
+    """Parse a volume string (name or integer)."""
+    lower = value.strip().lower()
+    if lower in _VOLUME_NAMES:
+        return _VOLUME_NAMES[lower]
+    try:
+        return int(value)
+    except ValueError:
+        return None
+
+
+def getEnvVolume() -> int | None:
+    """Read volume from TOKENABLE_VOLUME env var."""
+    val = os.environ.get("TOKENABLE_VOLUME")
+    if val:
+        return parseVolume(val)
+    return None
+
+
+def _find_config_file(start: Path | None = None) -> Path | None:
+    """Walk up from start directory looking for config file."""
+    current = start or Path.cwd()
+    current = current.resolve()
+    while True:
+        candidate = current / CONFIG_FILENAME
+        if candidate.is_file():
+            return candidate
+        parent = current.parent
+        if parent == current:
+            break
+        current = parent
+    return None
+
+
+def load_config(path: str | None = None) -> TokEnableConfig | None:
+    """Load config from explicit path, env var, or by walking up from CWD."""
+    config_path: Path | None = None
+
+    if path:
+        config_path = Path(path)
+    else:
+        env_path = os.environ.get("TOKENABLE_CONFIG")
+        config_path = Path(env_path) if env_path else _find_config_file()
+
+    if config_path is None or not config_path.is_file():
+        return None
+
+    raw = json.loads(config_path.read_text("utf-8"))
+    return _parse_config(raw)
+
+
+def _parse_config(raw: dict) -> TokEnableConfig:
+    """Parse raw JSON dict into TokEnableConfig, handling deprecated fields."""
+    overrides = [
+        OverrideConfig(pattern=o["pattern"], volume=o.get("volume"))
+        for o in raw.get("overrides", [])
+    ]
+
+    budgets: BudgetConfig | None = None
+    if "budgets" in raw:
+        b = raw["budgets"]
+        budgets = BudgetConfig(
+            warn=b.get("warn"),
+            block=b.get("block"),
+            require_approval=b.get("requireApproval"),
+            approvers=b.get("approvers", []),
+            max_monthly_cost=b.get("maxMonthlyCost"),
+        )
+
+    telemetry: TelemetryConfig | None = None
+    if "telemetry" in raw:
+        t = raw["telemetry"]
+        telemetry = TelemetryConfig(
+            backend=t["backend"],
+            endpoint=t["endpoint"],
+            headers=t.get("headers", {}),
+            api_key=t.get("apiKey"),
+        )
+    elif raw.get("apiUrl") or raw.get("apiKey"):
+        # Deprecated: map apiUrl/apiKey to inferwise-cloud backend
+        telemetry = TelemetryConfig(
+            backend="inferwise-cloud",
+            endpoint=raw.get("apiUrl", ""),
+            api_key=raw.get("apiKey"),
+        )
+
+    return TokEnableConfig(
+        default_volume=raw.get("defaultVolume"),
+        ignore=raw.get("ignore", []),
+        overrides=overrides,
+        budgets=budgets,
+        telemetry=telemetry,
+        api_url=raw.get("apiUrl"),
+        api_key=raw.get("apiKey"),
+    )
+
+
+def resolveVolume(
+    config: TokEnableConfig | None,
+    file_path: str,
+    cli_volume: int | None,
+    cli_volume_explicit: bool,
+) -> int:
+    """Resolve effective volume for a file path.
+
+    Priority: CLI explicit > override match > env > config default > 1000.
+    """
+    if cli_volume_explicit and cli_volume is not None:
+        return cli_volume
+
+    # Check overrides
+    if config and config.overrides:
+        for override in config.overrides:
+            if fnmatch(file_path, override.pattern) and override.volume is not None:
+                return override.volume
+
+    # Env var
+    env_vol = getEnvVolume()
+    if env_vol is not None:
+        return env_vol
+
+    # Config default
+    if config and config.default_volume is not None:
+        return config.default_volume
+
+    # CLI non-explicit (fallback)
+    if cli_volume is not None:
+        return cli_volume
+
+    return 1000

tokenable/enforcer/__init__.py

@@ -0,0 +1,124 @@
+"""Enforcement — pre-commit hooks, CI gate generation, budget policy."""
+
+from __future__ import annotations
+
+import os
+import stat
+from pathlib import Path
+
+PRE_COMMIT_HOOK = """#!/bin/sh
+# TokEnable cost check — blocks commits that exceed budget thresholds
+# Installed by: tokenable init
+# Config: tokenable.config.json
+
+tokenable check .
+"""
+
+PRE_PUSH_HOOK = """#!/bin/sh
+# TokEnable cost diff — compares HEAD against main before push
+# Installed by: tokenable init
+# Config: tokenable.config.json (budgets.block threshold)
+
+tokenable diff --format table 2>&1
+"""
+
+GITHUB_ACTION_WORKFLOW = """name: TokEnable Cost Check
+on:
+  pull_request:
+    branches: [main]
+
+jobs:
+  cost-check:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+      - uses: actions/setup-python@v5
+        with:
+          python-version: '3.11'
+      - run: pip install tokenable
+      - run: tokenable diff --base origin/main --head HEAD --format markdown
+"""
+
+CONFIG_TEMPLATE = {
+    "defaultVolume": 1000,
+    "ignore": ["node_modules", "dist", "build", "test", "__tests__", "*.test.ts", "*.spec.ts"],
+    "budgets": {"warn": 2000, "block": 50000},
+}
+
+
+def find_git_root(start_dir: str) -> str | None:
+    """Walk up from start_dir to find .git directory."""
+    current = Path(start_dir).resolve()
+    while True:
+        if (current / ".git").exists():
+            return str(current)
+        parent = current.parent
+        if parent == current:
+            return None
+        current = parent
+
+
+def detect_hook_manager(git_root: str) -> str:
+    """Detect hook manager: husky, lefthook, or git."""
+    root = Path(git_root)
+    if (root / ".husky").exists():
+        return "husky"
+    if (root / "lefthook.yml").exists() or (root / ".lefthook.yml").exists():
+        return "lefthook"
+    return "git"
+
+
+def install_hook(hooks_dir: str, hook_name: str, content: str) -> str:
+    """Install a git hook. Returns 'created', 'exists', or 'updated'."""
+    hook_path = Path(hooks_dir) / hook_name
+
+    if hook_path.exists():
+        existing = hook_path.read_text()
+        if "tokenable" in existing.lower():
+            return "exists"
+        hook_path.write_text(existing + "\n" + content)
+        hook_path.chmod(hook_path.stat().st_mode | stat.S_IEXEC)
+        return "updated"
+
+    hook_path.parent.mkdir(parents=True, exist_ok=True)
+    hook_path.write_text(content)
+    hook_path.chmod(hook_path.stat().st_mode | stat.S_IEXEC)
+    return "created"
+
+
+def setup_hooks(git_root: str, hook_type: str = "pre-commit") -> str:
+    """Install git hooks. Returns status message."""
+    manager = detect_hook_manager(git_root)
+    content = PRE_PUSH_HOOK if hook_type == "pre-push" else PRE_COMMIT_HOOK
+
+    if manager == "husky":
+        hooks_dir = os.path.join(git_root, ".husky")
+    elif manager == "lefthook":
+        return "lefthook detected — add tokenable check to your lefthook.yml manually"
+    else:
+        hooks_dir = os.path.join(git_root, ".git", "hooks")
+
+    result = install_hook(hooks_dir, hook_type, content)
+    return f"{result} {hook_type} hook ({manager})"
+
+
+def create_config(project_dir: str) -> str | None:
+    """Create tokenable.config.json if it doesn't exist. Returns path or None."""
+    import json
+
+    config_path = Path(project_dir) / "tokenable.config.json"
+    if config_path.exists():
+        return None
+    config_path.write_text(json.dumps(CONFIG_TEMPLATE, indent=2) + "\n")
+    return str(config_path)
+
+
+def generate_ci_workflow(project_dir: str) -> str:
+    """Generate GitHub Actions workflow file."""
+    workflow_dir = Path(project_dir) / ".github" / "workflows"
+    workflow_dir.mkdir(parents=True, exist_ok=True)
+    workflow_path = workflow_dir / "tokenable-cost-check.yml"
+    workflow_path.write_text(GITHUB_ACTION_WORKFLOW)
+    return str(workflow_path)

tokenable/estimator/__init__.py

@@ -0,0 +1,192 @@
+"""Cost estimation engine for TokEnable."""
+
+from __future__ import annotations
+
+import tiktoken
+
+from tokenable.models import (
+    CalibrationData,
+    EstimateRow,
+    ModelPricing,
+    ModelStats,
+    ScanResult,
+    TokEnableConfig,
+    TokenSource,
+)
+from tokenable.providers import calculate_cost as _calc_cost
+from tokenable.providers import get_model, get_provider_models
+
+# ── Token counting ───────────────────────────────────────────────────
+
+_encoding: tiktoken.Encoding | None = None
+
+
+def _get_encoding() -> tiktoken.Encoding:
+    global _encoding
+    if _encoding is None:
+        _encoding = tiktoken.get_encoding("cl100k_base")
+    return _encoding
+
+
+def count_tokens(provider: str, model: str, text: str) -> int:
+    """Count tokens in text using cl100k_base encoding."""
+    return len(_get_encoding().encode(text))
+
+
+def count_message_tokens(provider: str, model: str, system: str | None, user: str | None) -> int:
+    """Count tokens for a system+user message pair."""
+    total = 0
+    enc = _get_encoding()
+    if system:
+        total += len(enc.encode(system))
+    if user:
+        total += len(enc.encode(user))
+    return total
+
+
+# ── Typical token heuristics ─────────────────────────────────────────
+
+
+def typicalInputTokens(pricing: ModelPricing) -> int:
+    """Estimate typical input tokens for a model."""
+    if pricing.context_window < 16384:
+        return int(pricing.context_window * 0.25)
+    return 4096
+
+
+# Snake-case aliases
+typical_input_tokens = typicalInputTokens
+
+
+def typicalOutputTokens(pricing: ModelPricing) -> int:
+    """Estimate typical output tokens for a model."""
+    raw = int(pricing.max_output_tokens * 0.05)
+    return max(512, min(raw, 4096))
+
+
+# Snake-case aliases
+typical_output_tokens = typicalOutputTokens
+
+
+# ── Estimate builder ─────────────────────────────────────────────────
+
+
+def calculateCost(model: ModelPricing, input_tokens: int, output_tokens: int) -> float:
+    """Wrapper around providers.calculateCost."""
+    return _calc_cost(model, input_tokens, output_tokens)
+
+
+def buildEstimateRows(
+    results: list[ScanResult],
+    config: TokEnableConfig | None,
+    cli_volume: int | None,
+    cli_volume_explicit: bool,
+    stats_map: dict[str, ModelStats] | None = None,
+    calibration: CalibrationData | None = None,
+) -> tuple[list[EstimateRow], list[str]]:
+    """Build estimate rows from scan results.
+
+    Returns (rows, unknown_models).
+    """
+    from tokenable.config import resolveVolume
+
+    rows: list[EstimateRow] = []
+    unknown_models: list[str] = []
+
+    for result in results:
+        model_id = result.model
+        if not model_id:
+            unknown_models.append(f"{result.provider.value}/?")
+            continue
+
+        pricing = get_model(result.provider, model_id)
+        if pricing is None:
+            unknown_models.append(f"{result.provider.value}/{model_id}")
+            continue
+
+        volume = resolveVolume(config, result.file_path, cli_volume, cli_volume_explicit)
+
+        # Determine input tokens
+        input_tokens, input_source = _resolve_input_tokens(result, pricing, stats_map)
+        # Determine output tokens
+        output_tokens, output_source = _resolve_output_tokens(result, pricing, stats_map)
+
+        # Apply calibration for typical/model_limit sources
+        if calibration:
+            key = f"{result.provider.value}/{model_id}"
+            cal = calibration.models.get(key)
+            if cal:
+                if input_source in (TokenSource.TYPICAL, TokenSource.MODEL_LIMIT):
+                    input_tokens = int(input_tokens * cal.input_ratio)
+                    input_source = TokenSource.CALIBRATED
+                if output_source in (TokenSource.TYPICAL, TokenSource.MODEL_LIMIT):
+                    output_tokens = int(output_tokens * cal.output_ratio)
+                    output_source = TokenSource.CALIBRATED
+
+        cost_per_call = calculateCost(pricing, input_tokens, output_tokens)
+        monthly_cost = cost_per_call * volume
+
+        rows.append(
+            EstimateRow(
+                file=result.file_path,
+                line=result.line_number,
+                provider=result.provider.value,
+                model=pricing.id,
+                input_tokens=input_tokens,
+                input_token_source=input_source,
+                output_tokens=output_tokens,
+                output_token_source=output_source,
+                cost_per_call=cost_per_call,
+                monthly_cost=monthly_cost,
+                system_prompt=result.system_prompt,
+                user_prompt=result.user_prompt,
+            )
+        )
+
+    return rows, unknown_models
+
+
+def _resolve_input_tokens(
+    result: ScanResult,
+    pricing: ModelPricing,
+    stats_map: dict[str, ModelStats] | None,
+) -> tuple[int, TokenSource]:
+    """Determine input token count and source."""
+    # Production stats
+    if stats_map:
+        key = f"{result.provider.value}/{pricing.id}"
+        stats = stats_map.get(key)
+        if stats:
+            return int(stats.avg_input_tokens), TokenSource.PRODUCTION
+
+    # Count from code
+    if result.system_prompt or result.user_prompt:
+        tokens = count_message_tokens(
+            result.provider.value, pricing.id, result.system_prompt, result.user_prompt
+        )
+        if tokens > 0:
+            return tokens, TokenSource.CODE
+
+    # Typical heuristic
+    return typicalInputTokens(pricing), TokenSource.TYPICAL
+
+
+def _resolve_output_tokens(
+    result: ScanResult,
+    pricing: ModelPricing,
+    stats_map: dict[str, ModelStats] | None,
+) -> tuple[int, TokenSource]:
+    """Determine output token count and source."""
+    # Production stats
+    if stats_map:
+        key = f"{result.provider.value}/{pricing.id}"
+        stats = stats_map.get(key)
+        if stats:
+            return int(stats.avg_output_tokens), TokenSource.PRODUCTION
+
+    # Explicit max_output_tokens in code
+    if result.max_output_tokens is not None:
+        return result.max_output_tokens, TokenSource.MODEL_LIMIT
+
+    # Typical heuristic
+    return typicalOutputTokens(pricing), TokenSource.TYPICAL

tokenable/fixer/__init__.py

@@ -0,0 +1,101 @@
+"""Fixer module — applies model swap recommendations to source files."""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from pathlib import Path
+
+from tokenable.models import AppliedSwap, ApplyResult, ModelSwap, SkippedSwap
+
+SEARCH_WINDOW = 10
+
+
+def apply_model_swap(
+    file_content: str,
+    line_number: int,
+    current_model: str,
+    suggested_model: str,
+) -> dict[str, object] | None:
+    """Replace current_model with suggested_model near line_number, preserving quote style."""
+    lines = file_content.splitlines()
+    start = max(0, line_number - SEARCH_WINDOW)
+    end = min(len(lines), line_number + SEARCH_WINDOW)
+
+    for i in range(start, end):
+        line = lines[i]
+        for quote in ('"', "'", "`"):
+            token = f"{quote}{current_model}{quote}"
+            if token in line:
+                lines[i] = line.replace(token, f"{quote}{suggested_model}{quote}", 1)
+                return {"content": "\n".join(lines), "actual_line": i + 1}
+    return None
+
+
+def apply_recommendations(
+    swaps: list[ModelSwap],
+    base_path: str,
+    dry_run: bool = False,
+) -> ApplyResult:
+    """Apply swaps grouped by file, bottom-to-top to preserve line numbers."""
+    applied: list[AppliedSwap] = []
+    skipped: list[SkippedSwap] = []
+    total_savings = 0.0
+
+    by_file: dict[str, list[ModelSwap]] = defaultdict(list)
+    for s in swaps:
+        by_file[s.file].append(s)
+
+    for file_rel, file_swaps in by_file.items():
+        path = Path(base_path) / file_rel
+        if not path.exists():
+            for s in file_swaps:
+                skipped.append(
+                    SkippedSwap(
+                        file=s.file,
+                        line=s.line,
+                        from_model=s.current_model,
+                        to_model=s.suggested_model,
+                        reason="file not found",
+                    )
+                )
+            continue
+
+        content = path.read_text()
+        modified = False
+        sorted_swaps = sorted(file_swaps, key=lambda s: s.line, reverse=True)
+
+        for s in sorted_swaps:
+            result = apply_model_swap(content, s.line, s.current_model, s.suggested_model)
+            if result is None:
+                skipped.append(
+                    SkippedSwap(
+                        file=s.file,
+                        line=s.line,
+                        from_model=s.current_model,
+                        to_model=s.suggested_model,
+                        reason="model string not found as literal",
+                    )
+                )
+            else:
+                content = result["content"]  # type: ignore[assignment]
+                modified = True
+                applied.append(
+                    AppliedSwap(
+                        file=s.file,
+                        line=s.line,
+                        from_model=s.current_model,
+                        to_model=s.suggested_model,
+                    )
+                )
+                total_savings += s.monthly_savings
+
+        if modified and not dry_run:
+            path.write_text(content)
+
+    return ApplyResult(
+        applied=applied,
+        skipped=skipped,
+        total_applied=len(applied),
+        total_skipped=len(skipped),
+        estimated_monthly_savings=total_savings,
+    )