agentcontrolroom 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.
Files changed (25) hide show
  1. agentcontrolroom-0.1.0/PKG-INFO +38 -0
  2. agentcontrolroom-0.1.0/pyproject.toml +67 -0
  3. agentcontrolroom-0.1.0/setup.cfg +4 -0
  4. agentcontrolroom-0.1.0/src/agentcontrolroom/__init__.py +82 -0
  5. agentcontrolroom-0.1.0/src/agentcontrolroom/__main__.py +11 -0
  6. agentcontrolroom-0.1.0/src/agentcontrolroom/_cli.py +188 -0
  7. agentcontrolroom-0.1.0/src/agentcontrolroom/_config.py +40 -0
  8. agentcontrolroom-0.1.0/src/agentcontrolroom/client.py +149 -0
  9. agentcontrolroom-0.1.0/src/agentcontrolroom/cost.py +112 -0
  10. agentcontrolroom-0.1.0/src/agentcontrolroom/guardrails.py +154 -0
  11. agentcontrolroom-0.1.0/src/agentcontrolroom/instruments/__init__.py +21 -0
  12. agentcontrolroom-0.1.0/src/agentcontrolroom/instruments/crewai.py +60 -0
  13. agentcontrolroom-0.1.0/src/agentcontrolroom/instruments/langchain.py +406 -0
  14. agentcontrolroom-0.1.0/src/agentcontrolroom/instruments/llamaindex.py +66 -0
  15. agentcontrolroom-0.1.0/src/agentcontrolroom/spans.py +145 -0
  16. agentcontrolroom-0.1.0/src/agentcontrolroom/tracer.py +388 -0
  17. agentcontrolroom-0.1.0/src/agentcontrolroom.egg-info/PKG-INFO +38 -0
  18. agentcontrolroom-0.1.0/src/agentcontrolroom.egg-info/SOURCES.txt +23 -0
  19. agentcontrolroom-0.1.0/src/agentcontrolroom.egg-info/dependency_links.txt +1 -0
  20. agentcontrolroom-0.1.0/src/agentcontrolroom.egg-info/entry_points.txt +2 -0
  21. agentcontrolroom-0.1.0/src/agentcontrolroom.egg-info/requires.txt +20 -0
  22. agentcontrolroom-0.1.0/src/agentcontrolroom.egg-info/top_level.txt +1 -0
  23. agentcontrolroom-0.1.0/tests/test_cost.py +64 -0
  24. agentcontrolroom-0.1.0/tests/test_spans.py +104 -0
  25. agentcontrolroom-0.1.0/tests/test_tracer.py +87 -0
@@ -0,0 +1,38 @@
1
+ Metadata-Version: 2.4
2
+ Name: agentcontrolroom
3
+ Version: 0.1.0
4
+ Summary: The Reliability Layer for Autonomous AI Agents — full tracing, cost intelligence, guardrails, and quality evaluation.
5
+ Author-email: Darshan <darshan@agentcontrolroom.io>
6
+ License: MIT
7
+ Project-URL: Homepage, https://agentcontrolroom.io
8
+ Project-URL: Repository, https://github.com/Darshan-1812/Helm-AI
9
+ Project-URL: Documentation, https://github.com/Darshan-1812/Helm-AI#readme
10
+ Project-URL: Bug Tracker, https://github.com/Darshan-1812/Helm-AI/issues
11
+ Keywords: ai,agents,llm,observability,tracing,langchain,openai,monitoring,guardrails
12
+ Classifier: Development Status :: 4 - Beta
13
+ Classifier: Intended Audience :: Developers
14
+ Classifier: License :: OSI Approved :: MIT License
15
+ Classifier: Programming Language :: Python :: 3
16
+ Classifier: Programming Language :: Python :: 3.10
17
+ Classifier: Programming Language :: Python :: 3.11
18
+ Classifier: Programming Language :: Python :: 3.12
19
+ Classifier: Topic :: Software Development :: Libraries :: Python Modules
20
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
21
+ Requires-Python: >=3.10
22
+ Description-Content-Type: text/markdown
23
+ Requires-Dist: httpx>=0.27.0
24
+ Requires-Dist: pydantic>=2.0.0
25
+ Requires-Dist: opentelemetry-api>=1.25.0
26
+ Requires-Dist: opentelemetry-sdk>=1.25.0
27
+ Provides-Extra: langchain
28
+ Requires-Dist: langchain-core>=0.2.0; extra == "langchain"
29
+ Provides-Extra: crewai
30
+ Requires-Dist: crewai>=0.30.0; extra == "crewai"
31
+ Provides-Extra: llamaindex
32
+ Requires-Dist: llama-index-core>=0.10.0; extra == "llamaindex"
33
+ Provides-Extra: dev
34
+ Requires-Dist: pytest>=8.0.0; extra == "dev"
35
+ Requires-Dist: pytest-asyncio>=0.23.0; extra == "dev"
36
+ Requires-Dist: ruff>=0.5.0; extra == "dev"
37
+ Requires-Dist: build>=1.0.0; extra == "dev"
38
+ Requires-Dist: twine>=5.0.0; extra == "dev"
@@ -0,0 +1,67 @@
1
+ [build-system]
2
+ requires = ["setuptools>=68.0", "wheel"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "agentcontrolroom"
7
+ version = "0.1.0"
8
+ description = "The Reliability Layer for Autonomous AI Agents — full tracing, cost intelligence, guardrails, and quality evaluation."
9
+ readme = "README.md"
10
+ license = { text = "MIT" }
11
+ requires-python = ">=3.10"
12
+ authors = [
13
+ { name = "Darshan", email = "darshan@agentcontrolroom.io" },
14
+ ]
15
+ keywords = [
16
+ "ai", "agents", "llm", "observability", "tracing",
17
+ "langchain", "openai", "monitoring", "guardrails",
18
+ ]
19
+ classifiers = [
20
+ "Development Status :: 4 - Beta",
21
+ "Intended Audience :: Developers",
22
+ "License :: OSI Approved :: MIT License",
23
+ "Programming Language :: Python :: 3",
24
+ "Programming Language :: Python :: 3.10",
25
+ "Programming Language :: Python :: 3.11",
26
+ "Programming Language :: Python :: 3.12",
27
+ "Topic :: Software Development :: Libraries :: Python Modules",
28
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
29
+ ]
30
+
31
+ dependencies = [
32
+ "httpx>=0.27.0",
33
+ "pydantic>=2.0.0",
34
+ "opentelemetry-api>=1.25.0",
35
+ "opentelemetry-sdk>=1.25.0",
36
+ ]
37
+
38
+ [project.urls]
39
+ Homepage = "https://agentcontrolroom.io"
40
+ Repository = "https://github.com/Darshan-1812/Helm-AI"
41
+ Documentation = "https://github.com/Darshan-1812/Helm-AI#readme"
42
+ "Bug Tracker" = "https://github.com/Darshan-1812/Helm-AI/issues"
43
+
44
+ [project.optional-dependencies]
45
+ langchain = ["langchain-core>=0.2.0"]
46
+ crewai = ["crewai>=0.30.0"]
47
+ llamaindex = ["llama-index-core>=0.10.0"]
48
+ dev = [
49
+ "pytest>=8.0.0",
50
+ "pytest-asyncio>=0.23.0",
51
+ "ruff>=0.5.0",
52
+ "build>=1.0.0",
53
+ "twine>=5.0.0",
54
+ ]
55
+
56
+ [project.scripts]
57
+ acr-check = "agentcontrolroom._cli:main"
58
+
59
+ [tool.setuptools.packages.find]
60
+ where = ["src"]
61
+
62
+ [tool.setuptools.package-data]
63
+ "agentcontrolroom" = ["py.typed"]
64
+
65
+ [tool.ruff]
66
+ target-version = "py312"
67
+ line-length = 100
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,82 @@
1
+ """
2
+ Agent Control Room SDK
3
+ ======================
4
+
5
+ Reliability Layer for Autonomous AI Agents.
6
+
7
+ Quick Start (env-var based — zero code config):
8
+ export ACR_API_KEY=acr-dev-xxxx
9
+ export ACR_ENDPOINT=https://your-backend.onrender.com
10
+
11
+ from agentcontrolroom import trace
12
+
13
+ @trace.agent(name="my-agent")
14
+ def my_agent(query: str):
15
+ result = llm_call(query)
16
+ return result
17
+
18
+ @trace.tool(name="web-search")
19
+ def search(query: str):
20
+ ...
21
+
22
+ Manual config:
23
+ from agentcontrolroom import configure, trace
24
+
25
+ configure(api_key="acr-dev-xxxx", endpoint="https://...")
26
+
27
+ @trace.agent("my-agent")
28
+ def my_agent(query: str):
29
+ ...
30
+ """
31
+
32
+ from agentcontrolroom._config import get_api_key, get_endpoint, is_configured
33
+ from agentcontrolroom.tracer import trace, Tracer
34
+ from agentcontrolroom.client import ACRClient
35
+ from agentcontrolroom.spans import SpanKind
36
+ from agentcontrolroom.cost import CostCalculator
37
+ from agentcontrolroom.guardrails import Guardrails
38
+
39
+ __version__ = "0.1.0"
40
+
41
+
42
+ def configure(
43
+ api_key: str,
44
+ endpoint: str = "http://localhost:8000",
45
+ ) -> None:
46
+ """
47
+ Configure the global tracer with your API key and backend endpoint.
48
+
49
+ This is optional if you set ACR_API_KEY and ACR_ENDPOINT env vars.
50
+
51
+ Args:
52
+ api_key: Your Agent Control Room API key (e.g. "acr-dev-xxxx")
53
+ endpoint: URL of your ACR backend (default: http://localhost:8000)
54
+
55
+ Example:
56
+ from agentcontrolroom import configure
57
+ configure(api_key="acr-dev-xxxx", endpoint="https://acr.myapp.com")
58
+ """
59
+ trace.configure(api_key=api_key, endpoint=endpoint)
60
+
61
+
62
+ # ── Auto-configure from env vars at import time ───────────────────────────────
63
+ if is_configured():
64
+ configure(api_key=get_api_key(), endpoint=get_endpoint())
65
+
66
+
67
+ __all__ = [
68
+ # Core
69
+ "trace",
70
+ "Tracer",
71
+ "ACRClient",
72
+ "SpanKind",
73
+ "CostCalculator",
74
+ "Guardrails",
75
+ # Config helpers
76
+ "configure",
77
+ "get_api_key",
78
+ "get_endpoint",
79
+ "is_configured",
80
+ # Version
81
+ "__version__",
82
+ ]
@@ -0,0 +1,11 @@
1
+ """
2
+ Enable `python -m agentcontrolroom` entry point.
3
+
4
+ Usage:
5
+ python -m agentcontrolroom check
6
+ python -m agentcontrolroom check --api-key acr-dev-xxxx --endpoint https://...
7
+ """
8
+
9
+ from agentcontrolroom._cli import main
10
+
11
+ main()
@@ -0,0 +1,188 @@
1
+ """
2
+ acr-check CLI — verify your Agent Control Room setup is working.
3
+
4
+ Usage:
5
+ python -m agentcontrolroom check
6
+ python -m agentcontrolroom check --api-key acr-dev-xxxx --endpoint https://...
7
+ acr-check # if installed via pip
8
+ acr-check --api-key acr-dev-xxxx
9
+ """
10
+
11
+ import argparse
12
+ import sys
13
+ import os
14
+
15
+ # ── ANSI colours (disabled on Windows if no support) ─────────────────────────
16
+ GREEN = "\033[92m"
17
+ RED = "\033[91m"
18
+ YELLOW = "\033[93m"
19
+ CYAN = "\033[96m"
20
+ BOLD = "\033[1m"
21
+ RESET = "\033[0m"
22
+
23
+ OK = f"{GREEN}✅{RESET}"
24
+ FAIL = f"{RED}❌{RESET}"
25
+ WARN = f"{YELLOW}⚠️ {RESET}"
26
+ INFO = f"{CYAN}ℹ️ {RESET}"
27
+
28
+
29
+ def _print_header():
30
+ print(f"\n{BOLD}{'=' * 58}{RESET}")
31
+ print(f"{BOLD} Agent Control Room — System Health Check{RESET}")
32
+ print(f"{BOLD}{'=' * 58}{RESET}\n")
33
+
34
+
35
+ def _check_backend(client) -> bool:
36
+ """Check if the backend API is reachable and healthy."""
37
+ try:
38
+ health = client.health_check()
39
+ version = health.get("version", "unknown")
40
+ print(f" {OK} Backend reachable (version: {version})")
41
+ return True
42
+ except Exception as e:
43
+ print(f" {FAIL} Backend unreachable: {e}")
44
+ print(" → Is the backend running? Try: docker-compose up -d")
45
+ return False
46
+
47
+
48
+ def _check_auth(client) -> bool:
49
+ """Verify the API key is valid."""
50
+ try:
51
+ resp = client._client.get("/api/v1/runs")
52
+ if resp.status_code == 401:
53
+ print(f" {FAIL} API key invalid — got 401 Unauthorized")
54
+ print(" → Check your ACR_API_KEY value")
55
+ return False
56
+ elif resp.status_code in (200, 422):
57
+ print(f" {OK} API key valid")
58
+ return True
59
+ else:
60
+ print(f" {WARN} Unexpected status {resp.status_code} from /api/v1/runs")
61
+ return True
62
+ except Exception as e:
63
+ print(f" {FAIL} Auth check failed: {e}")
64
+ return False
65
+
66
+
67
+ def _check_trace_ingest(client) -> bool:
68
+ """Send a test trace and verify it is ingested."""
69
+ from agentcontrolroom.spans import RunData, SpanData, SpanKind
70
+
71
+ run = RunData(agent_name="acr-health-check", input_text="system health check")
72
+ span = SpanData(
73
+ name="health-check-span",
74
+ span_kind=SpanKind.TOOL,
75
+ run_id=run.run_id,
76
+ input_data="ping",
77
+ output_data="pong",
78
+ latency_ms=1.0,
79
+ )
80
+ span.finish()
81
+ run.add_span(span)
82
+ run.finish(output="ok", status="completed")
83
+
84
+ try:
85
+ result = client.send_run(run)
86
+ ingested = result.get("spans_ingested", 0)
87
+ print(f" {OK} Test trace ingested ({ingested} span(s) stored)")
88
+ return True
89
+ except Exception as e:
90
+ print(f" {FAIL} Trace ingest failed: {e}")
91
+ print(" → Is PostgreSQL + Dramatiq worker running?")
92
+ return False
93
+
94
+
95
+ def _check_dashboard(endpoint: str) -> bool:
96
+ """Verify the frontend dashboard is reachable."""
97
+ import httpx
98
+ # Try port 3001 (docker-compose) or 3000 (dev server)
99
+ for port in [3001, 3000]:
100
+ url = f"http://localhost:{port}"
101
+ try:
102
+ resp = httpx.get(url, timeout=3.0, follow_redirects=True)
103
+ if resp.status_code < 400:
104
+ print(f" {OK} Dashboard reachable at {url}")
105
+ return True
106
+ except Exception:
107
+ continue
108
+
109
+ print(f" {WARN} Dashboard not reachable on :3001 or :3000")
110
+ print(" → Run: docker-compose up -d frontend")
111
+ return False
112
+
113
+
114
+ def run_check(api_key: str, endpoint: str) -> int:
115
+ """Run all checks. Returns 0 if all pass, 1 if any fail."""
116
+ from agentcontrolroom.client import ACRClient
117
+
118
+ print(f" {INFO} Endpoint : {endpoint}")
119
+ print(f" {INFO} API Key : {api_key[:12]}{'*' * max(0, len(api_key) - 12)}\n")
120
+
121
+ results = []
122
+
123
+ client = ACRClient(api_key=api_key, endpoint=endpoint, auto_flush=False)
124
+
125
+ try:
126
+ results.append(_check_backend(client))
127
+ if results[-1]:
128
+ results.append(_check_auth(client))
129
+ if results[-1]:
130
+ results.append(_check_trace_ingest(client))
131
+ results.append(_check_dashboard(endpoint))
132
+ finally:
133
+ client.close()
134
+
135
+ print()
136
+ passed = sum(results)
137
+ total = len(results)
138
+
139
+ if all(results):
140
+ print(f"{BOLD}{GREEN}{'=' * 58}{RESET}")
141
+ print(f"{BOLD}{GREEN} 🎉 All checks passed! ({passed}/{total}){RESET}")
142
+ print(f"{BOLD}{GREEN} Your Agent Control Room is working correctly.{RESET}")
143
+ print(f"{BOLD}{GREEN}{'=' * 58}{RESET}\n")
144
+ return 0
145
+ else:
146
+ print(f"{BOLD}{RED}{'=' * 58}{RESET}")
147
+ print(f"{BOLD}{RED} {passed}/{total} checks passed — see issues above.{RESET}")
148
+ print(f"{BOLD}{RED}{'=' * 58}{RESET}\n")
149
+ return 1
150
+
151
+
152
+ def main():
153
+ _print_header()
154
+
155
+ parser = argparse.ArgumentParser(
156
+ prog="acr-check",
157
+ description="Verify your Agent Control Room setup is working correctly.",
158
+ )
159
+ parser.add_argument(
160
+ "--api-key",
161
+ default=os.environ.get("ACR_API_KEY", ""),
162
+ help="Your ACR API key (default: $ACR_API_KEY)",
163
+ )
164
+ parser.add_argument(
165
+ "--endpoint",
166
+ default=os.environ.get("ACR_ENDPOINT", "http://localhost:8000"),
167
+ help="Backend endpoint URL (default: $ACR_ENDPOINT or http://localhost:8000)",
168
+ )
169
+
170
+ # Support `python -m agentcontrolroom check` subcommand
171
+ if len(sys.argv) > 1 and sys.argv[1] == "check":
172
+ sys.argv.pop(1)
173
+
174
+ args = parser.parse_args()
175
+
176
+ if not args.api_key:
177
+ print(f" {FAIL} No API key provided.\n")
178
+ print(" Set it via:")
179
+ print(" export ACR_API_KEY=acr-dev-xxxx")
180
+ print(" or:")
181
+ print(" acr-check --api-key acr-dev-xxxx\n")
182
+ sys.exit(1)
183
+
184
+ sys.exit(run_check(api_key=args.api_key, endpoint=args.endpoint))
185
+
186
+
187
+ if __name__ == "__main__":
188
+ main()
@@ -0,0 +1,40 @@
1
+ """
2
+ Agent Control Room — centralized environment-variable configuration.
3
+
4
+ Reads ACR_API_KEY and ACR_ENDPOINT from the environment so users can configure
5
+ the SDK without writing any Python code:
6
+
7
+ export ACR_API_KEY=acr-dev-xxxx
8
+ export ACR_ENDPOINT=https://your-acr-backend.onrender.com
9
+
10
+ # Then in Python — just import, no configure() call needed:
11
+ from agentcontrolroom import trace
12
+
13
+ @trace.agent("my-agent")
14
+ def my_agent(query: str):
15
+ ...
16
+ """
17
+
18
+ import os
19
+ import logging
20
+
21
+ logger = logging.getLogger("agentcontrolroom")
22
+
23
+ # ── Public defaults (read once at import time) ────────────────────────────────
24
+ ACR_API_KEY: str = os.environ.get("ACR_API_KEY", "")
25
+ ACR_ENDPOINT: str = os.environ.get("ACR_ENDPOINT", "http://localhost:8000")
26
+
27
+
28
+ def get_api_key() -> str:
29
+ """Return the currently configured API key (env var or manually set)."""
30
+ return os.environ.get("ACR_API_KEY", ACR_API_KEY)
31
+
32
+
33
+ def get_endpoint() -> str:
34
+ """Return the currently configured backend endpoint."""
35
+ return os.environ.get("ACR_ENDPOINT", ACR_ENDPOINT)
36
+
37
+
38
+ def is_configured() -> bool:
39
+ """Return True if an API key is available."""
40
+ return bool(get_api_key())
@@ -0,0 +1,149 @@
1
+ """
2
+ HTTP Client — ships trace data to the Agent Control Room backend.
3
+ Supports batching and async flush.
4
+ """
5
+
6
+ import threading
7
+ import time
8
+ import logging
9
+ from typing import Optional
10
+
11
+ import httpx
12
+
13
+ from agentcontrolroom.spans import RunData
14
+
15
+ logger = logging.getLogger("agentcontrolroom")
16
+
17
+
18
+ class ACRClient:
19
+ """
20
+ HTTP client for Agent Control Room.
21
+
22
+ Usage:
23
+ client = ACRClient(
24
+ api_key="acr-dev-xxxx",
25
+ endpoint="http://localhost:8000",
26
+ )
27
+
28
+ # Ship a complete run
29
+ client.send_run(run_data)
30
+
31
+ # Or use auto-flush with batching
32
+ client.queue_run(run_data)
33
+ client.flush() # or let background thread flush
34
+ """
35
+
36
+ def __init__(
37
+ self,
38
+ api_key: str,
39
+ endpoint: str = "http://localhost:8000",
40
+ batch_size: int = 10,
41
+ flush_interval: float = 5.0,
42
+ timeout: float = 30.0,
43
+ auto_flush: bool = True,
44
+ ):
45
+ self.api_key = api_key
46
+ self.endpoint = endpoint.rstrip("/")
47
+ self.batch_size = batch_size
48
+ self.flush_interval = flush_interval
49
+ self.timeout = timeout
50
+
51
+ self._queue: list[RunData] = []
52
+ self._lock = threading.Lock()
53
+ self._client = httpx.Client(
54
+ base_url=self.endpoint,
55
+ headers={
56
+ "X-API-Key": self.api_key,
57
+ "Content-Type": "application/json",
58
+ },
59
+ timeout=timeout,
60
+ )
61
+
62
+ # Background flush thread
63
+ self._running = False
64
+ self._flush_thread: Optional[threading.Thread] = None
65
+ if auto_flush:
66
+ self._start_flush_thread()
67
+
68
+ def send_run(self, run: RunData) -> dict:
69
+ """Send a complete run synchronously to the backend."""
70
+ try:
71
+ payload = run.to_dict()
72
+ response = self._client.post(
73
+ "/api/v1/ingest/traces",
74
+ json=payload,
75
+ )
76
+ response.raise_for_status()
77
+ result = response.json()
78
+ logger.info(
79
+ f"Trace sent: run_id={result.get('run_id')}, "
80
+ f"spans={result.get('spans_ingested')}"
81
+ )
82
+ return result
83
+ except httpx.HTTPStatusError as e:
84
+ logger.error(f"Failed to send trace: {e.response.status_code} {e.response.text}")
85
+ raise
86
+ except Exception as e:
87
+ logger.error(f"Failed to send trace: {e}")
88
+ raise
89
+
90
+ def queue_run(self, run: RunData):
91
+ """Queue a run for batch sending."""
92
+ with self._lock:
93
+ self._queue.append(run)
94
+ if len(self._queue) >= self.batch_size:
95
+ self._flush_batch()
96
+
97
+ def flush(self):
98
+ """Flush all queued runs."""
99
+ with self._lock:
100
+ self._flush_batch()
101
+
102
+ def _flush_batch(self):
103
+ """Send all queued runs (called under lock)."""
104
+ if not self._queue:
105
+ return
106
+
107
+ batch = list(self._queue)
108
+ self._queue.clear()
109
+
110
+ for run in batch:
111
+ try:
112
+ self.send_run(run)
113
+ except Exception as e:
114
+ logger.error(f"Failed to flush run {run.run_id}: {e}")
115
+
116
+ def _start_flush_thread(self):
117
+ """Start background thread for periodic flushing."""
118
+ self._running = True
119
+ self._flush_thread = threading.Thread(
120
+ target=self._flush_loop, daemon=True, name="acr-flush"
121
+ )
122
+ self._flush_thread.start()
123
+
124
+ def _flush_loop(self):
125
+ """Background flush loop."""
126
+ while self._running:
127
+ time.sleep(self.flush_interval)
128
+ try:
129
+ self.flush()
130
+ except Exception as e:
131
+ logger.error(f"Background flush error: {e}")
132
+
133
+ def close(self):
134
+ """Flush remaining data and close the client."""
135
+ self._running = False
136
+ self.flush()
137
+ self._client.close()
138
+
139
+ def health_check(self) -> dict:
140
+ """Check if the backend is healthy."""
141
+ response = self._client.get("/health")
142
+ response.raise_for_status()
143
+ return response.json()
144
+
145
+ def __enter__(self):
146
+ return self
147
+
148
+ def __exit__(self, *args):
149
+ self.close()
@@ -0,0 +1,112 @@
1
+ """
2
+ Token-to-cost calculator with pricing tables for popular LLM models.
3
+ """
4
+
5
+ from typing import Optional
6
+ from dataclasses import dataclass
7
+
8
+
9
+ @dataclass
10
+ class ModelPricing:
11
+ """Pricing per 1K tokens for a model."""
12
+ prompt_cost_per_1k: float
13
+ completion_cost_per_1k: float
14
+
15
+
16
+ # ── Pricing Table (cost per 1K tokens in USD) ───────
17
+ MODEL_PRICING: dict[str, ModelPricing] = {
18
+ # OpenAI
19
+ "gpt-4": ModelPricing(0.03, 0.06),
20
+ "gpt-4-turbo": ModelPricing(0.01, 0.03),
21
+ "gpt-4-turbo-preview": ModelPricing(0.01, 0.03),
22
+ "gpt-4o": ModelPricing(0.005, 0.015),
23
+ "gpt-4o-mini": ModelPricing(0.00015, 0.0006),
24
+ "gpt-3.5-turbo": ModelPricing(0.0005, 0.0015),
25
+ "gpt-3.5-turbo-16k": ModelPricing(0.003, 0.004),
26
+ # Anthropic
27
+ "claude-3-opus": ModelPricing(0.015, 0.075),
28
+ "claude-3-opus-20240229": ModelPricing(0.015, 0.075),
29
+ "claude-3-sonnet": ModelPricing(0.003, 0.015),
30
+ "claude-3-sonnet-20240229": ModelPricing(0.003, 0.015),
31
+ "claude-3-haiku": ModelPricing(0.00025, 0.00125),
32
+ "claude-3-haiku-20240307": ModelPricing(0.00025, 0.00125),
33
+ "claude-3.5-sonnet": ModelPricing(0.003, 0.015),
34
+ "claude-3.5-sonnet-20240620": ModelPricing(0.003, 0.015),
35
+ "claude-3.5-haiku": ModelPricing(0.001, 0.005),
36
+ # Google
37
+ "gemini-pro": ModelPricing(0.00025, 0.0005),
38
+ "gemini-1.5-pro": ModelPricing(0.00125, 0.005),
39
+ "gemini-1.5-flash": ModelPricing(0.000075, 0.0003),
40
+ "gemini-2.0-flash": ModelPricing(0.0001, 0.0004),
41
+ # Meta (via API providers)
42
+ "llama-3-70b": ModelPricing(0.00059, 0.00079),
43
+ "llama-3-8b": ModelPricing(0.00005, 0.00008),
44
+ "llama-3.1-405b": ModelPricing(0.003, 0.003),
45
+ # Mistral
46
+ "mistral-large": ModelPricing(0.004, 0.012),
47
+ "mistral-medium": ModelPricing(0.0027, 0.0081),
48
+ "mistral-small": ModelPricing(0.001, 0.003),
49
+ "mixtral-8x7b": ModelPricing(0.0007, 0.0007),
50
+ }
51
+
52
+
53
+ class CostCalculator:
54
+ """
55
+ Calculate the cost of LLM calls based on model and token usage.
56
+
57
+ Usage:
58
+ calc = CostCalculator()
59
+ cost = calc.calculate("gpt-4o", prompt_tokens=500, completion_tokens=200)
60
+ print(f"Cost: ${cost:.6f}")
61
+
62
+ # Add custom model pricing
63
+ calc.add_model("my-model", 0.001, 0.002)
64
+ """
65
+
66
+ def __init__(self):
67
+ self._pricing = dict(MODEL_PRICING)
68
+
69
+ def add_model(
70
+ self,
71
+ model_name: str,
72
+ prompt_cost_per_1k: float,
73
+ completion_cost_per_1k: float,
74
+ ):
75
+ """Add or update pricing for a model."""
76
+ self._pricing[model_name] = ModelPricing(
77
+ prompt_cost_per_1k, completion_cost_per_1k
78
+ )
79
+
80
+ def calculate(
81
+ self,
82
+ model: str,
83
+ prompt_tokens: int = 0,
84
+ completion_tokens: int = 0,
85
+ ) -> Optional[float]:
86
+ """
87
+ Calculate cost in USD for a given model and token usage.
88
+ Returns None if model pricing is not found.
89
+ """
90
+ pricing = self._pricing.get(model)
91
+ if pricing is None:
92
+ # Try partial match (e.g., "gpt-4o-2024-05-13" → "gpt-4o")
93
+ for key, val in self._pricing.items():
94
+ if model.startswith(key):
95
+ pricing = val
96
+ break
97
+
98
+ if pricing is None:
99
+ return None
100
+
101
+ prompt_cost = (prompt_tokens / 1000) * pricing.prompt_cost_per_1k
102
+ completion_cost = (completion_tokens / 1000) * pricing.completion_cost_per_1k
103
+ return prompt_cost + completion_cost
104
+
105
+ def get_pricing(self, model: str) -> Optional[ModelPricing]:
106
+ """Get pricing info for a model."""
107
+ return self._pricing.get(model)
108
+
109
+ @property
110
+ def supported_models(self) -> list[str]:
111
+ """List all models with known pricing."""
112
+ return sorted(self._pricing.keys())