agentcogs 0.1.1__py3-none-any.whl
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.
- agentcogs/__init__.py +43 -0
- agentcogs/__main__.py +17 -0
- agentcogs/_shekel_compat.py +48 -0
- agentcogs/_version.py +1 -0
- agentcogs/budget.py +161 -0
- agentcogs/client.py +145 -0
- agentcogs/config.py +90 -0
- agentcogs/context.py +36 -0
- agentcogs/errors.py +32 -0
- agentcogs/integrations/__init__.py +1 -0
- agentcogs/integrations/fastapi.py +44 -0
- agentcogs/integrations/langgraph.py +38 -0
- agentcogs/integrations/litellm_event.py +67 -0
- agentcogs/outbox.py +121 -0
- agentcogs/ping.py +78 -0
- agentcogs/shutdown.py +38 -0
- agentcogs/tokens.py +4 -0
- agentcogs-0.1.1.dist-info/METADATA +149 -0
- agentcogs-0.1.1.dist-info/RECORD +21 -0
- agentcogs-0.1.1.dist-info/WHEEL +4 -0
- agentcogs-0.1.1.dist-info/licenses/LICENSE +21 -0
agentcogs/__init__.py
ADDED
|
@@ -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
|
+
]
|
agentcogs/__main__.py
ADDED
|
@@ -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
|
+
}
|
agentcogs/_version.py
ADDED
|
@@ -0,0 +1 @@
|
|
|
1
|
+
__version__ = "0.1.1"
|
agentcogs/budget.py
ADDED
|
@@ -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
|
agentcogs/client.py
ADDED
|
@@ -0,0 +1,145 @@
|
|
|
1
|
+
"""HTTP client. Fire-and-forget POST to ingest endpoint.
|
|
2
|
+
|
|
3
|
+
Never blocks user code. Failures route to outbox for retry.
|
|
4
|
+
"""
|
|
5
|
+
import logging
|
|
6
|
+
from concurrent.futures import ThreadPoolExecutor
|
|
7
|
+
from typing import Any, Dict, Optional, Tuple
|
|
8
|
+
|
|
9
|
+
import httpx
|
|
10
|
+
|
|
11
|
+
from ._version import __version__
|
|
12
|
+
from .config import get_config
|
|
13
|
+
from .errors import AgentCOGSError
|
|
14
|
+
from .outbox import drain, enqueue
|
|
15
|
+
|
|
16
|
+
log = logging.getLogger("agentcogs")
|
|
17
|
+
|
|
18
|
+
_client: Optional[httpx.Client] = None
|
|
19
|
+
_executor: Optional[ThreadPoolExecutor] = None
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
def _user_agent() -> str:
|
|
23
|
+
return f"agentcogs-python/{__version__}"
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
def _default_headers() -> Dict[str, str]:
|
|
27
|
+
return {
|
|
28
|
+
"Authorization": f"Bearer {get_config().api_key}",
|
|
29
|
+
"User-Agent": _user_agent(),
|
|
30
|
+
"Content-Type": "application/json",
|
|
31
|
+
"X-AgentCOGS-SDK-Version": __version__,
|
|
32
|
+
}
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
def reset_client() -> None:
|
|
36
|
+
"""Close httpx client (e.g. on re-init)."""
|
|
37
|
+
global _client
|
|
38
|
+
if _client is not None:
|
|
39
|
+
try:
|
|
40
|
+
_client.close()
|
|
41
|
+
except Exception:
|
|
42
|
+
pass
|
|
43
|
+
_client = None
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
def _get_client() -> httpx.Client:
|
|
47
|
+
global _client
|
|
48
|
+
if _client is None:
|
|
49
|
+
cfg = get_config()
|
|
50
|
+
_client = httpx.Client(
|
|
51
|
+
base_url=cfg.endpoint,
|
|
52
|
+
headers=_default_headers(),
|
|
53
|
+
timeout=cfg.timeout_seconds,
|
|
54
|
+
)
|
|
55
|
+
return _client
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
def _get_executor() -> ThreadPoolExecutor:
|
|
59
|
+
global _executor
|
|
60
|
+
if _executor is None:
|
|
61
|
+
cfg = get_config()
|
|
62
|
+
_executor = ThreadPoolExecutor(
|
|
63
|
+
max_workers=cfg.max_ingest_workers,
|
|
64
|
+
thread_name_prefix="agentcogs-ingest",
|
|
65
|
+
)
|
|
66
|
+
return _executor
|
|
67
|
+
|
|
68
|
+
|
|
69
|
+
def shutdown_executor(wait: bool = True, timeout: float = 10.0) -> None:
|
|
70
|
+
global _executor
|
|
71
|
+
if _executor is not None:
|
|
72
|
+
_executor.shutdown(wait=wait, cancel_futures=not wait)
|
|
73
|
+
_executor = None
|
|
74
|
+
|
|
75
|
+
|
|
76
|
+
def _post(event: Dict[str, Any]) -> None:
|
|
77
|
+
"""Synchronous POST. Raises on any failure."""
|
|
78
|
+
cfg = get_config()
|
|
79
|
+
if cfg.offline:
|
|
80
|
+
log.debug("offline mode — event %s skipped", event["run_id"])
|
|
81
|
+
return
|
|
82
|
+
resp = _get_client().post("/v1/ingest", json=event)
|
|
83
|
+
resp.raise_for_status()
|
|
84
|
+
|
|
85
|
+
|
|
86
|
+
def emit_event_sync(event: Dict[str, Any]) -> Tuple[bool, Optional[str]]:
|
|
87
|
+
"""Synchronous ingest. Returns (accepted, error_message)."""
|
|
88
|
+
try:
|
|
89
|
+
drain(_post, max_items=20)
|
|
90
|
+
except Exception as e:
|
|
91
|
+
log.debug("outbox drain failed: %s", e)
|
|
92
|
+
try:
|
|
93
|
+
_post(event)
|
|
94
|
+
return True, None
|
|
95
|
+
except Exception as e:
|
|
96
|
+
enqueue(event)
|
|
97
|
+
return False, str(e)
|
|
98
|
+
|
|
99
|
+
|
|
100
|
+
def emit_event(event: Dict[str, Any]) -> None:
|
|
101
|
+
"""Best-effort delivery via thread pool. Always returns quickly."""
|
|
102
|
+
|
|
103
|
+
def _work() -> None:
|
|
104
|
+
try:
|
|
105
|
+
drain(_post, max_items=20)
|
|
106
|
+
except Exception as e:
|
|
107
|
+
log.debug("outbox drain failed: %s", e)
|
|
108
|
+
try:
|
|
109
|
+
_post(event)
|
|
110
|
+
except Exception as e:
|
|
111
|
+
log.warning("ingest failed (%s) — queued for retry", e)
|
|
112
|
+
enqueue(event)
|
|
113
|
+
|
|
114
|
+
_get_executor().submit(_work)
|
|
115
|
+
|
|
116
|
+
|
|
117
|
+
def fetch_budget(workspace_id: str, customer_id: str) -> Dict[str, Any]:
|
|
118
|
+
"""Pre-flight budget check. Fail-open unless budget_mode=closed."""
|
|
119
|
+
cfg = get_config()
|
|
120
|
+
if cfg.offline:
|
|
121
|
+
return {
|
|
122
|
+
"status": "ok",
|
|
123
|
+
"spent_usd": 0.0,
|
|
124
|
+
"budget_usd": None,
|
|
125
|
+
"remaining_usd": float("inf"),
|
|
126
|
+
}
|
|
127
|
+
try:
|
|
128
|
+
resp = _get_client().get(
|
|
129
|
+
"/v1/budget",
|
|
130
|
+
params={"workspace": workspace_id, "customer": customer_id},
|
|
131
|
+
)
|
|
132
|
+
resp.raise_for_status()
|
|
133
|
+
return resp.json()
|
|
134
|
+
except Exception as e:
|
|
135
|
+
if cfg.budget_mode == "closed":
|
|
136
|
+
raise AgentCOGSError(
|
|
137
|
+
f"Budget check failed and budget_mode=closed: {e}"
|
|
138
|
+
) from e
|
|
139
|
+
log.warning("budget fetch failed (%s) — failing open", e)
|
|
140
|
+
return {
|
|
141
|
+
"status": "ok",
|
|
142
|
+
"spent_usd": 0.0,
|
|
143
|
+
"budget_usd": None,
|
|
144
|
+
"remaining_usd": float("inf"),
|
|
145
|
+
}
|
agentcogs/config.py
ADDED
|
@@ -0,0 +1,90 @@
|
|
|
1
|
+
import atexit
|
|
2
|
+
import os
|
|
3
|
+
from dataclasses import dataclass, field
|
|
4
|
+
from typing import Callable, Literal, Optional
|
|
5
|
+
|
|
6
|
+
from .errors import ConfigurationError
|
|
7
|
+
|
|
8
|
+
BudgetMode = Literal["open", "closed"]
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
@dataclass
|
|
12
|
+
class _Config:
|
|
13
|
+
api_key: Optional[str] = None
|
|
14
|
+
workspace_id: Optional[str] = None
|
|
15
|
+
endpoint: str = "https://api.agentcogs.dev"
|
|
16
|
+
timeout_seconds: float = 2.0
|
|
17
|
+
offline: bool = False
|
|
18
|
+
budget_mode: BudgetMode = "open"
|
|
19
|
+
max_ingest_workers: int = 4
|
|
20
|
+
register_atexit: bool = True
|
|
21
|
+
on_telemetry_error: Optional[Callable[[Exception], None]] = None
|
|
22
|
+
_atexit_registered: bool = field(default=False, repr=False)
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
_config = _Config()
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
def init(
|
|
29
|
+
api_key: Optional[str] = None,
|
|
30
|
+
workspace_id: Optional[str] = None,
|
|
31
|
+
endpoint: Optional[str] = None,
|
|
32
|
+
offline: bool = False,
|
|
33
|
+
timeout_seconds: float = 2.0,
|
|
34
|
+
budget_mode: BudgetMode = "open",
|
|
35
|
+
max_ingest_workers: int = 4,
|
|
36
|
+
register_atexit: bool = True,
|
|
37
|
+
on_telemetry_error: Optional[Callable[[Exception], None]] = None,
|
|
38
|
+
) -> None:
|
|
39
|
+
"""Initialise the AgentCOGS SDK.
|
|
40
|
+
|
|
41
|
+
Reads from env vars as fallback:
|
|
42
|
+
AGENTCOGS_API_KEY, AGENTCOGS_WORKSPACE_ID, AGENTCOGS_ENDPOINT,
|
|
43
|
+
AGENTCOGS_OFFLINE, AGENTCOGS_BUDGET_MODE (open|closed)
|
|
44
|
+
"""
|
|
45
|
+
from . import client as _client_mod
|
|
46
|
+
|
|
47
|
+
_client_mod.reset_client()
|
|
48
|
+
|
|
49
|
+
_config.api_key = api_key or os.environ.get("AGENTCOGS_API_KEY")
|
|
50
|
+
_config.workspace_id = workspace_id or os.environ.get("AGENTCOGS_WORKSPACE_ID")
|
|
51
|
+
_config.endpoint = (
|
|
52
|
+
endpoint
|
|
53
|
+
or os.environ.get("AGENTCOGS_ENDPOINT")
|
|
54
|
+
or "https://api.agentcogs.dev"
|
|
55
|
+
).rstrip("/")
|
|
56
|
+
_config.offline = offline or os.environ.get("AGENTCOGS_OFFLINE") == "1"
|
|
57
|
+
_config.timeout_seconds = timeout_seconds
|
|
58
|
+
env_budget = os.environ.get("AGENTCOGS_BUDGET_MODE", "").lower()
|
|
59
|
+
if env_budget in ("open", "closed"):
|
|
60
|
+
_config.budget_mode = env_budget # type: ignore[assignment]
|
|
61
|
+
else:
|
|
62
|
+
_config.budget_mode = budget_mode
|
|
63
|
+
_config.max_ingest_workers = max(1, max_ingest_workers)
|
|
64
|
+
_config.register_atexit = register_atexit
|
|
65
|
+
_config.on_telemetry_error = on_telemetry_error
|
|
66
|
+
|
|
67
|
+
if not _config.offline and not _config.api_key:
|
|
68
|
+
raise ConfigurationError(
|
|
69
|
+
"agentcogs.init() requires api_key (or AGENTCOGS_API_KEY env). "
|
|
70
|
+
"For local dev without backend, pass offline=True."
|
|
71
|
+
)
|
|
72
|
+
|
|
73
|
+
if register_atexit and not _config._atexit_registered:
|
|
74
|
+
atexit.register(_shutdown_on_exit)
|
|
75
|
+
_config._atexit_registered = True
|
|
76
|
+
|
|
77
|
+
|
|
78
|
+
def get_config() -> _Config:
|
|
79
|
+
if _config.api_key is None and not _config.offline:
|
|
80
|
+
init()
|
|
81
|
+
return _config
|
|
82
|
+
|
|
83
|
+
|
|
84
|
+
def _shutdown_on_exit() -> None:
|
|
85
|
+
try:
|
|
86
|
+
from .shutdown import shutdown
|
|
87
|
+
|
|
88
|
+
shutdown(drain_timeout=5.0)
|
|
89
|
+
except Exception:
|
|
90
|
+
pass
|
agentcogs/context.py
ADDED
|
@@ -0,0 +1,36 @@
|
|
|
1
|
+
"""Request-scoped customer/workflow attribution via contextvars."""
|
|
2
|
+
from contextvars import ContextVar
|
|
3
|
+
from typing import Optional
|
|
4
|
+
|
|
5
|
+
_customer_id: ContextVar[Optional[str]] = ContextVar("agentcogs_customer_id", default=None)
|
|
6
|
+
_workflow_id: ContextVar[Optional[str]] = ContextVar("agentcogs_workflow_id", default=None)
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
def set_customer(customer_id: Optional[str]) -> None:
|
|
10
|
+
"""Set billing customer for the current async/thread context."""
|
|
11
|
+
_customer_id.set(customer_id)
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
def set_workflow(workflow_id: Optional[str]) -> None:
|
|
15
|
+
"""Set default workflow id for subsequent run() calls in this context."""
|
|
16
|
+
_workflow_id.set(workflow_id)
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
def get_customer() -> Optional[str]:
|
|
20
|
+
return _customer_id.get()
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
def get_workflow() -> Optional[str]:
|
|
24
|
+
return _workflow_id.get()
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
def resolve_customer(explicit: Optional[str]) -> Optional[str]:
|
|
28
|
+
if explicit is not None:
|
|
29
|
+
return explicit
|
|
30
|
+
return get_customer()
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
def resolve_workflow(explicit: Optional[str], default: str = "default") -> str:
|
|
34
|
+
if explicit is not None:
|
|
35
|
+
return explicit
|
|
36
|
+
return get_workflow() or default
|
agentcogs/errors.py
ADDED
|
@@ -0,0 +1,32 @@
|
|
|
1
|
+
class AgentCOGSError(Exception):
|
|
2
|
+
"""Base exception for all AgentCOGS errors."""
|
|
3
|
+
|
|
4
|
+
|
|
5
|
+
class ConfigurationError(AgentCOGSError):
|
|
6
|
+
"""Raised when SDK is used without init() or with bad config."""
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
class PingError(AgentCOGSError):
|
|
10
|
+
"""Raised when agentcogs.ping() cannot reach or authenticate to the API."""
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
class CustomerBudgetExceededError(AgentCOGSError):
|
|
14
|
+
"""Raised when a customer's monthly budget cap is reached.
|
|
15
|
+
|
|
16
|
+
Raised on context entry BEFORE any LLM call fires, so zero
|
|
17
|
+
provider charges are incurred.
|
|
18
|
+
|
|
19
|
+
Attributes:
|
|
20
|
+
customer_id: External customer identifier
|
|
21
|
+
spent_usd: Current month-to-date spend
|
|
22
|
+
budget_usd: Configured monthly cap
|
|
23
|
+
"""
|
|
24
|
+
|
|
25
|
+
def __init__(self, customer_id: str, spent_usd: float, budget_usd: float):
|
|
26
|
+
self.customer_id = customer_id
|
|
27
|
+
self.spent_usd = spent_usd
|
|
28
|
+
self.budget_usd = budget_usd
|
|
29
|
+
super().__init__(
|
|
30
|
+
f"Customer '{customer_id}' exceeded budget: "
|
|
31
|
+
f"${spent_usd:.2f} / ${budget_usd:.2f}"
|
|
32
|
+
)
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
"""Adapters that turn another product's payload into an AgentCOGS ingest event."""
|
|
@@ -0,0 +1,44 @@
|
|
|
1
|
+
"""Starlette/FastAPI middleware — set tenant context once per request."""
|
|
2
|
+
from typing import Callable, Optional, Sequence
|
|
3
|
+
|
|
4
|
+
from starlette.middleware.base import BaseHTTPMiddleware
|
|
5
|
+
from starlette.requests import Request
|
|
6
|
+
from starlette.responses import Response
|
|
7
|
+
|
|
8
|
+
from ..context import set_customer, set_workflow
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
class AgentCOGSMiddleware(BaseHTTPMiddleware):
|
|
12
|
+
"""Populate agentcogs context from each HTTP request.
|
|
13
|
+
|
|
14
|
+
Example:
|
|
15
|
+
app.add_middleware(
|
|
16
|
+
AgentCOGSMiddleware,
|
|
17
|
+
customer_id=lambda req: req.state.tenant_id,
|
|
18
|
+
)
|
|
19
|
+
"""
|
|
20
|
+
|
|
21
|
+
def __init__(
|
|
22
|
+
self,
|
|
23
|
+
app,
|
|
24
|
+
customer_id: Callable[[Request], Optional[str]],
|
|
25
|
+
workflow_id: Optional[Callable[[Request], Optional[str]]] = None,
|
|
26
|
+
exclude_paths: Sequence[str] = ("/health", "/health/ready"),
|
|
27
|
+
):
|
|
28
|
+
super().__init__(app)
|
|
29
|
+
self._customer_id = customer_id
|
|
30
|
+
self._workflow_id = workflow_id
|
|
31
|
+
self._exclude = tuple(exclude_paths)
|
|
32
|
+
|
|
33
|
+
async def dispatch(self, request: Request, call_next) -> Response:
|
|
34
|
+
if request.url.path in self._exclude:
|
|
35
|
+
return await call_next(request)
|
|
36
|
+
try:
|
|
37
|
+
cid = self._customer_id(request)
|
|
38
|
+
set_customer(cid)
|
|
39
|
+
if self._workflow_id is not None:
|
|
40
|
+
set_workflow(self._workflow_id(request))
|
|
41
|
+
return await call_next(request)
|
|
42
|
+
finally:
|
|
43
|
+
set_customer(None)
|
|
44
|
+
set_workflow(None)
|
|
@@ -0,0 +1,38 @@
|
|
|
1
|
+
"""LangGraph helper — set tenant from configurable and wrap run()."""
|
|
2
|
+
from contextlib import contextmanager
|
|
3
|
+
from typing import Any, Iterator, Optional
|
|
4
|
+
|
|
5
|
+
from ..budget import RunContext, run
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
@contextmanager
|
|
9
|
+
def agentcogs_run(
|
|
10
|
+
config: Optional[dict] = None,
|
|
11
|
+
*,
|
|
12
|
+
customer_id: Optional[str] = None,
|
|
13
|
+
workflow_id: str = "default",
|
|
14
|
+
workspace_id: Optional[str] = None,
|
|
15
|
+
metadata: Optional[dict] = None,
|
|
16
|
+
) -> Iterator[RunContext]:
|
|
17
|
+
"""Context manager for LangGraph invoke blocks.
|
|
18
|
+
|
|
19
|
+
Reads tenant_id from config[\"configurable\"][\"tenant_id\"] when customer_id omitted.
|
|
20
|
+
|
|
21
|
+
Example:
|
|
22
|
+
with agentcogs_run(config, workflow_id=\"research_agent\") as ctx:
|
|
23
|
+
result = graph.invoke(state, config)
|
|
24
|
+
"""
|
|
25
|
+
cfg = config or {}
|
|
26
|
+
configurable = cfg.get("configurable") or {}
|
|
27
|
+
tenant = (
|
|
28
|
+
customer_id
|
|
29
|
+
or configurable.get("tenant_id")
|
|
30
|
+
or configurable.get("customer_id")
|
|
31
|
+
)
|
|
32
|
+
with run(
|
|
33
|
+
customer_id=str(tenant) if tenant else None,
|
|
34
|
+
workflow_id=workflow_id,
|
|
35
|
+
workspace_id=workspace_id,
|
|
36
|
+
metadata=metadata,
|
|
37
|
+
) as ctx:
|
|
38
|
+
yield ctx
|
|
@@ -0,0 +1,67 @@
|
|
|
1
|
+
"""Build an /v1/ingest body from a LiteLLM callback kwargs dict.
|
|
2
|
+
|
|
3
|
+
LiteLLM reports OpenAI-style usage (prompt_tokens, completion_tokens).
|
|
4
|
+
A raw provider response uses input_tokens and output_tokens. Both are accepted.
|
|
5
|
+
"""
|
|
6
|
+
from __future__ import annotations
|
|
7
|
+
|
|
8
|
+
import time
|
|
9
|
+
import uuid
|
|
10
|
+
from typing import Any
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
def _usage_int(usage: Any, *keys: str) -> int:
|
|
14
|
+
if usage is None:
|
|
15
|
+
return 0
|
|
16
|
+
if isinstance(usage, dict):
|
|
17
|
+
for key in keys:
|
|
18
|
+
value = usage.get(key)
|
|
19
|
+
if value:
|
|
20
|
+
return int(value)
|
|
21
|
+
return 0
|
|
22
|
+
for key in keys:
|
|
23
|
+
value = getattr(usage, key, None)
|
|
24
|
+
if value:
|
|
25
|
+
return int(value)
|
|
26
|
+
return 0
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
def build_ingest_event(
|
|
30
|
+
kwargs: dict,
|
|
31
|
+
*,
|
|
32
|
+
status: str,
|
|
33
|
+
error: str | None = None,
|
|
34
|
+
workspace_id: str | None = None,
|
|
35
|
+
) -> dict | None:
|
|
36
|
+
"""Return an ingest body, or None when the call has no customer."""
|
|
37
|
+
meta = kwargs.get("metadata") or {}
|
|
38
|
+
customer_id = kwargs.get("user") or meta.get("agentcogs_customer_id")
|
|
39
|
+
if not customer_id:
|
|
40
|
+
return None
|
|
41
|
+
|
|
42
|
+
model = kwargs.get("model") or "unknown"
|
|
43
|
+
cost = float(kwargs.get("response_cost") or 0)
|
|
44
|
+
usage = kwargs.get("usage")
|
|
45
|
+
start = kwargs.get("start_time")
|
|
46
|
+
ts = int(start.timestamp()) if start is not None and hasattr(start, "timestamp") else int(time.time())
|
|
47
|
+
event = {
|
|
48
|
+
"run_id": str(uuid.uuid4()),
|
|
49
|
+
"customer_id": str(customer_id),
|
|
50
|
+
"workflow_id": meta.get("agentcogs_workflow_id") or "default",
|
|
51
|
+
"ts": ts,
|
|
52
|
+
"status": status,
|
|
53
|
+
"total_usd": max(cost, 0.0),
|
|
54
|
+
"models": {
|
|
55
|
+
model: {
|
|
56
|
+
"input_tokens": _usage_int(usage, "prompt_tokens", "input_tokens"),
|
|
57
|
+
"output_tokens": _usage_int(usage, "completion_tokens", "output_tokens"),
|
|
58
|
+
"usd": max(cost, 0.0),
|
|
59
|
+
}
|
|
60
|
+
},
|
|
61
|
+
"node_costs": {},
|
|
62
|
+
"metadata": {"source": "litellm"},
|
|
63
|
+
"error": error,
|
|
64
|
+
}
|
|
65
|
+
if workspace_id:
|
|
66
|
+
event["workspace_id"] = workspace_id
|
|
67
|
+
return event
|
agentcogs/outbox.py
ADDED
|
@@ -0,0 +1,121 @@
|
|
|
1
|
+
"""Local SQLite outbox for offline-safe cost event delivery.
|
|
2
|
+
|
|
3
|
+
Pattern: write locally first, retry on next run().
|
|
4
|
+
Mirrors Stripe's recommended meter-event delivery pattern.
|
|
5
|
+
"""
|
|
6
|
+
import json
|
|
7
|
+
import pathlib
|
|
8
|
+
import sqlite3
|
|
9
|
+
import threading
|
|
10
|
+
import time
|
|
11
|
+
from typing import Any, Callable, Dict, List, Tuple
|
|
12
|
+
|
|
13
|
+
_DB_PATH = pathlib.Path.home() / ".agentcogs" / "outbox.db"
|
|
14
|
+
_LOCK = threading.Lock()
|
|
15
|
+
_MAX_ATTEMPTS = 8
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
def _ensure_db() -> sqlite3.Connection:
|
|
19
|
+
_DB_PATH.parent.mkdir(parents=True, exist_ok=True)
|
|
20
|
+
conn = sqlite3.connect(str(_DB_PATH), timeout=5.0, isolation_level=None)
|
|
21
|
+
conn.execute("PRAGMA journal_mode=WAL")
|
|
22
|
+
conn.execute(
|
|
23
|
+
"""
|
|
24
|
+
CREATE TABLE IF NOT EXISTS outbox (
|
|
25
|
+
run_id TEXT PRIMARY KEY,
|
|
26
|
+
payload TEXT NOT NULL,
|
|
27
|
+
attempts INTEGER NOT NULL DEFAULT 0,
|
|
28
|
+
next_retry INTEGER NOT NULL DEFAULT 0,
|
|
29
|
+
created_at INTEGER NOT NULL,
|
|
30
|
+
status TEXT NOT NULL DEFAULT 'pending'
|
|
31
|
+
)
|
|
32
|
+
"""
|
|
33
|
+
)
|
|
34
|
+
# Migrate older DBs without status column
|
|
35
|
+
cols = {row[1] for row in conn.execute("PRAGMA table_info(outbox)").fetchall()}
|
|
36
|
+
if "status" not in cols:
|
|
37
|
+
conn.execute("ALTER TABLE outbox ADD COLUMN status TEXT NOT NULL DEFAULT 'pending'")
|
|
38
|
+
return conn
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
def enqueue(event: dict) -> None:
|
|
42
|
+
with _LOCK:
|
|
43
|
+
conn = _ensure_db()
|
|
44
|
+
try:
|
|
45
|
+
conn.execute(
|
|
46
|
+
"INSERT OR IGNORE INTO outbox (run_id, payload, created_at, status) "
|
|
47
|
+
"VALUES (?,?,?, 'pending')",
|
|
48
|
+
(event["run_id"], json.dumps(event), int(time.time())),
|
|
49
|
+
)
|
|
50
|
+
finally:
|
|
51
|
+
conn.close()
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
def drain(post_fn: Callable[[dict], None], max_items: int = 50) -> Tuple[int, int]:
|
|
55
|
+
"""Try to deliver queued events. Returns (succeeded, failed)."""
|
|
56
|
+
now = int(time.time())
|
|
57
|
+
sent = failed = 0
|
|
58
|
+
with _LOCK:
|
|
59
|
+
conn = _ensure_db()
|
|
60
|
+
try:
|
|
61
|
+
rows: List[Tuple[str, str, int]] = conn.execute(
|
|
62
|
+
"SELECT run_id, payload, attempts FROM outbox "
|
|
63
|
+
"WHERE status = 'pending' AND attempts < ? AND next_retry <= ? "
|
|
64
|
+
"ORDER BY created_at LIMIT ?",
|
|
65
|
+
(_MAX_ATTEMPTS, now, max_items),
|
|
66
|
+
).fetchall()
|
|
67
|
+
|
|
68
|
+
for run_id, payload, attempts in rows:
|
|
69
|
+
try:
|
|
70
|
+
post_fn(json.loads(payload))
|
|
71
|
+
conn.execute("DELETE FROM outbox WHERE run_id=?", (run_id,))
|
|
72
|
+
sent += 1
|
|
73
|
+
except Exception:
|
|
74
|
+
new_attempts = attempts + 1
|
|
75
|
+
if new_attempts >= _MAX_ATTEMPTS:
|
|
76
|
+
conn.execute(
|
|
77
|
+
"UPDATE outbox SET attempts=?, status='dead' WHERE run_id=?",
|
|
78
|
+
(new_attempts, run_id),
|
|
79
|
+
)
|
|
80
|
+
else:
|
|
81
|
+
backoff = min(2**new_attempts, 3600)
|
|
82
|
+
conn.execute(
|
|
83
|
+
"UPDATE outbox SET attempts=?, next_retry=? WHERE run_id=?",
|
|
84
|
+
(new_attempts, now + backoff, run_id),
|
|
85
|
+
)
|
|
86
|
+
failed += 1
|
|
87
|
+
finally:
|
|
88
|
+
conn.close()
|
|
89
|
+
return sent, failed
|
|
90
|
+
|
|
91
|
+
|
|
92
|
+
def outbox_size() -> int:
|
|
93
|
+
return int(get_status()["pending"])
|
|
94
|
+
|
|
95
|
+
|
|
96
|
+
def clear_outbox() -> None:
|
|
97
|
+
with _LOCK:
|
|
98
|
+
conn = _ensure_db()
|
|
99
|
+
try:
|
|
100
|
+
conn.execute("DELETE FROM outbox")
|
|
101
|
+
finally:
|
|
102
|
+
conn.close()
|
|
103
|
+
|
|
104
|
+
|
|
105
|
+
def get_status() -> Dict[str, Any]:
|
|
106
|
+
with _LOCK:
|
|
107
|
+
conn = _ensure_db()
|
|
108
|
+
try:
|
|
109
|
+
pending = conn.execute(
|
|
110
|
+
"SELECT COUNT(*) FROM outbox WHERE status='pending'"
|
|
111
|
+
).fetchone()[0]
|
|
112
|
+
dead = conn.execute(
|
|
113
|
+
"SELECT COUNT(*) FROM outbox WHERE status='dead'"
|
|
114
|
+
).fetchone()[0]
|
|
115
|
+
oldest = conn.execute(
|
|
116
|
+
"SELECT MIN(created_at) FROM outbox WHERE status='pending'"
|
|
117
|
+
).fetchone()[0]
|
|
118
|
+
finally:
|
|
119
|
+
conn.close()
|
|
120
|
+
age = int(time.time()) - oldest if oldest else 0
|
|
121
|
+
return {"pending": pending, "dead": dead, "oldest_pending_age_sec": age}
|
agentcogs/ping.py
ADDED
|
@@ -0,0 +1,78 @@
|
|
|
1
|
+
"""Connectivity check — verify API key and endpoint before first run()."""
|
|
2
|
+
from dataclasses import dataclass
|
|
3
|
+
from typing import Optional
|
|
4
|
+
|
|
5
|
+
import httpx
|
|
6
|
+
|
|
7
|
+
from .config import get_config
|
|
8
|
+
from .errors import ConfigurationError, PingError
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
@dataclass(frozen=True)
|
|
12
|
+
class PingResult:
|
|
13
|
+
ok: bool
|
|
14
|
+
workspace_id: str
|
|
15
|
+
plan: str
|
|
16
|
+
server_time: Optional[int] = None
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
def ping(timeout_seconds: Optional[float] = None) -> PingResult:
|
|
20
|
+
"""Verify API key, workspace, and reachability.
|
|
21
|
+
|
|
22
|
+
Requires prior agentcogs.init() (or AGENTCOGS_* env vars).
|
|
23
|
+
"""
|
|
24
|
+
cfg = get_config()
|
|
25
|
+
if cfg.offline:
|
|
26
|
+
raise PingError(
|
|
27
|
+
"ping() is not available in offline mode. Set offline=False and provide api_key."
|
|
28
|
+
)
|
|
29
|
+
if not cfg.api_key:
|
|
30
|
+
raise ConfigurationError(
|
|
31
|
+
"agentcogs.init() requires api_key (or AGENTCOGS_API_KEY env)."
|
|
32
|
+
)
|
|
33
|
+
|
|
34
|
+
timeout = timeout_seconds if timeout_seconds is not None else cfg.timeout_seconds
|
|
35
|
+
try:
|
|
36
|
+
resp = httpx.get(
|
|
37
|
+
f"{cfg.endpoint}/v1/sdk/ping",
|
|
38
|
+
headers={
|
|
39
|
+
"Authorization": f"Bearer {cfg.api_key}",
|
|
40
|
+
"User-Agent": _user_agent(),
|
|
41
|
+
},
|
|
42
|
+
timeout=timeout,
|
|
43
|
+
)
|
|
44
|
+
except httpx.ConnectError as e:
|
|
45
|
+
raise PingError(
|
|
46
|
+
f"Cannot reach AgentCOGS API at {cfg.endpoint}. "
|
|
47
|
+
f"Check AGENTCOGS_ENDPOINT and network. ({e})"
|
|
48
|
+
) from e
|
|
49
|
+
except httpx.TimeoutException as e:
|
|
50
|
+
raise PingError(
|
|
51
|
+
f"AgentCOGS API timed out after {timeout}s ({cfg.endpoint})."
|
|
52
|
+
) from e
|
|
53
|
+
|
|
54
|
+
if resp.status_code == 401:
|
|
55
|
+
raise PingError(
|
|
56
|
+
"Invalid API key. Copy the key from Settings in the dashboard."
|
|
57
|
+
)
|
|
58
|
+
if resp.status_code == 404:
|
|
59
|
+
raise PingError(
|
|
60
|
+
f"SDK ping endpoint not found at {cfg.endpoint}. "
|
|
61
|
+
"Check AGENTCOGS_ENDPOINT (e.g. http://localhost:8000 for local dev)."
|
|
62
|
+
)
|
|
63
|
+
if resp.status_code >= 400:
|
|
64
|
+
raise PingError(f"ping failed: HTTP {resp.status_code} — {resp.text[:200]}")
|
|
65
|
+
|
|
66
|
+
data = resp.json()
|
|
67
|
+
return PingResult(
|
|
68
|
+
ok=bool(data.get("ok", True)),
|
|
69
|
+
workspace_id=str(data["workspace_id"]),
|
|
70
|
+
plan=str(data.get("plan", "free")),
|
|
71
|
+
server_time=data.get("server_time"),
|
|
72
|
+
)
|
|
73
|
+
|
|
74
|
+
|
|
75
|
+
def _user_agent() -> str:
|
|
76
|
+
from ._version import __version__
|
|
77
|
+
|
|
78
|
+
return f"agentcogs-python/{__version__}"
|
agentcogs/shutdown.py
ADDED
|
@@ -0,0 +1,38 @@
|
|
|
1
|
+
"""Graceful shutdown — drain outbox and close HTTP resources."""
|
|
2
|
+
from dataclasses import dataclass
|
|
3
|
+
|
|
4
|
+
from .client import emit_event_sync, reset_client, shutdown_executor
|
|
5
|
+
from .outbox import drain, get_status
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
@dataclass(frozen=True)
|
|
9
|
+
class ShutdownResult:
|
|
10
|
+
outbox_sent: int
|
|
11
|
+
outbox_failed: int
|
|
12
|
+
outbox_pending: int
|
|
13
|
+
outbox_dead: int
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
def shutdown(drain_timeout: float = 10.0) -> ShutdownResult:
|
|
17
|
+
"""Drain outbox, shut down ingest pool, close httpx client."""
|
|
18
|
+
from .config import get_config
|
|
19
|
+
from .client import _post
|
|
20
|
+
|
|
21
|
+
cfg = get_config()
|
|
22
|
+
sent = failed = 0
|
|
23
|
+
if not cfg.offline:
|
|
24
|
+
try:
|
|
25
|
+
sent, failed = drain(_post, max_items=500)
|
|
26
|
+
except Exception:
|
|
27
|
+
pass
|
|
28
|
+
|
|
29
|
+
shutdown_executor(wait=True, timeout=drain_timeout)
|
|
30
|
+
reset_client()
|
|
31
|
+
|
|
32
|
+
st = get_status()
|
|
33
|
+
return ShutdownResult(
|
|
34
|
+
outbox_sent=sent,
|
|
35
|
+
outbox_failed=failed,
|
|
36
|
+
outbox_pending=st["pending"],
|
|
37
|
+
outbox_dead=st["dead"],
|
|
38
|
+
)
|
agentcogs/tokens.py
ADDED
|
@@ -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)
|
|
40
|
+
[](https://www.python.org/downloads/)
|
|
41
|
+
[](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,21 @@
|
|
|
1
|
+
agentcogs/__init__.py,sha256=5qDdVACiAU4e7punzzeDUGJdEDnZFnQtnZnT2C4Wb4o,1056
|
|
2
|
+
agentcogs/__main__.py,sha256=rTZL4C-NQvvPx1P3bwws0GMGxFOvFPxLeAwG06Dn-CM,397
|
|
3
|
+
agentcogs/_shekel_compat.py,sha256=egePEUC0ql-OEebO6useAHSDLYoqXc5DR8g13QevMl4,1654
|
|
4
|
+
agentcogs/_version.py,sha256=rnObPjuBcEStqSO0S6gsdS_ot8ITOQjVj_-P1LUUYpg,22
|
|
5
|
+
agentcogs/budget.py,sha256=x7oauMAlkamAJ0v3nCOc0ekGeA-xrJ2qjKOmYAF_VUw,5546
|
|
6
|
+
agentcogs/client.py,sha256=v86F9UFOYahcjSKzqbJkYx8p46SR16Xzhd0OYLVXApE,4016
|
|
7
|
+
agentcogs/config.py,sha256=C9Rn3NSNfMydSNbBYYR9mjcrbG5VX57VLcAvyvcjXAw,2783
|
|
8
|
+
agentcogs/context.py,sha256=Kgly2n_w0Ho088k82fqjpTGvzreBLppFwgdPUDoWEmg,1098
|
|
9
|
+
agentcogs/errors.py,sha256=zkRL6A8zAz_8Y7Pd3Ka-SeP3KpyJ34UUHz0gn9C8Wb0,1032
|
|
10
|
+
agentcogs/outbox.py,sha256=hajnxk3WwbzUuejuXf159VOR__Rx5_CGlZE647Yb4Vk,4130
|
|
11
|
+
agentcogs/ping.py,sha256=rLPdWixwRnkHWmHh-nd0I1j1xWcf6l0YNyIzG8bg0UA,2389
|
|
12
|
+
agentcogs/shutdown.py,sha256=fXyGAdh_0x4XZNWypNevoExjXDncnVTgMnYIYppQ62I,980
|
|
13
|
+
agentcogs/tokens.py,sha256=zu9UU3cehK5pTY6X3sF5K1tuk4ucUrQ3KSUuV9lSe4E,163
|
|
14
|
+
agentcogs/integrations/__init__.py,sha256=X-JoNZJvuig45_VLsVqZetdlITO3dmZDbEy0m-tIG0g,83
|
|
15
|
+
agentcogs/integrations/fastapi.py,sha256=48pzjAwLPR_2bK0i37t6pG_RrRI9mDXgwp7QF45dmjs,1453
|
|
16
|
+
agentcogs/integrations/langgraph.py,sha256=8x0bzKMSdGqoqY14363LuHVurPsMgCF3oDQB1ypF9Lw,1128
|
|
17
|
+
agentcogs/integrations/litellm_event.py,sha256=SqcpkVCOKZ0JyxEmexBglBKtM3pyzD-wTOx-Iu2V5dM,2042
|
|
18
|
+
agentcogs-0.1.1.dist-info/METADATA,sha256=EjbjLRVmWhUBV1f54TuzN3RIzNkvAcBEouLYk60VpFk,6633
|
|
19
|
+
agentcogs-0.1.1.dist-info/WHEEL,sha256=2iZrJkDIkeDh0k9SikKMqFuRP-ua6mb4Ir1AsER4BEw,87
|
|
20
|
+
agentcogs-0.1.1.dist-info/licenses/LICENSE,sha256=TgFnrhRAuJjfSOjEXFD0SiLP7Wzy79NX1kIFz5spWGI,1066
|
|
21
|
+
agentcogs-0.1.1.dist-info/RECORD,,
|
|
@@ -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.
|