agentcogs 0.1.1__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,51 @@
1
+ __pycache__/
2
+ *.py[cod]
3
+ *$py.class
4
+ *.egg-info/
5
+ .eggs/
6
+ dist/
7
+ build/
8
+ .venv/
9
+ venv/
10
+ .env
11
+ .env.*
12
+ !.env.example
13
+ !**/.env.example
14
+ .pytest_cache/
15
+ .mypy_cache/
16
+ .ruff_cache/
17
+ *.db
18
+ node_modules/
19
+ dashboard/dist/
20
+ backend/.venv/
21
+
22
+ # Local-only planning (do not push to GitHub)
23
+ TECH_ROADMAP.local.md
24
+ *.local.md
25
+ tools/.demo_env
26
+
27
+ # Maintainer / GTM — keep on disk, not in public repo
28
+ docs/internal/
29
+ docs/launch/
30
+ docs/DISTRIBUTION_PLAYBOOK.md
31
+ docs/GTM_FOUNDER_DIRECT.md
32
+ docs/CREDIBILITY_PLAYBOOK.md
33
+ docs/CUSTOMER_JOURNEY_STATUS.md
34
+ docs/SLACK_ALERTS_TECHNICAL_PLAN.md
35
+ docs/SDK_CUSTOMER_JOURNEY_PLAN.md
36
+ docs/DEMO_CALL_SCRIPT.md
37
+ .run/
38
+
39
+ # Playwright / local auth state (when added)
40
+ playwright/.auth/
41
+ playwright-report/
42
+ test-results/
43
+
44
+ # OS / IDE
45
+ .DS_Store
46
+ .idea/
47
+ .vscode/
48
+ *.swp
49
+
50
+ # Cursor IDE (local agent skills, rules — not part of the product)
51
+ .cursor/
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 AgentCOGS
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,149 @@
1
+ Metadata-Version: 2.5
2
+ Name: agentcogs
3
+ Version: 0.1.1
4
+ Summary: Per-customer LLM cost attribution and runtime budget caps for agent SaaS
5
+ Project-URL: Homepage, https://agentcogs.dev
6
+ Project-URL: Repository, https://github.com/vaibhav11123/agentcogs
7
+ Project-URL: Issues, https://github.com/vaibhav11123/agentcogs/issues
8
+ Author-email: AgentCOGS <hello@agentcogs.dev>
9
+ License: MIT
10
+ License-File: LICENSE
11
+ Keywords: anthropic,attribution,billing,cost,langchain,langgraph,llm,openai
12
+ Classifier: Development Status :: 4 - Beta
13
+ Classifier: Intended Audience :: Developers
14
+ Classifier: License :: OSI Approved :: MIT License
15
+ Classifier: Programming Language :: Python :: 3.10
16
+ Classifier: Programming Language :: Python :: 3.11
17
+ Classifier: Programming Language :: Python :: 3.12
18
+ Classifier: Topic :: Software Development :: Libraries
19
+ Requires-Python: >=3.10
20
+ Requires-Dist: httpx>=0.27.0
21
+ Requires-Dist: shekel[anthropic,openai]<2,>=1.2
22
+ Requires-Dist: starlette>=0.27.0
23
+ Provides-Extra: dev
24
+ Requires-Dist: pytest-asyncio>=0.23; extra == 'dev'
25
+ Requires-Dist: pytest>=8.0; extra == 'dev'
26
+ Requires-Dist: respx>=0.20; extra == 'dev'
27
+ Requires-Dist: ruff>=0.6; extra == 'dev'
28
+ Requires-Dist: starlette>=0.27.0; extra == 'dev'
29
+ Provides-Extra: integrations
30
+ Requires-Dist: starlette>=0.27.0; extra == 'integrations'
31
+ Description-Content-Type: text/markdown
32
+
33
+ # AgentCOGS
34
+
35
+ **Which customer stops making money this month.**
36
+
37
+ You sell an agent to other companies. The model provider sends one invoice. AgentCOGS splits that spend by the tenant who caused it, sets it against the monthly fee, and shows where the month is headed if the calls keep coming.
38
+
39
+ [![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](LICENSE)
40
+ [![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/)
41
+ [![CI](https://github.com/vaibhav11123/agentcogs/actions/workflows/ci.yml/badge.svg)](https://github.com/vaibhav11123/agentcogs/actions/workflows/ci.yml)
42
+
43
+ The sample company is [Patternstack](https://agentcogs.vercel.app). One click, nothing to type. It is a snapshot of a fictional book of business, shipped in the dashboard, so it opens when no API is running. A hosted AgentCOGS API is not serving. To run your own customers, self-host.
44
+
45
+ ---
46
+
47
+ ## The book, and the account inside it
48
+
49
+ Patternstack, as of 14 Nov 2026. Forty customers. The month so far looks healthy.
50
+
51
+ | | So far | If the calls continue |
52
+ | --- | ---: | ---: |
53
+ | AI cost | $39,613.25 | $84,885.54 |
54
+ | Fees | $155,918.00 | $155,918.00 |
55
+ | Margin | 74.6% | 45.6% |
56
+
57
+ The fee is the monthly amount on the account. Pace is a straight line from the cost so far. It does not forecast a change in behavior.
58
+
59
+ Two accounts are where the month breaks. Vellum Research has cost $5,822.84 against an $8,200 fee: **29.0%** so far, 2,724 runs, and the largest node is `synthesize`. If that spend continues, the month ends near **−52.2%**. Pinnacle Claims is 35.0% so far and ends near −39.3%. The blended 74.6% does not show either of those.
60
+
61
+ <p align="center">
62
+ <img src="docs/assets/screenshots/patternstack-customers.png" alt="Patternstack customers. Blended margin 74.6%. Vellum Research is 29.0% and the month ends near −52.2%." width="880" />
63
+ </p>
64
+
65
+ Open the account and the same numbers sit next to the call that posted them.
66
+
67
+ <p align="center">
68
+ <img src="docs/assets/screenshots/vellum-research.png" alt="Vellum Research. Margin 29.0%, AI cost $5,822.84, fee $8,200, and the set_customer call that posts each run." width="880" />
69
+ </p>
70
+
71
+ ```python
72
+ import agentcogs
73
+
74
+ agentcogs.init() # AGENTCOGS_API_KEY, AGENTCOGS_WORKSPACE_ID, AGENTCOGS_ENDPOINT
75
+ agentcogs.set_customer("vellum_research") # your tenant id, once per request
76
+
77
+ with agentcogs.run(workflow_id="research_agent"):
78
+ result = agent.invoke(state)
79
+ ```
80
+
81
+ The SDK does not call a model and does not invent the fee. It names the customer and posts the cost of the call you already made. You type the fee, or you pull the subscription from that customer's Stripe account. If no customer is set, the run stays unattributed.
82
+
83
+ `python examples/one_run.py` posts one recorded $0.84 cost for Vellum Research so you can see a row land. It does not call a model.
84
+
85
+ ---
86
+
87
+ ## Look, then run it on yours
88
+
89
+ **Look.** [agentcogs.vercel.app](https://agentcogs.vercel.app) → **Open Patternstack**.
90
+
91
+ **Run.** From this repo:
92
+
93
+ ```bash
94
+ cp .env.selfhost.example .env
95
+ docker compose up -d --build
96
+ ```
97
+
98
+ Open http://localhost:3000 and choose **Create my workspace**. No email. The next screen has the install command with your key in it. The package is not on PyPI yet:
99
+
100
+ ```bash
101
+ pip install "agentcogs @ git+https://github.com/vaibhav11123/agentcogs.git"
102
+ ```
103
+
104
+ Point `AGENTCOGS_ENDPOINT` at `http://localhost:8000`. Send one call, open that customer, and type the monthly fee. Margin is the fee minus the cost.
105
+
106
+ Full steps: [docs/quickstart.md](docs/quickstart.md). Self-host notes: [docs/DEPLOY.md](docs/DEPLOY.md).
107
+
108
+ ---
109
+
110
+ ## Who it is for
111
+
112
+ A founder of a product that sells one agent to many companies, priced as a monthly fee, with the model bill arriving as a single invoice. Support, research, claims, intake, document pipelines.
113
+
114
+ | You already have | AgentCOGS adds |
115
+ | --- | --- |
116
+ | Langfuse or LangSmith | Which tenant the trace belonged to, and whether that tenant is still profitable |
117
+ | LiteLLM | A callback that posts the tenant and the cost. The proxy stays |
118
+ | Stripe | The subscription as the fee, and an optional meter for the usage you already attributed |
119
+ | The provider invoice | The join from that one number to the customer list |
120
+
121
+ ---
122
+
123
+ ## What you can connect
124
+
125
+ | | |
126
+ | --- | --- |
127
+ | [FastAPI](docs/integrations/fastapi.md) | Set the tenant from the request |
128
+ | [LangGraph](docs/integrations/langgraph.md) | Wrap the graph |
129
+ | [LiteLLM](docs/integrations/litellm.md) | Post proxy traffic without wrapping the app |
130
+ | [Customer import](docs/integrations/customer-import.md) | Load names and fees before the first run |
131
+ | [Slack](docs/integrations/slack.md) | A spike opens that customer |
132
+ | Stripe | Pull subscription fees. Push a daily usage meter |
133
+ | Budget | `monthly_budget_usd` raises `CustomerBudgetExceededError` before the next call |
134
+
135
+ Terminal proof for the SDK, the outbox, and a live call: [docs/demo-paths.md](docs/demo-paths.md).
136
+
137
+ ---
138
+
139
+ ## Repo
140
+
141
+ | Path | |
142
+ | --- | --- |
143
+ | `src/agentcogs/` | The SDK |
144
+ | `backend/` | Ingest, margin, pace, budgets |
145
+ | `dashboard/` | The book |
146
+ | `examples/one_run.py` | One recorded cost, no model call |
147
+ | `docs/` | [Index](docs/README.md) |
148
+
149
+ Python SDK, FastAPI, Postgres, Redis, React. [MIT](LICENSE). [Contributing](CONTRIBUTING.md). [Security](SECURITY.md).
@@ -0,0 +1,117 @@
1
+ # AgentCOGS
2
+
3
+ **Which customer stops making money this month.**
4
+
5
+ You sell an agent to other companies. The model provider sends one invoice. AgentCOGS splits that spend by the tenant who caused it, sets it against the monthly fee, and shows where the month is headed if the calls keep coming.
6
+
7
+ [![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](LICENSE)
8
+ [![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/)
9
+ [![CI](https://github.com/vaibhav11123/agentcogs/actions/workflows/ci.yml/badge.svg)](https://github.com/vaibhav11123/agentcogs/actions/workflows/ci.yml)
10
+
11
+ The sample company is [Patternstack](https://agentcogs.vercel.app). One click, nothing to type. It is a snapshot of a fictional book of business, shipped in the dashboard, so it opens when no API is running. A hosted AgentCOGS API is not serving. To run your own customers, self-host.
12
+
13
+ ---
14
+
15
+ ## The book, and the account inside it
16
+
17
+ Patternstack, as of 14 Nov 2026. Forty customers. The month so far looks healthy.
18
+
19
+ | | So far | If the calls continue |
20
+ | --- | ---: | ---: |
21
+ | AI cost | $39,613.25 | $84,885.54 |
22
+ | Fees | $155,918.00 | $155,918.00 |
23
+ | Margin | 74.6% | 45.6% |
24
+
25
+ The fee is the monthly amount on the account. Pace is a straight line from the cost so far. It does not forecast a change in behavior.
26
+
27
+ Two accounts are where the month breaks. Vellum Research has cost $5,822.84 against an $8,200 fee: **29.0%** so far, 2,724 runs, and the largest node is `synthesize`. If that spend continues, the month ends near **−52.2%**. Pinnacle Claims is 35.0% so far and ends near −39.3%. The blended 74.6% does not show either of those.
28
+
29
+ <p align="center">
30
+ <img src="docs/assets/screenshots/patternstack-customers.png" alt="Patternstack customers. Blended margin 74.6%. Vellum Research is 29.0% and the month ends near −52.2%." width="880" />
31
+ </p>
32
+
33
+ Open the account and the same numbers sit next to the call that posted them.
34
+
35
+ <p align="center">
36
+ <img src="docs/assets/screenshots/vellum-research.png" alt="Vellum Research. Margin 29.0%, AI cost $5,822.84, fee $8,200, and the set_customer call that posts each run." width="880" />
37
+ </p>
38
+
39
+ ```python
40
+ import agentcogs
41
+
42
+ agentcogs.init() # AGENTCOGS_API_KEY, AGENTCOGS_WORKSPACE_ID, AGENTCOGS_ENDPOINT
43
+ agentcogs.set_customer("vellum_research") # your tenant id, once per request
44
+
45
+ with agentcogs.run(workflow_id="research_agent"):
46
+ result = agent.invoke(state)
47
+ ```
48
+
49
+ The SDK does not call a model and does not invent the fee. It names the customer and posts the cost of the call you already made. You type the fee, or you pull the subscription from that customer's Stripe account. If no customer is set, the run stays unattributed.
50
+
51
+ `python examples/one_run.py` posts one recorded $0.84 cost for Vellum Research so you can see a row land. It does not call a model.
52
+
53
+ ---
54
+
55
+ ## Look, then run it on yours
56
+
57
+ **Look.** [agentcogs.vercel.app](https://agentcogs.vercel.app) → **Open Patternstack**.
58
+
59
+ **Run.** From this repo:
60
+
61
+ ```bash
62
+ cp .env.selfhost.example .env
63
+ docker compose up -d --build
64
+ ```
65
+
66
+ Open http://localhost:3000 and choose **Create my workspace**. No email. The next screen has the install command with your key in it. The package is not on PyPI yet:
67
+
68
+ ```bash
69
+ pip install "agentcogs @ git+https://github.com/vaibhav11123/agentcogs.git"
70
+ ```
71
+
72
+ Point `AGENTCOGS_ENDPOINT` at `http://localhost:8000`. Send one call, open that customer, and type the monthly fee. Margin is the fee minus the cost.
73
+
74
+ Full steps: [docs/quickstart.md](docs/quickstart.md). Self-host notes: [docs/DEPLOY.md](docs/DEPLOY.md).
75
+
76
+ ---
77
+
78
+ ## Who it is for
79
+
80
+ A founder of a product that sells one agent to many companies, priced as a monthly fee, with the model bill arriving as a single invoice. Support, research, claims, intake, document pipelines.
81
+
82
+ | You already have | AgentCOGS adds |
83
+ | --- | --- |
84
+ | Langfuse or LangSmith | Which tenant the trace belonged to, and whether that tenant is still profitable |
85
+ | LiteLLM | A callback that posts the tenant and the cost. The proxy stays |
86
+ | Stripe | The subscription as the fee, and an optional meter for the usage you already attributed |
87
+ | The provider invoice | The join from that one number to the customer list |
88
+
89
+ ---
90
+
91
+ ## What you can connect
92
+
93
+ | | |
94
+ | --- | --- |
95
+ | [FastAPI](docs/integrations/fastapi.md) | Set the tenant from the request |
96
+ | [LangGraph](docs/integrations/langgraph.md) | Wrap the graph |
97
+ | [LiteLLM](docs/integrations/litellm.md) | Post proxy traffic without wrapping the app |
98
+ | [Customer import](docs/integrations/customer-import.md) | Load names and fees before the first run |
99
+ | [Slack](docs/integrations/slack.md) | A spike opens that customer |
100
+ | Stripe | Pull subscription fees. Push a daily usage meter |
101
+ | Budget | `monthly_budget_usd` raises `CustomerBudgetExceededError` before the next call |
102
+
103
+ Terminal proof for the SDK, the outbox, and a live call: [docs/demo-paths.md](docs/demo-paths.md).
104
+
105
+ ---
106
+
107
+ ## Repo
108
+
109
+ | Path | |
110
+ | --- | --- |
111
+ | `src/agentcogs/` | The SDK |
112
+ | `backend/` | Ingest, margin, pace, budgets |
113
+ | `dashboard/` | The book |
114
+ | `examples/one_run.py` | One recorded cost, no model call |
115
+ | `docs/` | [Index](docs/README.md) |
116
+
117
+ Python SDK, FastAPI, Postgres, Redis, React. [MIT](LICENSE). [Contributing](CONTRIBUTING.md). [Security](SECURITY.md).
@@ -0,0 +1,55 @@
1
+ [build-system]
2
+ requires = ["hatchling>=1.21"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "agentcogs"
7
+ dynamic = ["version"]
8
+ description = "Per-customer LLM cost attribution and runtime budget caps for agent SaaS"
9
+ readme = "README.md"
10
+ license = { text = "MIT" }
11
+ requires-python = ">=3.10"
12
+ authors = [{ name = "AgentCOGS", email = "hello@agentcogs.dev" }]
13
+ keywords = ["llm", "cost", "attribution", "langgraph", "langchain", "openai", "anthropic", "billing"]
14
+ classifiers = [
15
+ "Development Status :: 4 - Beta",
16
+ "Intended Audience :: Developers",
17
+ "License :: OSI Approved :: MIT License",
18
+ "Programming Language :: Python :: 3.10",
19
+ "Programming Language :: Python :: 3.11",
20
+ "Programming Language :: Python :: 3.12",
21
+ "Topic :: Software Development :: Libraries",
22
+ ]
23
+ dependencies = [
24
+ "shekel[openai,anthropic]>=1.2,<2",
25
+ "httpx>=0.27.0",
26
+ "starlette>=0.27.0",
27
+ ]
28
+
29
+ [project.optional-dependencies]
30
+ dev = ["pytest>=8.0", "pytest-asyncio>=0.23", "respx>=0.20", "ruff>=0.6", "starlette>=0.27.0"]
31
+ integrations = ["starlette>=0.27.0"]
32
+
33
+ [project.urls]
34
+ Homepage = "https://agentcogs.dev"
35
+ Repository = "https://github.com/vaibhav11123/agentcogs"
36
+ Issues = "https://github.com/vaibhav11123/agentcogs/issues"
37
+
38
+ [tool.hatch.version]
39
+ path = "src/agentcogs/_version.py"
40
+
41
+ [tool.hatch.build.targets.wheel]
42
+ packages = ["src/agentcogs"]
43
+
44
+ # pip install prefers the wheel. The sdist is the fallback and must stay the
45
+ # SDK. The backend, dashboard, and demo tools are the rest of this repo.
46
+ [tool.hatch.build.targets.sdist]
47
+ only-include = ["src/agentcogs", "README.md", "LICENSE"]
48
+
49
+ [tool.ruff]
50
+ line-length = 100
51
+ target-version = "py310"
52
+
53
+ [tool.pytest.ini_options]
54
+ asyncio_mode = "auto"
55
+ testpaths = ["tests"]
@@ -0,0 +1,43 @@
1
+ """AgentCOGS — per-customer LLM cost attribution.
2
+
3
+ Quickstart:
4
+ import agentcogs
5
+ agentcogs.init(api_key="acg_live_...", workspace_id="...")
6
+
7
+ agentcogs.set_customer("tenant_42") # once per request
8
+ with agentcogs.run(workflow_id="summarize"):
9
+ # any openai / anthropic / langgraph code here
10
+ ...
11
+ """
12
+ from ._version import __version__
13
+ from .budget import IngestStatus, RunContext, run
14
+ from .config import init
15
+ from .context import get_customer, get_workflow, set_customer, set_workflow
16
+ from .errors import (
17
+ AgentCOGSError,
18
+ ConfigurationError,
19
+ CustomerBudgetExceededError,
20
+ PingError,
21
+ )
22
+ from .ping import PingResult, ping
23
+ from .shutdown import ShutdownResult, shutdown
24
+
25
+ __all__ = [
26
+ "__version__",
27
+ "init",
28
+ "run",
29
+ "ping",
30
+ "shutdown",
31
+ "set_customer",
32
+ "set_workflow",
33
+ "get_customer",
34
+ "get_workflow",
35
+ "RunContext",
36
+ "IngestStatus",
37
+ "PingResult",
38
+ "ShutdownResult",
39
+ "AgentCOGSError",
40
+ "CustomerBudgetExceededError",
41
+ "ConfigurationError",
42
+ "PingError",
43
+ ]
@@ -0,0 +1,17 @@
1
+ """CLI: python -m agentcogs outbox status"""
2
+ import json
3
+ import sys
4
+
5
+ from .outbox import get_status
6
+
7
+
8
+ def main() -> None:
9
+ if len(sys.argv) >= 3 and sys.argv[1] == "outbox" and sys.argv[2] == "status":
10
+ print(json.dumps(get_status(), indent=2))
11
+ return
12
+ print("Usage: python -m agentcogs outbox status", file=sys.stderr)
13
+ sys.exit(1)
14
+
15
+
16
+ if __name__ == "__main__":
17
+ main()
@@ -0,0 +1,48 @@
1
+ """Thin adapter over shekel 1.x — single place for version-specific field mapping."""
2
+ from __future__ import annotations
3
+
4
+ from typing import Any, Dict
5
+
6
+ from shekel import budget as _shekel_budget
7
+
8
+
9
+ def budget(**kwargs: Any):
10
+ """Return shekel budget context manager (signature-stable for agentcogs)."""
11
+ return _shekel_budget(**kwargs)
12
+
13
+
14
+ def normalize_summary(by_model: Dict[str, Dict[str, Any]]) -> Dict[str, Dict[str, Any]]:
15
+ """Convert Shekel's by_model summary into a uniform schema.
16
+
17
+ Anthropic prompt caching splits input tokens across three fields; all must be summed.
18
+ """
19
+ out: Dict[str, Dict[str, Any]] = {}
20
+ for model, stats in by_model.items():
21
+ model_lower = model.lower()
22
+ if "claude" in model_lower or "anthropic" in model_lower:
23
+ input_total = (
24
+ int(stats.get("input_tokens", 0) or 0)
25
+ + int(stats.get("cache_read_input_tokens", 0) or 0)
26
+ + int(stats.get("cache_creation_input_tokens", 0) or 0)
27
+ )
28
+ else:
29
+ input_total = int(stats.get("input_tokens", 0) or 0)
30
+
31
+ out[model] = {
32
+ "input_tokens": input_total,
33
+ "output_tokens": int(stats.get("output_tokens", 0) or 0),
34
+ "usd": float(stats.get("cost", 0) or 0),
35
+ }
36
+ return out
37
+
38
+
39
+ def normalize_summary_data(raw: dict) -> dict:
40
+ """Normalize summary_data() from a shekel budget context."""
41
+ total = raw.get("total_cost")
42
+ if total is None:
43
+ total = raw.get("total_spent", 0)
44
+ by_model = raw.get("by_model") or {}
45
+ return {
46
+ "total_cost": float(total or 0),
47
+ "by_model": by_model,
48
+ }
@@ -0,0 +1 @@
1
+ __version__ = "0.1.1"
@@ -0,0 +1,161 @@
1
+ """The agentcogs.run() context manager — the SDK's entire public surface."""
2
+ import json
3
+ import logging
4
+ import time
5
+ import uuid
6
+ from contextlib import contextmanager
7
+ from dataclasses import dataclass
8
+ from typing import Dict, Iterator, Optional
9
+
10
+ from ._shekel_compat import budget as shekel_budget, normalize_summary, normalize_summary_data
11
+ from .client import emit_event, emit_event_sync, fetch_budget
12
+ from .config import get_config
13
+ from .context import resolve_customer, resolve_workflow
14
+ from .errors import ConfigurationError, CustomerBudgetExceededError
15
+
16
+ log = logging.getLogger("agentcogs")
17
+
18
+ _MAX_METADATA_KEYS = 100
19
+ _MAX_METADATA_BYTES = 16_384
20
+
21
+
22
+ def _cap_metadata(metadata: Optional[dict]) -> dict:
23
+ if not metadata:
24
+ return {}
25
+ capped = dict(list(metadata.items())[:_MAX_METADATA_KEYS])
26
+ if len(capped) < len(metadata):
27
+ log.debug("metadata truncated to %d keys", _MAX_METADATA_KEYS)
28
+ while capped and len(json.dumps(capped, default=str)) > _MAX_METADATA_BYTES:
29
+ capped.pop(next(iter(capped)))
30
+ if len(capped) < len(metadata):
31
+ log.debug("metadata size capped to %d bytes", _MAX_METADATA_BYTES)
32
+ return capped
33
+
34
+
35
+ @dataclass
36
+ class IngestStatus:
37
+ accepted: bool
38
+ duplicate: bool = False
39
+ error: Optional[str] = None
40
+
41
+
42
+ class RunContext:
43
+ """Handle returned from agentcogs.run()."""
44
+
45
+ def __init__(
46
+ self,
47
+ shekel_ctx: object,
48
+ run_id: str,
49
+ customer_id: str,
50
+ workflow_id: str,
51
+ ):
52
+ self._shekel = shekel_ctx
53
+ self.run_id = run_id
54
+ self.customer_id = customer_id
55
+ self.workflow_id = workflow_id
56
+ self._event: Optional[dict] = None
57
+
58
+ def summary_data(self) -> dict:
59
+ return self._shekel.summary_data() # type: ignore[attr-defined]
60
+
61
+ def wait_for_ingest(self, timeout: float = 5.0) -> IngestStatus:
62
+ """Block until event is POSTed (onboarding/tests). Uses sync ingest."""
63
+ if self._event is None:
64
+ return IngestStatus(accepted=False, error="no event built yet")
65
+ ok, err = emit_event_sync(self._event)
66
+ return IngestStatus(accepted=ok, error=err)
67
+
68
+ def __getattr__(self, name: str):
69
+ return getattr(self._shekel, name)
70
+
71
+
72
+ @contextmanager
73
+ def run(
74
+ customer_id: Optional[str] = None,
75
+ workflow_id: Optional[str] = None,
76
+ workspace_id: Optional[str] = None,
77
+ metadata: Optional[dict] = None,
78
+ price_per_1k_tokens: Optional[Dict[str, float]] = None,
79
+ ) -> Iterator[RunContext]:
80
+ """Wrap any block of LLM-using code with per-customer cost tracking.
81
+
82
+ customer_id may be omitted if agentcogs.set_customer() was called for this context.
83
+ workflow_id defaults to set_workflow() context value, else "default".
84
+ """
85
+ cfg = get_config()
86
+ ws_id = workspace_id or cfg.workspace_id
87
+ if not ws_id and not cfg.offline:
88
+ raise ConfigurationError(
89
+ "workspace_id required: pass workspace_id= or set in init()"
90
+ )
91
+
92
+ resolved_customer = resolve_customer(customer_id)
93
+ if not resolved_customer:
94
+ raise ConfigurationError(
95
+ "customer_id required: pass customer_id= to run() or call "
96
+ "set_customer() in middleware before agent code. "
97
+ "See docs/concepts/customer-id.md"
98
+ )
99
+ resolved_workflow = resolve_workflow(workflow_id)
100
+
101
+ run_id = str(uuid.uuid4())
102
+
103
+ b = fetch_budget(ws_id or "offline", resolved_customer)
104
+ if b["status"] == "exceeded":
105
+ raise CustomerBudgetExceededError(
106
+ customer_id=resolved_customer,
107
+ spent_usd=b["spent_usd"],
108
+ budget_usd=b["budget_usd"],
109
+ )
110
+
111
+ remaining = b.get("remaining_usd")
112
+ max_usd = float(remaining) if remaining and remaining != float("inf") else None
113
+
114
+ status = "completed"
115
+ error_msg: Optional[str] = None
116
+
117
+ kwargs: dict = {"name": resolved_customer}
118
+ if max_usd is not None:
119
+ kwargs["max_usd"] = max_usd
120
+ if price_per_1k_tokens is not None:
121
+ kwargs["price_per_1k_tokens"] = price_per_1k_tokens
122
+
123
+ run_ctx: Optional[RunContext] = None
124
+
125
+ with shekel_budget(**kwargs) as ctx:
126
+ try:
127
+ run_ctx = RunContext(ctx, run_id, resolved_customer, resolved_workflow)
128
+ yield run_ctx
129
+ except CustomerBudgetExceededError:
130
+ status = "budget_exceeded"
131
+ raise
132
+ except Exception as e:
133
+ status = "error"
134
+ error_msg = str(e)[:500]
135
+ raise
136
+ finally:
137
+ try:
138
+ summary = normalize_summary_data(ctx.summary_data())
139
+ event = {
140
+ "run_id": run_id,
141
+ "customer_id": resolved_customer,
142
+ "workflow_id": resolved_workflow,
143
+ "ts": int(time.time()),
144
+ "status": status,
145
+ "total_usd": float(summary.get("total_cost") or 0),
146
+ "models": normalize_summary(summary.get("by_model", {})),
147
+ "node_costs": getattr(ctx, "node_costs", {}),
148
+ "metadata": _cap_metadata(metadata),
149
+ "error": error_msg,
150
+ }
151
+ if run_ctx is not None:
152
+ run_ctx._event = event
153
+ emit_event(event)
154
+ except Exception as e:
155
+ log.warning("telemetry failed", exc_info=True)
156
+ cb = cfg.on_telemetry_error
157
+ if cb is not None:
158
+ try:
159
+ cb(e)
160
+ except Exception:
161
+ pass