diff --git a/docs/my-website/docs/observability/agentcogs.md b/docs/my-website/docs/observability/agentcogs.md new file mode 100644 index 00000000000..9a972baf768 --- /dev/null +++ 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": "", + "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/__init__.py b/litellm/__init__.py index c868ae55b4f..e1e42ea057e 100644 --- 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", diff --git a/litellm/integrations/agentcogs.py b/litellm/integrations/agentcogs.py new file mode 100644 index 00000000000..3ec736a6a1e --- /dev/null +++ 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) diff --git a/litellm/litellm_core_utils/custom_logger_registry.py b/litellm/litellm_core_utils/custom_logger_registry.py index fd402b90d88..bf8f6b3bda9 100644 --- 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, diff --git a/tests/test_litellm/integrations/test_agentcogs.py b/tests/test_litellm/integrations/test_agentcogs.py new file mode 100644 index 00000000000..5b5ee0e9d2b --- /dev/null +++ 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)