From f12939e53818de788a7139387a8b6981e9b64ebe Mon Sep 17 00:00:00 2001 From: vaibhavsurendra-web Date: Tue, 19 May 2026 19:55:00 +0530 Subject: [PATCH] fix(integrations): address AgentCOGS PR review feedback - Remove in-repo docs (belongs in BerriAI/litellm-docs) - Use litellm_call_id for ingest run_id idempotency - Lago-style AGENTCOGS_CHARGE_BY proxy attribution; ignore client metadata on proxy requests - Log callback HTTP errors at warning level Co-authored-by: Cursor --- .../docs/observability/agentcogs.md | 83 ------------------- litellm/integrations/agentcogs.py | 58 +++++++++++-- .../integrations/test_agentcogs.py | 49 ++++++++++- 3 files changed, 98 insertions(+), 92 deletions(-) delete mode 100644 docs/my-website/docs/observability/agentcogs.md diff --git a/docs/my-website/docs/observability/agentcogs.md b/docs/my-website/docs/observability/agentcogs.md deleted file mode 100644 index 9a972baf768..00000000000 --- a/docs/my-website/docs/observability/agentcogs.md +++ /dev/null @@ -1,83 +0,0 @@ -# AgentCOGS - Per-customer margin - -[AgentCOGS](https://github.com/vaibhav11123/agentcogs) tracks per-customer LLM cost and gross margin for B2B SaaS (cost + revenue), alongside your existing proxy and observability stack. - -## Quick Start - -Use one line to send successful completion cost to AgentCOGS: - -Get your AgentCOGS [API key and workspace id](https://github.com/vaibhav11123/agentcogs/blob/main/docs/quickstart.md). - -```python -import os -import litellm - -os.environ["AGENTCOGS_API_KEY"] = "" -os.environ["AGENTCOGS_WORKSPACE_ID"] = "" -# optional — defaults to https://api.agentcogs.dev -os.environ["AGENTCOGS_ENDPOINT"] = "https://api.agentcogs.dev" - -litellm.success_callback = ["agentcogs"] - -response = litellm.completion( - model="gpt-4o-mini", - messages=[{"role": "user", "content": "Hi"}], - user="your_customer_id", # 👈 B2B tenant id → AgentCOGS customer_id - metadata={"agentcogs_workflow_id": "support_bot"}, -) -``` - -- SDK -- PROXY - -```yaml -litellm_settings: - callbacks: ["agentcogs"] -``` - -Set `AGENTCOGS_API_KEY`, `AGENTCOGS_WORKSPACE_ID`, and optionally `AGENTCOGS_ENDPOINT` in the proxy environment. - -## Environment variables - -| Variable | Required | Description | -|----------|----------|-------------| -| `AGENTCOGS_API_KEY` | Yes | Workspace API key (`acg_live_...`) | -| `AGENTCOGS_WORKSPACE_ID` | Yes | Workspace UUID | -| `AGENTCOGS_ENDPOINT` | No | API base URL (default `https://api.agentcogs.dev`) | - -## Tenant attribution - -Pass `user=` on each completion (same pattern as [Lago](./lago.md)). LiteLLM maps it to AgentCOGS `customer_id`. - -Alternatively set `metadata.agentcogs_customer_id` if you cannot use `user`. - -Completions without a customer id are skipped (proxy traffic is not blocked). - -## What LiteLLM sends - -Each successful or failed completion POSTs to `POST /v1/ingest`: - -```json -{ - "run_id": "", - "workspace_id": "", - "customer_id": "", - "workflow_id": "default", - "ts": 1710000000, - "status": "completed", - "total_usd": 0.0012, - "models": { - "gpt-4o-mini": { - "input_tokens": 10, - "output_tokens": 5, - "usd": 0.0012 - } - }, - "metadata": { "source": "litellm" } -} -``` - -## Learn more - -- [AgentCOGS quickstart](https://github.com/vaibhav11123/agentcogs/blob/main/docs/quickstart.md) -- [LiteLLM callback (user-landed)](https://github.com/vaibhav11123/agentcogs/blob/main/docs/integrations/litellm.md) diff --git a/litellm/integrations/agentcogs.py b/litellm/integrations/agentcogs.py index 3ec736a6a1e..9580f7d0675 100644 --- a/litellm/integrations/agentcogs.py +++ b/litellm/integrations/agentcogs.py @@ -3,7 +3,7 @@ import json import os import time -from typing import Any, Optional +from typing import Any, Literal, Optional import litellm from litellm._logging import verbose_logger @@ -57,6 +57,50 @@ def _metadata_from_kwargs(kwargs: dict) -> dict: return meta if isinstance(meta, dict) else {} +def _resolve_customer_id(kwargs: dict) -> Optional[str]: + """ + Resolve tenant id for AgentCOGS ingest. + + Proxy mode (default): same attribution sources as Lago — end_user_id from proxy + body, or authenticated key metadata (user_id / team_id). Client-supplied + metadata.agentcogs_customer_id is not trusted on proxy requests. + + Direct SDK mode (no proxy_server_request): kwargs user or metadata.agentcogs_customer_id. + """ + litellm_params = kwargs.get("litellm_params", {}) or {} + meta = litellm_params.get("metadata", {}) or {} + proxy_server_request = litellm_params.get("proxy_server_request") + is_proxy = bool(proxy_server_request) + + proxy_body = (proxy_server_request or {}).get("body") or {} + end_user_id = proxy_body.get("user") or kwargs.get("user") + user_id = meta.get("user_api_key_user_id") + team_id = meta.get("user_api_key_team_id") + + charge_by: Literal["end_user_id", "team_id", "user_id"] = "end_user_id" + if os.getenv("AGENTCOGS_CHARGE_BY") is not None and isinstance( + os.environ["AGENTCOGS_CHARGE_BY"], str + ): + if os.environ["AGENTCOGS_CHARGE_BY"] in ("end_user_id", "user_id", "team_id"): + charge_by = os.environ["AGENTCOGS_CHARGE_BY"] # type: ignore + else: + raise Exception("invalid AGENTCOGS_CHARGE_BY set") + + if charge_by == "end_user_id": + customer_id = end_user_id + elif charge_by == "team_id": + customer_id = team_id + else: + customer_id = user_id + + if customer_id is None and not is_proxy: + customer_id = _metadata_from_kwargs(kwargs).get("agentcogs_customer_id") + + if customer_id is None: + return None + return str(customer_id) + + class AgentCOGSLogger(CustomLogger): """POST per-completion cost to AgentCOGS /v1/ingest for B2B per-customer margin.""" @@ -97,11 +141,11 @@ class AgentCOGSLogger(CustomLogger): start_time: Any = None, error: Optional[str] = None, ) -> Optional[dict]: - meta = _metadata_from_kwargs(kwargs) - customer_id = kwargs.get("user") or meta.get("agentcogs_customer_id") + customer_id = _resolve_customer_id(kwargs) if not customer_id: return None + meta = _metadata_from_kwargs(kwargs) model = kwargs.get("model") or "unknown" cost = float(kwargs.get("response_cost") or 0) usage = _extract_usage(kwargs, response_obj) @@ -113,9 +157,9 @@ class AgentCOGSLogger(CustomLogger): ts = int(time.time()) return { - "run_id": str(uuid.uuid4()), + "run_id": str(kwargs.get("litellm_call_id") or uuid.uuid4()), "workspace_id": os.environ["AGENTCOGS_WORKSPACE_ID"], - "customer_id": str(customer_id), + "customer_id": customer_id, "workflow_id": meta.get("agentcogs_workflow_id", "default"), "ts": ts, "status": status, @@ -141,7 +185,7 @@ class AgentCOGSLogger(CustomLogger): ) response.raise_for_status() except Exception as e: - verbose_logger.debug("AgentCOGS callback error: {}".format(e)) + verbose_logger.warning("AgentCOGS callback error: {}".format(e)) def _sync_post(self, event: dict) -> None: try: @@ -152,7 +196,7 @@ class AgentCOGSLogger(CustomLogger): ) response.raise_for_status() except Exception as e: - verbose_logger.debug("AgentCOGS callback error: {}".format(e)) + verbose_logger.warning("AgentCOGS callback error: {}".format(e)) def log_success_event(self, kwargs, response_obj, start_time, end_time): event = self._build_event( diff --git a/tests/test_litellm/integrations/test_agentcogs.py b/tests/test_litellm/integrations/test_agentcogs.py index 5b5ee0e9d2b..f5d1577fff9 100644 --- a/tests/test_litellm/integrations/test_agentcogs.py +++ b/tests/test_litellm/integrations/test_agentcogs.py @@ -4,7 +4,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest -from litellm.integrations.agentcogs import AgentCOGSLogger +from litellm.integrations.agentcogs import AgentCOGSLogger, _resolve_customer_id class TestAgentCOGSIntegration: @@ -18,6 +18,7 @@ class TestAgentCOGSIntegration: "AGENTCOGS_API_KEY", "AGENTCOGS_WORKSPACE_ID", "AGENTCOGS_ENDPOINT", + "AGENTCOGS_CHARGE_BY", ): os.environ.pop(key, None) @@ -58,7 +59,8 @@ class TestAgentCOGSIntegration: } assert logger._build_event(kwargs, {}, status="completed") is None - def test_build_event_customer_from_metadata(self): + def test_build_event_customer_from_metadata_sdk_mode(self): + """SDK/direct calls (no proxy): metadata.agentcogs_customer_id is allowed.""" logger = AgentCOGSLogger() kwargs = { "litellm_params": { @@ -71,6 +73,49 @@ class TestAgentCOGSIntegration: assert event is not None assert event["customer_id"] == "meta_tenant" + def test_run_id_uses_litellm_call_id(self): + logger = AgentCOGSLogger() + kwargs = { + "user": "acme", + "litellm_call_id": "call-abc-123", + "model": "gpt-4", + "response_cost": 0.001, + } + event = logger._build_event(kwargs, {}, status="completed") + assert event is not None + assert event["run_id"] == "call-abc-123" + + def test_proxy_uses_end_user_from_body(self): + kwargs = { + "litellm_params": { + "proxy_server_request": {"body": {"user": "proxy_tenant"}}, + "metadata": { + "agentcogs_customer_id": "malicious_tenant", + "user_api_key_user_id": "key_user", + }, + }, + } + assert _resolve_customer_id(kwargs) == "proxy_tenant" + + def test_proxy_ignores_client_metadata_customer_id(self): + kwargs = { + "litellm_params": { + "proxy_server_request": {"body": {}}, + "metadata": {"agentcogs_customer_id": "malicious_tenant"}, + }, + } + assert _resolve_customer_id(kwargs) is None + + def test_proxy_charge_by_team_id(self): + os.environ["AGENTCOGS_CHARGE_BY"] = "team_id" + kwargs = { + "litellm_params": { + "proxy_server_request": {"body": {"user": "end_user"}}, + "metadata": {"user_api_key_team_id": "team-99"}, + }, + } + assert _resolve_customer_id(kwargs) == "team-99" + def test_build_event_error_status(self): logger = AgentCOGSLogger() kwargs = {