feat(integrations): AgentCOGS per-customer margin callback

Add AgentCOGSLogger (Lago/OpenMeter pattern) posting per-completion cost to AgentCOGS /v1/ingest with user= tenant attribution, unit tests, and docs. Co-authored-by: Cursor <cursoragent@cursor.com>
2026-09-27 01:22:18 +00:00 · 2026-05-19 19:37:58 +05:30 · 2026-05-19 19:37:58 +05:30 · d3627e4523
commit d3627e4523
parent cff3e0b75e
5 changed files with 433 additions and 0 deletions
--- a/docs/my-website/docs/observability/agentcogs.md
+++ b/docs/my-website/docs/observability/agentcogs.md
@ -0,0 +1,83 @@
+# 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": "<uuid>",
+  "workspace_id": "<AGENTCOGS_WORKSPACE_ID>",
+  "customer_id": "<user or metadata.agentcogs_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)
--- a/litellm/init.py
+++ b/litellm/init.py
@ -100,6 +100,7 @@ service_callback: List[CALLBACK_TYPES] = []
 audit_log_callbacks: List[CALLBACK_TYPES] = []
 # logging_callback_manager is lazy-loaded via __getattr__
 _custom_logger_compatible_callbacks_literal = Literal[
+    "agentcogs",
    "lago",
    "openmeter",
    "logfire",
--- a/litellm/integrations/agentcogs.py
+++ b/litellm/integrations/agentcogs.py
@ -0,0 +1,193 @@
+# AgentCOGS per-customer margin callback — https://github.com/vaibhav11123/agentcogs
+
+import json
+import os
+import time
+from typing import Any, Optional
+
+import litellm
+from litellm._logging import verbose_logger
+from litellm._uuid import uuid
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.llms.custom_httpx.http_handler import (
+    HTTPHandler,
+    get_async_httpx_client,
+    httpxSpecialProvider,
+)
+
+
+def _usage_int(usage: Any, key: str) -> int:
+    if usage is None:
+        return 0
+    if isinstance(usage, dict):
+        return int(usage.get(key, 0) or 0)
+    return int(getattr(usage, key, 0) or 0)
+
+
+def _extract_usage(kwargs: dict, response_obj: Any) -> dict:
+    usage = kwargs.get("usage")
+    if usage is not None:
+        return {
+            "prompt_tokens": _usage_int(usage, "prompt_tokens"),
+            "completion_tokens": _usage_int(usage, "completion_tokens"),
+        }
+    if (
+        isinstance(response_obj, litellm.ModelResponse)
+        or isinstance(response_obj, litellm.EmbeddingResponse)
+    ) and hasattr(response_obj, "usage"):
+        u = response_obj.usage
+        if isinstance(u, dict):
+            return {
+                "prompt_tokens": int(u.get("prompt_tokens", 0) or 0),
+                "completion_tokens": int(u.get("completion_tokens", 0) or 0),
+            }
+        return {
+            "prompt_tokens": int(getattr(u, "prompt_tokens", 0) or 0),
+            "completion_tokens": int(getattr(u, "completion_tokens", 0) or 0),
+        }
+    return {"prompt_tokens": 0, "completion_tokens": 0}
+
+
+def _metadata_from_kwargs(kwargs: dict) -> dict:
+    meta = kwargs.get("metadata")
+    if isinstance(meta, dict):
+        return meta
+    litellm_params = kwargs.get("litellm_params", {}) or {}
+    meta = litellm_params.get("metadata", {})
+    return meta if isinstance(meta, dict) else {}
+
+
+class AgentCOGSLogger(CustomLogger):
+    """POST per-completion cost to AgentCOGS /v1/ingest for B2B per-customer margin."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.validate_environment()
+        self.async_http_handler = get_async_httpx_client(
+            llm_provider=httpxSpecialProvider.LoggingCallback
+        )
+        self.sync_http_handler = HTTPHandler()
+
+    def validate_environment(self) -> None:
+        missing_keys = []
+        if os.getenv("AGENTCOGS_API_KEY") is None:
+            missing_keys.append("AGENTCOGS_API_KEY")
+        if os.getenv("AGENTCOGS_WORKSPACE_ID") is None:
+            missing_keys.append("AGENTCOGS_WORKSPACE_ID")
+        if len(missing_keys) > 0:
+            raise Exception("Missing keys={} in environment.".format(missing_keys))
+
+    def _endpoint(self) -> str:
+        base = os.getenv("AGENTCOGS_ENDPOINT", "https://api.agentcogs.dev").rstrip("/")
+        return f"{base}/v1/ingest"
+
+    def _headers(self) -> dict:
+        return {
+            "Content-Type": "application/json",
+            "Authorization": "Bearer {}".format(os.environ["AGENTCOGS_API_KEY"]),
+            "X-AgentCOGS-SDK-Version": "litellm-callback/0.1.0",
+        }
+
+    def _build_event(
+        self,
+        kwargs: dict,
+        response_obj: Any,
+        *,
+        status: str,
+        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")
+        if not customer_id:
+            return None
+
+        model = kwargs.get("model") or "unknown"
+        cost = float(kwargs.get("response_cost") or 0)
+        usage = _extract_usage(kwargs, response_obj)
+
+        ts_source = start_time or kwargs.get("start_time")
+        if ts_source is not None and hasattr(ts_source, "timestamp"):
+            ts = int(ts_source.timestamp())
+        else:
+            ts = int(time.time())
+
+        return {
+            "run_id": str(uuid.uuid4()),
+            "workspace_id": os.environ["AGENTCOGS_WORKSPACE_ID"],
+            "customer_id": str(customer_id),
+            "workflow_id": meta.get("agentcogs_workflow_id", "default"),
+            "ts": ts,
+            "status": status,
+            "total_usd": cost,
+            "models": {
+                model: {
+                    "input_tokens": usage["prompt_tokens"],
+                    "output_tokens": usage["completion_tokens"],
+                    "usd": cost,
+                }
+            },
+            "node_costs": {},
+            "metadata": {"source": "litellm"},
+            "error": error,
+        }
+
+    async def _async_post(self, event: dict) -> None:
+        try:
+            response = await self.async_http_handler.post(
+                url=self._endpoint(),
+                data=json.dumps(event),
+                headers=self._headers(),
+            )
+            response.raise_for_status()
+        except Exception as e:
+            verbose_logger.debug("AgentCOGS callback error: {}".format(e))
+
+    def _sync_post(self, event: dict) -> None:
+        try:
+            response = self.sync_http_handler.post(
+                url=self._endpoint(),
+                data=json.dumps(event),
+                headers=self._headers(),
+            )
+            response.raise_for_status()
+        except Exception as e:
+            verbose_logger.debug("AgentCOGS callback error: {}".format(e))
+
+    def log_success_event(self, kwargs, response_obj, start_time, end_time):
+        event = self._build_event(
+            kwargs, response_obj, status="completed", start_time=start_time
+        )
+        if event:
+            self._sync_post(event)
+
+    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+        event = self._build_event(
+            kwargs, response_obj, status="completed", start_time=start_time
+        )
+        if event:
+            await self._async_post(event)
+
+    def log_failure_event(self, kwargs, response_obj, start_time, end_time):
+        err = str(kwargs.get("exception", "error"))[:500]
+        event = self._build_event(
+            kwargs,
+            response_obj,
+            status="error",
+            start_time=start_time,
+            error=err,
+        )
+        if event:
+            self._sync_post(event)
+
+    async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
+        err = str(kwargs.get("exception", "error"))[:500]
+        event = self._build_event(
+            kwargs,
+            response_obj,
+            status="error",
+            start_time=start_time,
+            error=err,
+        )
+        if event:
+            await self._async_post(event)
--- a/litellm/litellm_core_utils/custom_logger_registry.py
+++ b/litellm/litellm_core_utils/custom_logger_registry.py
@ -31,6 +31,7 @@ from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger
 from litellm.integrations.gcs_pubsub.pub_sub import GcsPubSubLogger
 from litellm.integrations.gitlab import GitLabPromptManager
 from litellm.integrations.humanloop import HumanloopLogger
+from litellm.integrations.agentcogs import AgentCOGSLogger
 from litellm.integrations.lago import LagoLogger
 from litellm.integrations.langfuse.langfuse_prompt_management import (
    LangfusePromptManagement,
@ -59,6 +60,7 @@ class CustomLoggerRegistry:
    """

    CALLBACK_CLASS_STR_TO_CLASS_TYPE = {
+        "agentcogs": AgentCOGSLogger,
        "lago": LagoLogger,
        "openmeter": OpenMeterLogger,
        "braintrust": BraintrustLogger,
--- a/tests/test_litellm/integrations/test_agentcogs.py
+++ b/tests/test_litellm/integrations/test_agentcogs.py
@ -0,0 +1,154 @@
+import json
+import os
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from litellm.integrations.agentcogs import AgentCOGSLogger
+
+
+class TestAgentCOGSIntegration:
+    def setup_method(self):
+        os.environ["AGENTCOGS_API_KEY"] = "test-api-key"
+        os.environ["AGENTCOGS_WORKSPACE_ID"] = "ws-test-uuid"
+        os.environ["AGENTCOGS_ENDPOINT"] = "https://api.agentcogs.test"
+
+    def teardown_method(self):
+        for key in (
+            "AGENTCOGS_API_KEY",
+            "AGENTCOGS_WORKSPACE_ID",
+            "AGENTCOGS_ENDPOINT",
+        ):
+            os.environ.pop(key, None)
+
+    def test_logger_initialization(self):
+        logger = AgentCOGSLogger()
+        assert logger is not None
+
+    def test_logger_missing_api_key(self):
+        os.environ.pop("AGENTCOGS_API_KEY", None)
+        with pytest.raises(Exception, match="Missing keys.*AGENTCOGS_API_KEY"):
+            AgentCOGSLogger()
+
+    def test_build_event_with_user(self):
+        logger = AgentCOGSLogger()
+        kwargs = {
+            "user": "acme_corp",
+            "model": "gpt-4o-mini",
+            "response_cost": 0.002,
+            "usage": {"prompt_tokens": 10, "completion_tokens": 5},
+            "metadata": {"agentcogs_workflow_id": "support_bot"},
+        }
+        event = logger._build_event(kwargs, {}, status="completed")
+        assert event is not None
+        assert event["customer_id"] == "acme_corp"
+        assert event["workspace_id"] == "ws-test-uuid"
+        assert event["status"] == "completed"
+        assert event["workflow_id"] == "support_bot"
+        assert event["ts"] == pytest.approx(int(__import__("time").time()), abs=5)
+        assert event["models"]["gpt-4o-mini"]["input_tokens"] == 10
+        assert event["models"]["gpt-4o-mini"]["usd"] == 0.002
+        assert event["metadata"]["source"] == "litellm"
+
+    def test_build_event_skips_without_customer(self):
+        logger = AgentCOGSLogger()
+        kwargs = {
+            "model": "gpt-4o-mini",
+            "response_cost": 0.001,
+        }
+        assert logger._build_event(kwargs, {}, status="completed") is None
+
+    def test_build_event_customer_from_metadata(self):
+        logger = AgentCOGSLogger()
+        kwargs = {
+            "litellm_params": {
+                "metadata": {"agentcogs_customer_id": "meta_tenant"},
+            },
+            "model": "gpt-4",
+            "response_cost": 0.01,
+        }
+        event = logger._build_event(kwargs, {}, status="completed")
+        assert event is not None
+        assert event["customer_id"] == "meta_tenant"
+
+    def test_build_event_error_status(self):
+        logger = AgentCOGSLogger()
+        kwargs = {
+            "user": "acme",
+            "model": "gpt-4",
+            "response_cost": 0,
+            "exception": "rate limited",
+        }
+        event = logger._build_event(
+            kwargs, {}, status="error", error="rate limited"
+        )
+        assert event is not None
+        assert event["status"] == "error"
+        assert event["error"] == "rate limited"
+
+    @patch("litellm.integrations.agentcogs.HTTPHandler")
+    def test_log_success_event_posts(self, mock_http_handler):
+        mock_post = MagicMock()
+        mock_http_handler.return_value.post = mock_post
+
+        logger = AgentCOGSLogger()
+        kwargs = {
+            "user": "test-user",
+            "model": "gpt-3.5-turbo",
+            "response_cost": 0.001,
+            "usage": {"prompt_tokens": 3, "completion_tokens": 2},
+        }
+        logger.log_success_event(kwargs, {}, None, None)
+        mock_post.assert_called_once()
+        payload = json.loads(mock_post.call_args[1]["data"])
+        assert payload["customer_id"] == "test-user"
+        assert payload["status"] == "completed"
+
+    @patch("litellm.integrations.agentcogs.get_async_httpx_client")
+    @pytest.mark.asyncio
+    async def test_async_log_success_event_posts(self, mock_get_client):
+        mock_post = AsyncMock()
+        mock_client = MagicMock()
+        mock_client.post = mock_post
+        mock_get_client.return_value = mock_client
+
+        logger = AgentCOGSLogger()
+        kwargs = {
+            "user": "async-user",
+            "model": "gpt-4",
+            "response_cost": 0.002,
+            "usage": {"prompt_tokens": 20, "completion_tokens": 10},
+        }
+        await logger.async_log_success_event(kwargs, {}, None, None)
+        mock_post.assert_called_once()
+        payload = json.loads(mock_post.call_args[1]["data"])
+        assert payload["customer_id"] == "async-user"
+
+    @patch("litellm.integrations.agentcogs.get_async_httpx_client")
+    @pytest.mark.asyncio
+    async def test_async_log_skips_without_customer(self, mock_get_client):
+        mock_post = AsyncMock()
+        mock_client = MagicMock()
+        mock_client.post = mock_post
+        mock_get_client.return_value = mock_client
+
+        logger = AgentCOGSLogger()
+        kwargs = {"model": "gpt-4", "response_cost": 0.001}
+        await logger.async_log_success_event(kwargs, {}, None, None)
+        mock_post.assert_not_called()
+
+    @patch("litellm.integrations.agentcogs.get_async_httpx_client")
+    @pytest.mark.asyncio
+    async def test_async_log_swallows_post_errors(self, mock_get_client):
+        mock_post = AsyncMock(side_effect=Exception("network down"))
+        mock_client = MagicMock()
+        mock_client.post = mock_post
+        mock_get_client.return_value = mock_client
+
+        logger = AgentCOGSLogger()
+        kwargs = {
+            "user": "u1",
+            "model": "gpt-4",
+            "response_cost": 0.001,
+        }
+        await logger.async_log_success_event(kwargs, {}, None, None)