compute-cfo 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,33 @@
1
+ # Python
2
+ __pycache__/
3
+ *.py[cod]
4
+ *.egg-info/
5
+ dist/
6
+ build/
7
+ .eggs/
8
+ *.egg
9
+ .venv/
10
+ venv/
11
+ .pytest_cache/
12
+ .mypy_cache/
13
+ *.jsonl
14
+
15
+ # TypeScript
16
+ node_modules/
17
+ dist/
18
+ *.js
19
+ *.d.ts
20
+ *.js.map
21
+ !jest.config.js
22
+ !tsconfig.json
23
+
24
+ # IDE
25
+ .vscode/
26
+ .idea/
27
+ *.swp
28
+ *.swo
29
+ .DS_Store
30
+
31
+ # Environment
32
+ .env
33
+ .env.local
@@ -0,0 +1,35 @@
1
+ Metadata-Version: 2.4
2
+ Name: compute-cfo
3
+ Version: 0.1.0
4
+ Summary: Cost tracking, attribution, and budget enforcement for AI inference APIs
5
+ Project-URL: Homepage, https://github.com/YanLukashin/compute-cfo
6
+ Project-URL: Documentation, https://github.com/YanLukashin/compute-cfo#readme
7
+ Project-URL: Repository, https://github.com/YanLukashin/compute-cfo
8
+ Project-URL: Issues, https://github.com/YanLukashin/compute-cfo/issues
9
+ Author-email: Compute CFO <hello@computecfo.com>
10
+ License-Expression: Apache-2.0
11
+ Keywords: ai,anthropic,budget,cost,inference,llm,openai,tracking
12
+ Classifier: Development Status :: 3 - Alpha
13
+ Classifier: Intended Audience :: Developers
14
+ Classifier: License :: OSI Approved :: Apache Software License
15
+ Classifier: Programming Language :: Python :: 3
16
+ Classifier: Programming Language :: Python :: 3.9
17
+ Classifier: Programming Language :: Python :: 3.10
18
+ Classifier: Programming Language :: Python :: 3.11
19
+ Classifier: Programming Language :: Python :: 3.12
20
+ Classifier: Programming Language :: Python :: 3.13
21
+ Classifier: Topic :: Software Development :: Libraries
22
+ Requires-Python: >=3.9
23
+ Provides-Extra: all
24
+ Requires-Dist: anthropic>=0.20.0; extra == 'all'
25
+ Requires-Dist: openai>=1.0.0; extra == 'all'
26
+ Provides-Extra: anthropic
27
+ Requires-Dist: anthropic>=0.20.0; extra == 'anthropic'
28
+ Provides-Extra: dev
29
+ Requires-Dist: pytest-asyncio>=0.21.0; extra == 'dev'
30
+ Requires-Dist: pytest>=7.0; extra == 'dev'
31
+ Provides-Extra: openai
32
+ Requires-Dist: openai>=1.0.0; extra == 'openai'
33
+ Description-Content-Type: text/markdown
34
+
35
+ See the main [README](../README.md).
@@ -0,0 +1 @@
1
+ See the main [README](../README.md).
@@ -0,0 +1,45 @@
1
+ [build-system]
2
+ requires = ["hatchling"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "compute-cfo"
7
+ version = "0.1.0"
8
+ description = "Cost tracking, attribution, and budget enforcement for AI inference APIs"
9
+ readme = "README.md"
10
+ license = "Apache-2.0"
11
+ requires-python = ">=3.9"
12
+ authors = [
13
+ { name = "Compute CFO", email = "hello@computecfo.com" },
14
+ ]
15
+ keywords = ["llm", "cost", "tracking", "openai", "anthropic", "budget", "inference", "ai"]
16
+ classifiers = [
17
+ "Development Status :: 3 - Alpha",
18
+ "Intended Audience :: Developers",
19
+ "License :: OSI Approved :: Apache Software License",
20
+ "Programming Language :: Python :: 3",
21
+ "Programming Language :: Python :: 3.9",
22
+ "Programming Language :: Python :: 3.10",
23
+ "Programming Language :: Python :: 3.11",
24
+ "Programming Language :: Python :: 3.12",
25
+ "Programming Language :: Python :: 3.13",
26
+ "Topic :: Software Development :: Libraries",
27
+ ]
28
+
29
+ [project.optional-dependencies]
30
+ openai = ["openai>=1.0.0"]
31
+ anthropic = ["anthropic>=0.20.0"]
32
+ all = ["openai>=1.0.0", "anthropic>=0.20.0"]
33
+ dev = ["pytest>=7.0", "pytest-asyncio>=0.21.0"]
34
+
35
+ [project.urls]
36
+ Homepage = "https://github.com/YanLukashin/compute-cfo"
37
+ Documentation = "https://github.com/YanLukashin/compute-cfo#readme"
38
+ Repository = "https://github.com/YanLukashin/compute-cfo"
39
+ Issues = "https://github.com/YanLukashin/compute-cfo/issues"
40
+
41
+ [tool.hatch.build.targets.wheel]
42
+ packages = ["src/compute_cfo"]
43
+
44
+ [tool.pytest.ini_options]
45
+ testpaths = ["tests"]
@@ -0,0 +1,23 @@
1
+ """compute-cfo: Cost tracking, attribution, and budget enforcement for AI inference APIs."""
2
+
3
+ from .budget import BudgetPolicy
4
+ from .exporters import console_exporter, jsonl_exporter, webhook_exporter
5
+ from .pricing import get_cost, get_price
6
+ from .tracker import CostTracker
7
+ from .types import BudgetExceededError, CostEvent
8
+ from .wrapper import wrap
9
+
10
+ __version__ = "0.1.0"
11
+
12
+ __all__ = [
13
+ "wrap",
14
+ "CostTracker",
15
+ "CostEvent",
16
+ "BudgetPolicy",
17
+ "BudgetExceededError",
18
+ "get_cost",
19
+ "get_price",
20
+ "console_exporter",
21
+ "jsonl_exporter",
22
+ "webhook_exporter",
23
+ ]
@@ -0,0 +1,91 @@
1
+ """Budget policy enforcement."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import warnings
6
+ from dataclasses import dataclass, field
7
+ from datetime import datetime, timezone
8
+ from typing import Callable, Dict, List, Literal, Optional
9
+
10
+ from .types import BudgetExceededError, CostEvent
11
+
12
+
13
+ @dataclass
14
+ class BudgetPolicy:
15
+ """Defines spending limits and enforcement behavior.
16
+
17
+ Args:
18
+ max_cost: Maximum allowed spend in USD.
19
+ window: Time window for the budget. "total" means lifetime of the tracker.
20
+ on_exceed: Action when budget is exceeded.
21
+ - "raise": raise BudgetExceededError
22
+ - "warn": emit a warning
23
+ - "callback": call the on_exceed_callback function
24
+ on_exceed_callback: Called when on_exceed="callback". Receives the CostEvent
25
+ that would exceed the budget.
26
+ tags: If set, this budget applies only to events matching ALL these tags.
27
+ """
28
+
29
+ max_cost: float
30
+ window: Literal["hourly", "daily", "monthly", "total"] = "total"
31
+ on_exceed: Literal["raise", "warn", "callback"] = "raise"
32
+ on_exceed_callback: Optional[Callable[[CostEvent, float], None]] = field(
33
+ default=None, repr=False
34
+ )
35
+ tags: Optional[Dict[str, str]] = None
36
+
37
+ def _get_window_start(self, now: datetime) -> Optional[datetime]:
38
+ if self.window == "total":
39
+ return None
40
+ if self.window == "hourly":
41
+ return now.replace(minute=0, second=0, microsecond=0)
42
+ if self.window == "daily":
43
+ return now.replace(hour=0, minute=0, second=0, microsecond=0)
44
+ if self.window == "monthly":
45
+ return now.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
46
+ return None
47
+
48
+ def _matches_tags(self, event: CostEvent) -> bool:
49
+ if self.tags is None:
50
+ return True
51
+ return all(event.tags.get(k) == v for k, v in self.tags.items())
52
+
53
+ def current_spend(self, events: List[CostEvent]) -> float:
54
+ """Calculate current spend within the active window."""
55
+ now = datetime.now(timezone.utc)
56
+ window_start = self._get_window_start(now)
57
+ total = 0.0
58
+ for e in events:
59
+ if window_start and e.timestamp < window_start:
60
+ continue
61
+ if not self._matches_tags(e):
62
+ continue
63
+ if e.cost_usd is not None:
64
+ total += e.cost_usd
65
+ return total
66
+
67
+ def check(self, events: List[CostEvent], pending_event: CostEvent) -> None:
68
+ """Check if recording pending_event would exceed the budget.
69
+
70
+ Raises BudgetExceededError or warns depending on on_exceed setting.
71
+ """
72
+ if not self._matches_tags(pending_event):
73
+ return
74
+
75
+ current = self.current_spend(events)
76
+ pending_cost = pending_event.cost_usd or 0.0
77
+ projected = current + pending_cost
78
+
79
+ if projected <= self.max_cost:
80
+ return
81
+
82
+ if self.on_exceed == "raise":
83
+ raise BudgetExceededError(self.max_cost, projected, self.window)
84
+ elif self.on_exceed == "warn":
85
+ warnings.warn(
86
+ f"[compute-cfo] Budget warning: ${projected:.4f} / ${self.max_cost:.4f} "
87
+ f"({self.window} window)",
88
+ stacklevel=3,
89
+ )
90
+ elif self.on_exceed == "callback" and self.on_exceed_callback:
91
+ self.on_exceed_callback(pending_event, projected)
@@ -0,0 +1,87 @@
1
+ """Cost event exporters."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import json
6
+ import sys
7
+ import urllib.request
8
+ from typing import Callable
9
+
10
+ from .types import CostEvent
11
+
12
+ Exporter = Callable[[CostEvent], None]
13
+
14
+
15
+ def console_exporter(event: CostEvent) -> None:
16
+ """Print cost event to stderr."""
17
+ cost_str = f"${event.cost_usd:.6f}" if event.cost_usd is not None else "$?.??????"
18
+ total_tokens = event.input_tokens + event.output_tokens
19
+
20
+ parts = [
21
+ f"\033[90m[compute-cfo]\033[0m",
22
+ f"\033[1m{event.model}\033[0m",
23
+ f"{total_tokens} tokens",
24
+ f"\033[33m{cost_str}\033[0m",
25
+ ]
26
+
27
+ if event.tags:
28
+ tag_str = " ".join(f"{k}:{v}" for k, v in event.tags.items())
29
+ parts.append(f"\033[90m{tag_str}\033[0m")
30
+
31
+ print(" | ".join(parts), file=sys.stderr)
32
+
33
+
34
+ def jsonl_exporter(path: str = "compute_cfo_events.jsonl") -> Exporter:
35
+ """Return an exporter that appends JSON lines to a file."""
36
+
37
+ def _export(event: CostEvent) -> None:
38
+ with open(path, "a") as f:
39
+ f.write(json.dumps(event.to_dict()) + "\n")
40
+
41
+ return _export
42
+
43
+
44
+ def webhook_exporter(url: str) -> Exporter:
45
+ """Return an exporter that POSTs cost events to a URL."""
46
+
47
+ def _export(event: CostEvent) -> None:
48
+ data = json.dumps(event.to_dict()).encode("utf-8")
49
+ req = urllib.request.Request(
50
+ url,
51
+ data=data,
52
+ headers={"Content-Type": "application/json"},
53
+ method="POST",
54
+ )
55
+ try:
56
+ urllib.request.urlopen(req, timeout=5)
57
+ except Exception:
58
+ pass # fire-and-forget
59
+
60
+ return _export
61
+
62
+
63
+ def get_exporter(spec: str) -> Exporter:
64
+ """Parse an exporter specification string and return the exporter.
65
+
66
+ Supported formats:
67
+ "console" — print to stderr
68
+ "jsonl" — write to compute_cfo_events.jsonl
69
+ "jsonl:/path/to/file.jsonl" — write to custom path
70
+ "webhook:https://example.com/hook" — POST to URL
71
+ """
72
+ if spec == "console":
73
+ return console_exporter
74
+
75
+ if spec == "jsonl":
76
+ return jsonl_exporter()
77
+
78
+ if spec.startswith("jsonl:"):
79
+ return jsonl_exporter(spec[6:])
80
+
81
+ if spec.startswith("webhook:"):
82
+ return webhook_exporter(spec[8:])
83
+
84
+ raise ValueError(
85
+ f"Unknown exporter: {spec!r}. "
86
+ f"Use 'console', 'jsonl', 'jsonl:/path', or 'webhook:URL'."
87
+ )
@@ -0,0 +1,96 @@
1
+ """Model pricing database for OpenAI and Anthropic.
2
+
3
+ Prices are in USD per 1 million tokens. Updated March 2026.
4
+ """
5
+
6
+ from __future__ import annotations
7
+
8
+ from typing import Optional, Tuple
9
+
10
+ # {model_id: (input_price_per_1M, output_price_per_1M)}
11
+ MODEL_PRICES: dict[str, Tuple[float, float]] = {
12
+ # ── OpenAI ──────────────────────────────────────────────
13
+ # GPT-4.1 family
14
+ "gpt-4.1": (2.00, 8.00),
15
+ "gpt-4.1-mini": (0.40, 1.60),
16
+ "gpt-4.1-nano": (0.10, 0.40),
17
+ # GPT-4o family
18
+ "gpt-4o": (2.50, 10.00),
19
+ "gpt-4o-mini": (0.15, 0.60),
20
+ "gpt-4o-audio-preview": (2.50, 10.00),
21
+ # GPT-4 legacy
22
+ "gpt-4-turbo": (10.00, 30.00),
23
+ "gpt-4": (30.00, 60.00),
24
+ # GPT-3.5
25
+ "gpt-3.5-turbo": (0.50, 1.50),
26
+ # o-series reasoning
27
+ "o3": (2.00, 8.00),
28
+ "o3-mini": (1.10, 4.40),
29
+ "o4-mini": (1.10, 4.40),
30
+ "o1": (15.00, 60.00),
31
+ "o1-mini": (1.10, 4.40),
32
+ "o1-preview": (15.00, 60.00),
33
+ # Embeddings
34
+ "text-embedding-3-small": (0.02, 0.0),
35
+ "text-embedding-3-large": (0.13, 0.0),
36
+ "text-embedding-ada-002": (0.10, 0.0),
37
+
38
+ # ── Anthropic ───────────────────────────────────────────
39
+ "claude-opus-4-20250514": (15.00, 75.00),
40
+ "claude-sonnet-4-20250514": (3.00, 15.00),
41
+ "claude-3-7-sonnet-20250219": (3.00, 15.00),
42
+ "claude-3-5-sonnet-20241022": (3.00, 15.00),
43
+ "claude-3-5-sonnet-20240620": (3.00, 15.00),
44
+ "claude-3-5-haiku-20241022": (0.80, 4.00),
45
+ "claude-3-opus-20240229": (15.00, 75.00),
46
+ "claude-3-sonnet-20240229": (3.00, 15.00),
47
+ "claude-3-haiku-20240307": (0.25, 1.25),
48
+ # Aliases
49
+ "claude-opus-4-0": (15.00, 75.00),
50
+ "claude-sonnet-4-0": (3.00, 15.00),
51
+ "claude-3.7-sonnet": (3.00, 15.00),
52
+ "claude-3.5-sonnet": (3.00, 15.00),
53
+ "claude-3.5-haiku": (0.80, 4.00),
54
+ "claude-3-opus": (15.00, 75.00),
55
+ "claude-3-sonnet": (3.00, 15.00),
56
+ "claude-3-haiku": (0.25, 1.25),
57
+ }
58
+
59
+ # Common aliases mapping
60
+ _ALIASES: dict[str, str] = {
61
+ "gpt-4o-2024-11-20": "gpt-4o",
62
+ "gpt-4o-2024-08-06": "gpt-4o",
63
+ "gpt-4o-2024-05-13": "gpt-4o",
64
+ "gpt-4o-mini-2024-07-18": "gpt-4o-mini",
65
+ "gpt-4-turbo-2024-04-09": "gpt-4-turbo",
66
+ "gpt-4-turbo-preview": "gpt-4-turbo",
67
+ "gpt-4-0125-preview": "gpt-4-turbo",
68
+ "gpt-4-1106-preview": "gpt-4-turbo",
69
+ "gpt-3.5-turbo-0125": "gpt-3.5-turbo",
70
+ "gpt-3.5-turbo-1106": "gpt-3.5-turbo",
71
+ "o3-2025-04-16": "o3",
72
+ "o4-mini-2025-04-16": "o4-mini",
73
+ }
74
+
75
+
76
+ def resolve_model(model: str) -> str:
77
+ """Resolve model aliases to canonical name."""
78
+ return _ALIASES.get(model, model)
79
+
80
+
81
+ def get_price(model: str) -> Optional[Tuple[float, float]]:
82
+ """Return (input_price, output_price) per 1M tokens, or None if unknown."""
83
+ canonical = resolve_model(model)
84
+ return MODEL_PRICES.get(canonical)
85
+
86
+
87
+ def get_cost(model: str, input_tokens: int, output_tokens: int) -> Optional[float]:
88
+ """Calculate cost in USD for a given model and token counts.
89
+
90
+ Returns None if model pricing is unknown.
91
+ """
92
+ price = get_price(model)
93
+ if price is None:
94
+ return None
95
+ input_price, output_price = price
96
+ return (input_tokens * input_price + output_tokens * output_price) / 1_000_000
@@ -0,0 +1,103 @@
1
+ """CostTracker — core cost accumulation and querying."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import threading
6
+ from collections import defaultdict
7
+ from typing import Dict, List, Optional
8
+
9
+ from .budget import BudgetPolicy
10
+ from .exporters import Exporter, get_exporter
11
+ from .types import CostEvent
12
+
13
+
14
+ class CostTracker:
15
+ """Accumulates cost events and provides spend queries.
16
+
17
+ Thread-safe by default.
18
+
19
+ Args:
20
+ budget: Optional budget policy to enforce.
21
+ export: Exporter specification. Can be:
22
+ - "console" (default): print to stdout
23
+ - "jsonl": append to compute_cfo_events.jsonl
24
+ - "jsonl:/path/to/file.jsonl": custom file path
25
+ - None: no automatic export
26
+ - A callable that receives CostEvent
27
+ quiet: If True, suppress console output even if export="console".
28
+ """
29
+
30
+ def __init__(
31
+ self,
32
+ budget: Optional[BudgetPolicy] = None,
33
+ export: Optional[str] = "console",
34
+ quiet: bool = False,
35
+ ):
36
+ self._events: List[CostEvent] = []
37
+ self._lock = threading.Lock()
38
+ self._budget = budget
39
+ self._exporters: List[Exporter] = []
40
+
41
+ if quiet:
42
+ export = None
43
+
44
+ if export is not None:
45
+ self._exporters.append(get_exporter(export))
46
+
47
+ def record(self, event: CostEvent) -> None:
48
+ """Record a cost event. Checks budget and exports."""
49
+ with self._lock:
50
+ if self._budget:
51
+ self._budget.check(self._events, event)
52
+
53
+ remaining = None
54
+ if self._budget and event.cost_usd is not None:
55
+ spent = self._budget.current_spend(self._events) + event.cost_usd
56
+ remaining = max(0, self._budget.max_cost - spent)
57
+ event.budget_remaining_usd = remaining
58
+
59
+ self._events.append(event)
60
+
61
+ for exporter in self._exporters:
62
+ exporter(event)
63
+
64
+ @property
65
+ def events(self) -> List[CostEvent]:
66
+ """Return a copy of all recorded events."""
67
+ with self._lock:
68
+ return list(self._events)
69
+
70
+ @property
71
+ def total_cost(self) -> float:
72
+ """Total cost across all events."""
73
+ with self._lock:
74
+ return sum(e.cost_usd for e in self._events if e.cost_usd is not None)
75
+
76
+ def cost_by(self, key: str) -> Dict[str, float]:
77
+ """Aggregate cost by a tag key or by 'model'/'provider'.
78
+
79
+ Args:
80
+ key: Tag name, or "model" / "provider" to group by those fields.
81
+
82
+ Returns:
83
+ Dict mapping key values to total cost.
84
+ """
85
+ result: Dict[str, float] = defaultdict(float)
86
+ with self._lock:
87
+ for e in self._events:
88
+ if e.cost_usd is None:
89
+ continue
90
+ if key == "model":
91
+ result[e.model] += e.cost_usd
92
+ elif key == "provider":
93
+ result[e.provider] += e.cost_usd
94
+ else:
95
+ tag_val = e.tags.get(key)
96
+ if tag_val is not None:
97
+ result[tag_val] += e.cost_usd
98
+ return dict(result)
99
+
100
+ def reset(self) -> None:
101
+ """Clear all recorded events."""
102
+ with self._lock:
103
+ self._events.clear()
@@ -0,0 +1,55 @@
1
+ """Core data types for compute-cfo."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from dataclasses import dataclass, field
6
+ from datetime import datetime
7
+ from typing import Any, Dict, Optional
8
+
9
+
10
+ @dataclass
11
+ class CostEvent:
12
+ """A single inference cost event."""
13
+
14
+ timestamp: datetime
15
+ provider: str
16
+ model: str
17
+ operation: str
18
+ input_tokens: int
19
+ output_tokens: int
20
+ cost_usd: Optional[float]
21
+ latency_ms: Optional[float] = None
22
+ tags: Dict[str, str] = field(default_factory=dict)
23
+ budget_remaining_usd: Optional[float] = None
24
+ raw_response: Any = field(default=None, repr=False)
25
+
26
+ def to_dict(self) -> Dict[str, Any]:
27
+ """Serialize to a JSON-compatible dict."""
28
+ d: Dict[str, Any] = {
29
+ "timestamp": self.timestamp.isoformat(),
30
+ "provider": self.provider,
31
+ "model": self.model,
32
+ "operation": self.operation,
33
+ "input_tokens": self.input_tokens,
34
+ "output_tokens": self.output_tokens,
35
+ "cost_usd": self.cost_usd,
36
+ }
37
+ if self.latency_ms is not None:
38
+ d["latency_ms"] = self.latency_ms
39
+ if self.tags:
40
+ d["tags"] = self.tags
41
+ if self.budget_remaining_usd is not None:
42
+ d["budget_remaining_usd"] = self.budget_remaining_usd
43
+ return d
44
+
45
+
46
+ class BudgetExceededError(Exception):
47
+ """Raised when a budget limit has been exceeded."""
48
+
49
+ def __init__(self, limit: float, current: float, window: str):
50
+ self.limit = limit
51
+ self.current = current
52
+ self.window = window
53
+ super().__init__(
54
+ f"Budget exceeded: ${current:.4f} / ${limit:.4f} ({window} window)"
55
+ )
@@ -0,0 +1,173 @@
1
+ """Drop-in wrapper for OpenAI and Anthropic SDK clients."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import time
6
+ from datetime import datetime, timezone
7
+ from typing import Any, Dict, Optional
8
+
9
+ from .pricing import get_cost
10
+ from .tracker import CostTracker
11
+ from .types import CostEvent
12
+
13
+ _DEFAULT_TRACKER: Optional[CostTracker] = None
14
+
15
+
16
+ def _get_or_create_default_tracker() -> CostTracker:
17
+ global _DEFAULT_TRACKER
18
+ if _DEFAULT_TRACKER is None:
19
+ _DEFAULT_TRACKER = CostTracker(export="console")
20
+ return _DEFAULT_TRACKER
21
+
22
+
23
+ def wrap(client: Any, tracker: Optional[CostTracker] = None) -> Any:
24
+ """Wrap an OpenAI or Anthropic client with cost tracking.
25
+
26
+ Args:
27
+ client: An instance of openai.OpenAI or anthropic.Anthropic.
28
+ tracker: Optional CostTracker. If not provided, a default one
29
+ with console output is used.
30
+
31
+ Returns:
32
+ A wrapped client that tracks costs transparently.
33
+ """
34
+ tracker = tracker or _get_or_create_default_tracker()
35
+
36
+ client_type = type(client).__module__.split(".")[0]
37
+
38
+ if client_type == "openai":
39
+ return _OpenAIWrapper(client, tracker)
40
+ elif client_type == "anthropic":
41
+ return _AnthropicWrapper(client, tracker)
42
+ else:
43
+ raise TypeError(
44
+ f"Unsupported client type: {type(client).__name__}. "
45
+ f"Supported: openai.OpenAI, anthropic.Anthropic"
46
+ )
47
+
48
+
49
+ class _TrackedCompletions:
50
+ """Proxy for client.chat.completions with cost tracking."""
51
+
52
+ def __init__(self, completions: Any, tracker: CostTracker):
53
+ self._completions = completions
54
+ self._tracker = tracker
55
+
56
+ def create(self, **kwargs: Any) -> Any:
57
+ tags = kwargs.pop("metadata", None) or {}
58
+ if not isinstance(tags, dict):
59
+ tags = {}
60
+
61
+ model = kwargs.get("model", "unknown")
62
+ start = time.monotonic()
63
+ response = self._completions.create(**kwargs)
64
+ latency_ms = (time.monotonic() - start) * 1000
65
+
66
+ usage = getattr(response, "usage", None)
67
+ input_tokens = getattr(usage, "prompt_tokens", 0) if usage else 0
68
+ output_tokens = getattr(usage, "completion_tokens", 0) if usage else 0
69
+
70
+ actual_model = getattr(response, "model", model)
71
+ cost = get_cost(actual_model, input_tokens, output_tokens)
72
+
73
+ event = CostEvent(
74
+ timestamp=datetime.now(timezone.utc),
75
+ provider="openai",
76
+ model=actual_model,
77
+ operation="chat.completions",
78
+ input_tokens=input_tokens,
79
+ output_tokens=output_tokens,
80
+ cost_usd=cost,
81
+ latency_ms=round(latency_ms, 1),
82
+ tags=tags,
83
+ )
84
+ self._tracker.record(event)
85
+ return response
86
+
87
+ def __getattr__(self, name: str) -> Any:
88
+ return getattr(self._completions, name)
89
+
90
+
91
+ class _TrackedChat:
92
+ """Proxy for client.chat with tracked completions."""
93
+
94
+ def __init__(self, chat: Any, tracker: CostTracker):
95
+ self._chat = chat
96
+ self._tracker = tracker
97
+ self.completions = _TrackedCompletions(chat.completions, tracker)
98
+
99
+ def __getattr__(self, name: str) -> Any:
100
+ return getattr(self._chat, name)
101
+
102
+
103
+ class _OpenAIWrapper:
104
+ """Proxy for OpenAI client."""
105
+
106
+ def __init__(self, client: Any, tracker: CostTracker):
107
+ self._client = client
108
+ self._tracker = tracker
109
+ self.chat = _TrackedChat(client.chat, tracker)
110
+
111
+ def __getattr__(self, name: str) -> Any:
112
+ return getattr(self._client, name)
113
+
114
+
115
+ class _TrackedMessages:
116
+ """Proxy for client.messages (Anthropic) with cost tracking."""
117
+
118
+ def __init__(self, messages: Any, tracker: CostTracker):
119
+ self._messages = messages
120
+ self._tracker = tracker
121
+
122
+ def create(self, **kwargs: Any) -> Any:
123
+ metadata = kwargs.get("metadata", None)
124
+ tags: Dict[str, str] = {}
125
+ if isinstance(metadata, dict):
126
+ # Anthropic metadata has a user_id field; we extract custom tags
127
+ tags = {k: str(v) for k, v in metadata.items() if k != "user_id"}
128
+
129
+ # Also support a compute_cfo_tags kwarg for explicit tagging
130
+ explicit_tags = kwargs.pop("compute_cfo_tags", None)
131
+ if isinstance(explicit_tags, dict):
132
+ tags.update(explicit_tags)
133
+
134
+ model = kwargs.get("model", "unknown")
135
+ start = time.monotonic()
136
+ response = self._messages.create(**kwargs)
137
+ latency_ms = (time.monotonic() - start) * 1000
138
+
139
+ usage = getattr(response, "usage", None)
140
+ input_tokens = getattr(usage, "input_tokens", 0) if usage else 0
141
+ output_tokens = getattr(usage, "output_tokens", 0) if usage else 0
142
+
143
+ actual_model = getattr(response, "model", model)
144
+ cost = get_cost(actual_model, input_tokens, output_tokens)
145
+
146
+ event = CostEvent(
147
+ timestamp=datetime.now(timezone.utc),
148
+ provider="anthropic",
149
+ model=actual_model,
150
+ operation="messages",
151
+ input_tokens=input_tokens,
152
+ output_tokens=output_tokens,
153
+ cost_usd=cost,
154
+ latency_ms=round(latency_ms, 1),
155
+ tags=tags,
156
+ )
157
+ self._tracker.record(event)
158
+ return response
159
+
160
+ def __getattr__(self, name: str) -> Any:
161
+ return getattr(self._messages, name)
162
+
163
+
164
+ class _AnthropicWrapper:
165
+ """Proxy for Anthropic client."""
166
+
167
+ def __init__(self, client: Any, tracker: CostTracker):
168
+ self._client = client
169
+ self._tracker = tracker
170
+ self.messages = _TrackedMessages(client.messages, tracker)
171
+
172
+ def __getattr__(self, name: str) -> Any:
173
+ return getattr(self._client, name)
@@ -0,0 +1,112 @@
1
+ import warnings
2
+ from datetime import datetime, timezone
3
+
4
+ import pytest
5
+
6
+ from compute_cfo.budget import BudgetPolicy
7
+ from compute_cfo.tracker import CostTracker
8
+ from compute_cfo.types import BudgetExceededError, CostEvent
9
+
10
+
11
+ def _make_event(cost: float = 0.01) -> CostEvent:
12
+ return CostEvent(
13
+ timestamp=datetime.now(timezone.utc),
14
+ provider="openai",
15
+ model="gpt-4o",
16
+ operation="chat.completions",
17
+ input_tokens=100,
18
+ output_tokens=50,
19
+ cost_usd=cost,
20
+ )
21
+
22
+
23
+ def test_budget_raise_on_exceed():
24
+ tracker = CostTracker(
25
+ budget=BudgetPolicy(max_cost=0.05, on_exceed="raise"),
26
+ quiet=True,
27
+ )
28
+ tracker.record(_make_event(cost=0.03))
29
+ with pytest.raises(BudgetExceededError) as exc_info:
30
+ tracker.record(_make_event(cost=0.03))
31
+ assert exc_info.value.limit == 0.05
32
+ assert exc_info.value.current > 0.05
33
+
34
+
35
+ def test_budget_warn_on_exceed():
36
+ tracker = CostTracker(
37
+ budget=BudgetPolicy(max_cost=0.05, on_exceed="warn"),
38
+ quiet=True,
39
+ )
40
+ tracker.record(_make_event(cost=0.03))
41
+ with warnings.catch_warnings(record=True) as w:
42
+ warnings.simplefilter("always")
43
+ tracker.record(_make_event(cost=0.03))
44
+ assert len(w) == 1
45
+ assert "Budget warning" in str(w[0].message)
46
+
47
+
48
+ def test_budget_callback_on_exceed():
49
+ exceeded_events = []
50
+
51
+ def on_exceed(event: CostEvent, projected: float):
52
+ exceeded_events.append((event, projected))
53
+
54
+ tracker = CostTracker(
55
+ budget=BudgetPolicy(
56
+ max_cost=0.05,
57
+ on_exceed="callback",
58
+ on_exceed_callback=on_exceed,
59
+ ),
60
+ quiet=True,
61
+ )
62
+ tracker.record(_make_event(cost=0.03))
63
+ tracker.record(_make_event(cost=0.03))
64
+ assert len(exceeded_events) == 1
65
+ assert exceeded_events[0][1] > 0.05
66
+
67
+
68
+ def test_budget_not_exceeded():
69
+ tracker = CostTracker(
70
+ budget=BudgetPolicy(max_cost=1.00, on_exceed="raise"),
71
+ quiet=True,
72
+ )
73
+ tracker.record(_make_event(cost=0.03))
74
+ tracker.record(_make_event(cost=0.03))
75
+ assert abs(tracker.total_cost - 0.06) < 1e-10
76
+
77
+
78
+ def test_budget_remaining_tracked():
79
+ tracker = CostTracker(
80
+ budget=BudgetPolicy(max_cost=1.00),
81
+ quiet=True,
82
+ )
83
+ tracker.record(_make_event(cost=0.30))
84
+ event = tracker.events[-1]
85
+ assert event.budget_remaining_usd is not None
86
+ assert abs(event.budget_remaining_usd - 0.70) < 1e-10
87
+
88
+
89
+ def test_budget_with_tag_filter():
90
+ tracker = CostTracker(
91
+ budget=BudgetPolicy(
92
+ max_cost=0.05,
93
+ on_exceed="raise",
94
+ tags={"project": "expensive"},
95
+ ),
96
+ quiet=True,
97
+ )
98
+ # This should not count toward the budget
99
+ cheap_event = _make_event(cost=0.10)
100
+ cheap_event.tags = {"project": "cheap"}
101
+ tracker.record(cheap_event)
102
+
103
+ # This should count
104
+ exp_event = _make_event(cost=0.03)
105
+ exp_event.tags = {"project": "expensive"}
106
+ tracker.record(exp_event)
107
+
108
+ # This should exceed
109
+ exp_event2 = _make_event(cost=0.03)
110
+ exp_event2.tags = {"project": "expensive"}
111
+ with pytest.raises(BudgetExceededError):
112
+ tracker.record(exp_event2)
@@ -0,0 +1,45 @@
1
+ from compute_cfo.pricing import get_cost, get_price, resolve_model
2
+
3
+
4
+ def test_known_model_price():
5
+ price = get_price("gpt-4o")
6
+ assert price == (2.50, 10.00)
7
+
8
+
9
+ def test_anthropic_model_price():
10
+ price = get_price("claude-sonnet-4-20250514")
11
+ assert price == (3.00, 15.00)
12
+
13
+
14
+ def test_unknown_model_returns_none():
15
+ assert get_price("nonexistent-model") is None
16
+
17
+
18
+ def test_alias_resolution():
19
+ assert resolve_model("gpt-4o-2024-11-20") == "gpt-4o"
20
+ assert resolve_model("gpt-4o") == "gpt-4o" # not an alias, returns as-is
21
+
22
+
23
+ def test_get_cost_calculation():
24
+ # gpt-4o: $2.50 input, $10.00 output per 1M tokens
25
+ cost = get_cost("gpt-4o", input_tokens=1000, output_tokens=500)
26
+ assert cost is not None
27
+ expected = (1000 * 2.50 + 500 * 10.00) / 1_000_000
28
+ assert abs(cost - expected) < 1e-10
29
+
30
+
31
+ def test_get_cost_unknown_model():
32
+ assert get_cost("nonexistent", 100, 100) is None
33
+
34
+
35
+ def test_get_cost_alias():
36
+ cost_alias = get_cost("gpt-4o-2024-11-20", 1000, 500)
37
+ cost_direct = get_cost("gpt-4o", 1000, 500)
38
+ assert cost_alias == cost_direct
39
+
40
+
41
+ def test_embedding_model_no_output_cost():
42
+ cost = get_cost("text-embedding-3-small", input_tokens=1000, output_tokens=0)
43
+ assert cost is not None
44
+ expected = (1000 * 0.02) / 1_000_000
45
+ assert abs(cost - expected) < 1e-10
@@ -0,0 +1,80 @@
1
+ from datetime import datetime, timezone
2
+
3
+ from compute_cfo.tracker import CostTracker
4
+ from compute_cfo.types import CostEvent
5
+
6
+
7
+ def _make_event(
8
+ cost: float = 0.01,
9
+ model: str = "gpt-4o",
10
+ provider: str = "openai",
11
+ tags: dict | None = None,
12
+ ) -> CostEvent:
13
+ return CostEvent(
14
+ timestamp=datetime.now(timezone.utc),
15
+ provider=provider,
16
+ model=model,
17
+ operation="chat.completions",
18
+ input_tokens=100,
19
+ output_tokens=50,
20
+ cost_usd=cost,
21
+ tags=tags or {},
22
+ )
23
+
24
+
25
+ def test_empty_tracker():
26
+ t = CostTracker(quiet=True)
27
+ assert t.total_cost == 0.0
28
+ assert t.events == []
29
+
30
+
31
+ def test_record_and_total():
32
+ t = CostTracker(quiet=True)
33
+ t.record(_make_event(cost=0.05))
34
+ t.record(_make_event(cost=0.03))
35
+ assert abs(t.total_cost - 0.08) < 1e-10
36
+
37
+
38
+ def test_cost_by_model():
39
+ t = CostTracker(quiet=True)
40
+ t.record(_make_event(cost=0.05, model="gpt-4o"))
41
+ t.record(_make_event(cost=0.01, model="gpt-4o-mini"))
42
+ t.record(_make_event(cost=0.03, model="gpt-4o"))
43
+ by_model = t.cost_by("model")
44
+ assert abs(by_model["gpt-4o"] - 0.08) < 1e-10
45
+ assert abs(by_model["gpt-4o-mini"] - 0.01) < 1e-10
46
+
47
+
48
+ def test_cost_by_tag():
49
+ t = CostTracker(quiet=True)
50
+ t.record(_make_event(cost=0.05, tags={"project": "search"}))
51
+ t.record(_make_event(cost=0.03, tags={"project": "chat"}))
52
+ t.record(_make_event(cost=0.02, tags={"project": "search"}))
53
+ by_project = t.cost_by("project")
54
+ assert abs(by_project["search"] - 0.07) < 1e-10
55
+ assert abs(by_project["chat"] - 0.03) < 1e-10
56
+
57
+
58
+ def test_cost_by_provider():
59
+ t = CostTracker(quiet=True)
60
+ t.record(_make_event(cost=0.05, provider="openai"))
61
+ t.record(_make_event(cost=0.03, provider="anthropic"))
62
+ by_provider = t.cost_by("provider")
63
+ assert abs(by_provider["openai"] - 0.05) < 1e-10
64
+ assert abs(by_provider["anthropic"] - 0.03) < 1e-10
65
+
66
+
67
+ def test_reset():
68
+ t = CostTracker(quiet=True)
69
+ t.record(_make_event(cost=0.05))
70
+ t.reset()
71
+ assert t.total_cost == 0.0
72
+ assert t.events == []
73
+
74
+
75
+ def test_events_returns_copy():
76
+ t = CostTracker(quiet=True)
77
+ t.record(_make_event(cost=0.05))
78
+ events = t.events
79
+ events.clear()
80
+ assert len(t.events) == 1 # original not affected
@@ -0,0 +1,147 @@
1
+ """Tests for the wrap() function using mock SDK clients."""
2
+
3
+ from datetime import datetime, timezone
4
+ from types import SimpleNamespace
5
+ from unittest.mock import MagicMock
6
+
7
+ from compute_cfo.tracker import CostTracker
8
+ from compute_cfo.wrapper import wrap
9
+
10
+
11
+ def _make_openai_client():
12
+ """Create a mock OpenAI client."""
13
+ client = MagicMock()
14
+ client.__class__.__module__ = "openai.client"
15
+ # Simulate response
16
+ response = SimpleNamespace(
17
+ model="gpt-4o",
18
+ usage=SimpleNamespace(prompt_tokens=100, completion_tokens=50),
19
+ choices=[SimpleNamespace(message=SimpleNamespace(content="Hello!"))],
20
+ )
21
+ client.chat.completions.create.return_value = response
22
+ return client
23
+
24
+
25
+ def _make_anthropic_client():
26
+ """Create a mock Anthropic client."""
27
+ client = MagicMock()
28
+ client.__class__.__module__ = "anthropic.client"
29
+ response = SimpleNamespace(
30
+ model="claude-sonnet-4-20250514",
31
+ usage=SimpleNamespace(input_tokens=100, output_tokens=50),
32
+ content=[SimpleNamespace(text="Hello!")],
33
+ )
34
+ client.messages.create.return_value = response
35
+ return client
36
+
37
+
38
+ def test_wrap_openai_tracks_cost():
39
+ tracker = CostTracker(quiet=True)
40
+ client = _make_openai_client()
41
+ wrapped = wrap(client, tracker=tracker)
42
+
43
+ response = wrapped.chat.completions.create(
44
+ model="gpt-4o",
45
+ messages=[{"role": "user", "content": "hi"}],
46
+ )
47
+
48
+ assert response.model == "gpt-4o"
49
+ assert len(tracker.events) == 1
50
+ event = tracker.events[0]
51
+ assert event.provider == "openai"
52
+ assert event.model == "gpt-4o"
53
+ assert event.input_tokens == 100
54
+ assert event.output_tokens == 50
55
+ assert event.cost_usd is not None
56
+ assert event.cost_usd > 0
57
+
58
+
59
+ def test_wrap_openai_with_tags():
60
+ tracker = CostTracker(quiet=True)
61
+ client = _make_openai_client()
62
+ wrapped = wrap(client, tracker=tracker)
63
+
64
+ wrapped.chat.completions.create(
65
+ model="gpt-4o",
66
+ messages=[{"role": "user", "content": "hi"}],
67
+ metadata={"project": "search", "agent": "summarizer"},
68
+ )
69
+
70
+ event = tracker.events[0]
71
+ assert event.tags == {"project": "search", "agent": "summarizer"}
72
+
73
+
74
+ def test_wrap_anthropic_tracks_cost():
75
+ tracker = CostTracker(quiet=True)
76
+ client = _make_anthropic_client()
77
+ wrapped = wrap(client, tracker=tracker)
78
+
79
+ response = wrapped.messages.create(
80
+ model="claude-sonnet-4-20250514",
81
+ max_tokens=1024,
82
+ messages=[{"role": "user", "content": "hi"}],
83
+ )
84
+
85
+ assert len(tracker.events) == 1
86
+ event = tracker.events[0]
87
+ assert event.provider == "anthropic"
88
+ assert event.model == "claude-sonnet-4-20250514"
89
+ assert event.input_tokens == 100
90
+ assert event.output_tokens == 50
91
+ assert event.cost_usd is not None
92
+
93
+
94
+ def test_wrap_anthropic_with_tags():
95
+ tracker = CostTracker(quiet=True)
96
+ client = _make_anthropic_client()
97
+ wrapped = wrap(client, tracker=tracker)
98
+
99
+ wrapped.messages.create(
100
+ model="claude-sonnet-4-20250514",
101
+ max_tokens=1024,
102
+ messages=[{"role": "user", "content": "hi"}],
103
+ compute_cfo_tags={"project": "search"},
104
+ )
105
+
106
+ event = tracker.events[0]
107
+ assert event.tags == {"project": "search"}
108
+
109
+
110
+ def test_wrap_passthrough_attributes():
111
+ """Non-tracked attributes should pass through to the original client."""
112
+ tracker = CostTracker(quiet=True)
113
+ client = _make_openai_client()
114
+ client.models = MagicMock()
115
+ client.models.list.return_value = ["gpt-4o"]
116
+
117
+ wrapped = wrap(client, tracker=tracker)
118
+ result = wrapped.models.list()
119
+ assert result == ["gpt-4o"]
120
+
121
+
122
+ def test_wrap_unsupported_client():
123
+ """Should raise TypeError for unsupported clients."""
124
+ import pytest
125
+
126
+ tracker = CostTracker(quiet=True)
127
+
128
+ class FakeClient:
129
+ pass
130
+
131
+ with pytest.raises(TypeError, match="Unsupported client type"):
132
+ wrap(FakeClient(), tracker=tracker)
133
+
134
+
135
+ def test_wrap_records_latency():
136
+ tracker = CostTracker(quiet=True)
137
+ client = _make_openai_client()
138
+ wrapped = wrap(client, tracker=tracker)
139
+
140
+ wrapped.chat.completions.create(
141
+ model="gpt-4o",
142
+ messages=[{"role": "user", "content": "hi"}],
143
+ )
144
+
145
+ event = tracker.events[0]
146
+ assert event.latency_ms is not None
147
+ assert event.latency_ms >= 0