Merge pull request #4 from MukundaKatta/codex/tokenwise-pricing-and-budgeting

feat: add versioned pricing and budget tracking

Author: MukundaKatta

README.md

@@ -103,6 +103,22 @@ tracker.track(
 print(f"Total spend: ${tracker.total_cost():.6f}")
 ```
 
+### Multi-step Budget Breakdown
+
+```python
+from tokenwise import BudgetTracker
+
+tracker = BudgetTracker()
+tracker.add_step("draft", request="Write a landing page headline", response="Fast AI workflows for teams.")
+tracker.add_step("review", request="Critique the headline", response="Shorten the second clause.")
+
+report = tracker.get_report(warning_threshold_usd=0.01)
+print(report.total_cost)
+print(report.pricing_version)
+for step in report.steps:
+    print(step.name, step.total_tokens, step.total_cost)
+```
+
 ### CLI
 
 ```bash
@@ -139,6 +155,18 @@ summary = batch.batch_summary(unique)
 print(f"Saved {summary['total_tokens_saved']} tokens across {summary['prompt_count']} prompts")
 ```
 
+## Pricing Data
+
+Model pricing now lives in a versioned package data file at `src/tokenwise/data/model_pricing.v1.json`.
+
+That gives TokenWise a safer update workflow:
+
+- pricing changes are separated from estimator logic
+- the catalog carries an explicit version
+- historical reports can point back to the pricing version used at the time
+
+To update pricing, edit the JSON catalog, keep the schema consistent, and run the test suite before publishing.
+
 ## Pricing Table
 
 | Model | Input (per 1K tokens) | Output (per 1K tokens) |

pyproject.toml

@@ -49,6 +49,9 @@ Issues = "https://github.com/MukundaKatta/TokenWise/issues"
 [tool.setuptools.packages.find]
 where = ["src"]
 
+[tool.setuptools.package-data]
+tokenwise = ["data/*.json"]
+
 [tool.pytest.ini_options]
 testpaths = ["tests"]
 pythonpath = ["src"]

src/tokenwise/__init__.py

@@ -4,18 +4,21 @@
 
 from tokenwise.core import (
     BatchOptimizer,
+    BudgetTracker,
     CostEstimator,
     TokenCounter,
     TokenOptimizer,
     UsageTracker,
 )
-from tokenwise.config import TokenWiseConfig
+from tokenwise.config import PRICING_VERSION, TokenWiseConfig
 
 __all__ = [
     "TokenCounter",
     "TokenOptimizer",
     "CostEstimator",
     "UsageTracker",
+    "BudgetTracker",
     "BatchOptimizer",
     "TokenWiseConfig",
+    "PRICING_VERSION",
 ]

src/tokenwise/config.py

@@ -2,30 +2,27 @@
 
 from __future__ import annotations
 
+import json
 import os
+from functools import lru_cache
+from importlib.resources import files
 from typing import Optional
 
 from pydantic import BaseModel, Field
 
 
-# Per-token pricing in USD (per 1K tokens)
+@lru_cache(maxsize=1)
+def load_pricing_catalog() -> dict:
+    """Load the versioned pricing catalog from package data."""
+    catalog_path = files("tokenwise").joinpath("data/model_pricing.v1.json")
+    return json.loads(catalog_path.read_text(encoding="utf-8"))
+
+
+PRICING_CATALOG = load_pricing_catalog()
+PRICING_VERSION = PRICING_CATALOG["version"]
 MODEL_PRICING: dict[str, dict[str, float]] = {
-    "gpt-4": {"input": 0.03, "output": 0.06},
-    "gpt-4-turbo": {"input": 0.01, "output": 0.03},
-    "gpt-4o": {"input": 0.005, "output": 0.015},
-    "gpt-3.5-turbo": {"input": 0.0005, "output": 0.0015},
-    "claude-3-opus": {"input": 0.015, "output": 0.075},
-    "claude-3-sonnet": {"input": 0.003, "output": 0.015},
-    "claude-3-haiku": {"input": 0.00025, "output": 0.00125},
-    "claude-3.5-sonnet": {"input": 0.003, "output": 0.015},
-    "claude-4-opus": {"input": 0.015, "output": 0.075},
-    "claude-4-sonnet": {"input": 0.003, "output": 0.015},
-    "gemini-1.5-pro": {"input": 0.00125, "output": 0.005},
-    "gemini-1.5-flash": {"input": 0.000075, "output": 0.0003},
-    "llama-3-70b": {"input": 0.00059, "output": 0.00079},
-    "llama-3-8b": {"input": 0.00005, "output": 0.00008},
-    "mistral-large": {"input": 0.004, "output": 0.012},
-    "mistral-small": {"input": 0.001, "output": 0.003},
+    model: {"input": details["input"], "output": details["output"]}
+    for model, details in PRICING_CATALOG["models"].items()
 }
 
 # Characters-per-token ratio heuristics by model family
@@ -38,24 +35,9 @@
     "default": 3.7,
 }
 
-# Default context window sizes
 MODEL_CONTEXT_WINDOWS: dict[str, int] = {
-    "gpt-4": 8192,
-    "gpt-4-turbo": 128000,
-    "gpt-4o": 128000,
-    "gpt-3.5-turbo": 16385,
-    "claude-3-opus": 200000,
-    "claude-3-sonnet": 200000,
-    "claude-3-haiku": 200000,
-    "claude-3.5-sonnet": 200000,
-    "claude-4-opus": 200000,
-    "claude-4-sonnet": 200000,
-    "gemini-1.5-pro": 1000000,
-    "gemini-1.5-flash": 1000000,
-    "llama-3-70b": 8192,
-    "llama-3-8b": 8192,
-    "mistral-large": 32000,
-    "mistral-small": 32000,
+    model: details["context_window"]
+    for model, details in PRICING_CATALOG["models"].items()
 }
 
 # Default budget settings
@@ -84,6 +66,7 @@ class TokenWiseConfig(BaseModel):
     monthly_budget_usd: float = Field(default=DEFAULT_BUDGET["monthly_limit_usd"])
     alert_threshold_pct: int = Field(default=DEFAULT_BUDGET["alert_threshold_pct"])
     custom_pricing: Optional[dict[str, dict[str, float]]] = None
+    pricing_version: str = Field(default=PRICING_VERSION)
 
     def get_pricing(self, model: str) -> dict[str, float]:
         """Return pricing dict for a model, checking custom overrides first."""
@@ -101,3 +84,11 @@ def get_tokenizer_ratio(self, model: str) -> float:
             if family in model.lower():
                 return ratio
         return TOKENIZER_RATIOS["default"]
+
+    def get_context_window(self, model: str) -> int:
+        """Return the context window for a model."""
+        if model in MODEL_CONTEXT_WINDOWS:
+            return MODEL_CONTEXT_WINDOWS[model]
+        raise ValueError(
+            f"Unknown model '{model}'. Available: {', '.join(MODEL_CONTEXT_WINDOWS.keys())}"
+        )

src/tokenwise/core.py

@@ -11,6 +11,7 @@
 from tokenwise.config import (
     MODEL_CONTEXT_WINDOWS,
     MODEL_PRICING,
+    PRICING_VERSION,
     TokenWiseConfig,
 )
 from tokenwise.utils import (
@@ -34,6 +35,34 @@ class UsageRecord(BaseModel):
     total_tokens: int
     model: str
     estimated_cost: float
+    pricing_version: str
+
+
+class BudgetStep(BaseModel):
+    """One step in a multi-step budget workflow."""
+
+    name: str
+    model: str
+    request_tokens: int
+    response_tokens: int
+    total_tokens: int
+    input_cost: float
+    output_cost: float
+    total_cost: float
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+
+class BudgetReport(BaseModel):
+    """Aggregate report for a multi-step budget workflow."""
+
+    pricing_version: str
+    total_steps: int
+    total_tokens: int
+    total_cost: float
+    warning_threshold_usd: float | None = None
+    warning_triggered: bool = False
+    by_model: dict[str, dict[str, float]]
+    steps: list[BudgetStep]
 
 
 class BudgetAlert(BaseModel):
@@ -253,6 +282,7 @@ def track(self, request: str, response: str, model: Optional[str] = None) -> Usa
             total_tokens=req_tokens + res_tokens,
             model=model,
             estimated_cost=round(input_cost + output_cost, 8),
+            pricing_version=self.config.pricing_version,
         )
         self._records.append(record)
         return record
@@ -300,6 +330,7 @@ def get_report(self) -> dict[str, Any]:
             "total_requests": len(self._records),
             "total_tokens": self.total_tokens(),
             "estimated_total_cost": self.total_cost(),
+            "pricing_version": self.config.pricing_version,
             "by_model": by_model,
         }
 
@@ -368,3 +399,70 @@ def deduplicate_prompts(self, prompts: list[str]) -> list[str]:
                 seen.add(normalized)
                 unique.append(p)
         return unique
+
+
+class BudgetTracker:
+    """Track token and cost breakdowns across multi-step tasks."""
+
+    def __init__(self, config: Optional[TokenWiseConfig] = None) -> None:
+        self.config = config or TokenWiseConfig()
+        self._counter = TokenCounter(self.config)
+        self._estimator = CostEstimator(self.config)
+        self._steps: list[BudgetStep] = []
+
+    def add_step(
+        self,
+        name: str,
+        request: str,
+        response: str = "",
+        model: Optional[str] = None,
+        metadata: Optional[dict[str, Any]] = None,
+    ) -> BudgetStep:
+        """Add one step to the budget report."""
+        model = model or self.config.default_model
+        request_tokens = self._counter.count(request, model)
+        response_tokens = self._counter.count(response, model) if response else 0
+        input_cost = self._estimator.estimate(request_tokens, model, "input")
+        output_cost = self._estimator.estimate(response_tokens, model, "output")
+        step = BudgetStep(
+            name=name,
+            model=model,
+            request_tokens=request_tokens,
+            response_tokens=response_tokens,
+            total_tokens=request_tokens + response_tokens,
+            input_cost=input_cost,
+            output_cost=output_cost,
+            total_cost=round(input_cost + output_cost, 8),
+            metadata=metadata or {},
+        )
+        self._steps.append(step)
+        return step
+
+    def get_report(self, warning_threshold_usd: float | None = None) -> BudgetReport:
+        """Return the full workflow budget report."""
+        by_model: dict[str, dict[str, float]] = {}
+        for step in self._steps:
+            if step.model not in by_model:
+                by_model[step.model] = {"steps": 0.0, "tokens": 0.0, "cost": 0.0}
+            by_model[step.model]["steps"] += 1
+            by_model[step.model]["tokens"] += step.total_tokens
+            by_model[step.model]["cost"] += step.total_cost
+
+        total_cost = round(sum(step.total_cost for step in self._steps), 8)
+        total_tokens = sum(step.total_tokens for step in self._steps)
+        return BudgetReport(
+            pricing_version=self.config.pricing_version,
+            total_steps=len(self._steps),
+            total_tokens=total_tokens,
+            total_cost=total_cost,
+            warning_threshold_usd=warning_threshold_usd,
+            warning_triggered=(
+                warning_threshold_usd is not None and total_cost >= warning_threshold_usd
+            ),
+            by_model=by_model,
+            steps=list(self._steps),
+        )
+
+    def reset(self) -> None:
+        """Clear tracked steps."""
+        self._steps.clear()

src/tokenwise/data/model_pricing.v1.json

@@ -0,0 +1,21 @@
+{
+  "version": "2026-04-20",
+  "models": {
+    "gpt-4": {"input": 0.03, "output": 0.06, "context_window": 8192},
+    "gpt-4-turbo": {"input": 0.01, "output": 0.03, "context_window": 128000},
+    "gpt-4o": {"input": 0.005, "output": 0.015, "context_window": 128000},
+    "gpt-3.5-turbo": {"input": 0.0005, "output": 0.0015, "context_window": 16385},
+    "claude-3-opus": {"input": 0.015, "output": 0.075, "context_window": 200000},
+    "claude-3-sonnet": {"input": 0.003, "output": 0.015, "context_window": 200000},
+    "claude-3-haiku": {"input": 0.00025, "output": 0.00125, "context_window": 200000},
+    "claude-3.5-sonnet": {"input": 0.003, "output": 0.015, "context_window": 200000},
+    "claude-4-opus": {"input": 0.015, "output": 0.075, "context_window": 200000},
+    "claude-4-sonnet": {"input": 0.003, "output": 0.015, "context_window": 200000},
+    "gemini-1.5-pro": {"input": 0.00125, "output": 0.005, "context_window": 1000000},
+    "gemini-1.5-flash": {"input": 0.000075, "output": 0.0003, "context_window": 1000000},
+    "llama-3-70b": {"input": 0.00059, "output": 0.00079, "context_window": 8192},
+    "llama-3-8b": {"input": 0.00005, "output": 0.00008, "context_window": 8192},
+    "mistral-large": {"input": 0.004, "output": 0.012, "context_window": 32000},
+    "mistral-small": {"input": 0.001, "output": 0.003, "context_window": 32000}
+  }
+}

tests/test_core.py

@@ -6,12 +6,13 @@
 
 from tokenwise.core import (
     BatchOptimizer,
+    BudgetTracker,
     CostEstimator,
     TokenCounter,
     TokenOptimizer,
     UsageTracker,
 )
-from tokenwise.config import TokenWiseConfig
+from tokenwise.config import PRICING_VERSION, TokenWiseConfig, load_pricing_catalog
 
 
 class TestTokenCounter:
@@ -95,6 +96,9 @@ def test_compare_models_sorted_by_cost(self) -> None:
         costs = [info["cost"] for info in result.values()]
         assert costs == sorted(costs)
 
+    def test_pricing_version_exposed_from_catalog(self) -> None:
+        assert PRICING_VERSION == load_pricing_catalog()["version"]
+
 
 class TestUsageTracker:
     """Tests for usage tracking."""
@@ -108,6 +112,7 @@ def test_track_single_request(self) -> None:
         assert record.request_tokens > 0
         assert record.response_tokens > 0
         assert record.total_tokens == record.request_tokens + record.response_tokens
+        assert record.pricing_version == PRICING_VERSION
 
     def test_report_aggregates_correctly(self) -> None:
         tracker = UsageTracker()
@@ -118,6 +123,7 @@ def test_report_aggregates_correctly(self) -> None:
         assert report["total_requests"] == 2
         assert report["total_tokens"] > 0
         assert report["estimated_total_cost"] > 0
+        assert report["pricing_version"] == PRICING_VERSION
 
     def test_reset_clears_log(self) -> None:
         tracker = UsageTracker()
@@ -197,3 +203,21 @@ def test_deduplicate(self) -> None:
         prompts = ["Hello", "hello", "World", "Hello"]
         unique = batch.deduplicate_prompts(prompts)
         assert len(unique) == 2
+
+
+class TestBudgetTracker:
+    """Tests for multi-step budget reporting."""
+
+    def test_budget_report_breaks_costs_down_by_step(self) -> None:
+        tracker = BudgetTracker()
+        tracker.add_step("draft", request="Write a summary", response="Here is a draft", model="gpt-4o")
+        tracker.add_step("review", request="Critique the draft", response="Needs more detail", model="gpt-4o")
+
+        report = tracker.get_report(warning_threshold_usd=0.00000001)
+
+        assert report.total_steps == 2
+        assert report.total_tokens > 0
+        assert report.total_cost > 0
+        assert report.warning_triggered is True
+        assert report.pricing_version == PRICING_VERSION
+        assert len(report.steps) == 2