diff --git a/litellm/__init__.py b/litellm/__init__.py index e6c30e12286..0d89d73ee66 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -153,6 +153,7 @@ _custom_logger_compatible_callbacks_literal = Literal[ "bitbucket", "gitlab", "cloudzero", + "faros", "focus", "vantage", "posthog", diff --git a/litellm/integrations/faros/__init__.py b/litellm/integrations/faros/__init__.py new file mode 100644 index 00000000000..2c499dbce5d --- /dev/null +++ b/litellm/integrations/faros/__init__.py @@ -0,0 +1,3 @@ +from litellm.integrations.faros.faros_logger import FarosLogger + +__all__ = ["FarosLogger"] diff --git a/litellm/integrations/faros/faros_logger.py b/litellm/integrations/faros/faros_logger.py new file mode 100644 index 00000000000..dd373f8b625 --- /dev/null +++ b/litellm/integrations/faros/faros_logger.py @@ -0,0 +1,177 @@ +""" +Faros AI (https://www.faros.ai/) integration. + +Sends LiteLLM usage data to a Faros graph so LLM usage shows up alongside the +rest of an engineering org's productivity data. + +Each successful LLM request is recorded as a Faros canonical +``vcs_UserToolUsage`` row (with its backing ``vcs_UserTool`` and ``vcs_User`` +rows) via the Faros GraphQL API - the same write path Faros' own first-party +usage trackers use (see faros-ai/faros-vscode-extension). +""" + +import asyncio +import os +from datetime import datetime, timezone +from typing import Any, Dict, List, Optional + +from litellm._logging import verbose_logger +from litellm.constants import LITELLM_PROXY_ADMIN_NAME +from litellm.integrations.custom_batch_logger import CustomBatchLogger +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + get_async_httpx_client, + httpxSpecialProvider, +) +from litellm.types.utils import StandardLoggingPayload + +FAROS_DEFAULT_API_URL = "https://prod.api.faros.ai" +FAROS_DEFAULT_GRAPH = "default" +FAROS_DEFAULT_ORIGIN = "litellm" +FAROS_DEFAULT_USER_SOURCE = "LiteLLM" +FAROS_DEFAULT_TOOL_CATEGORY = "LiteLLM" + +USAGE_MUTATION = ( + "mutation LiteLLMUserToolUsage($usages: [vcs_UserToolUsage_insert_input!]!) { " + "insert_vcs_UserToolUsage(objects: $usages, on_conflict: {" + "constraint: vcs_UserToolUsage_pkey, update_columns: [refreshedAt, origin]" + "}) { affected_rows } }" +) + + +class FarosLogger(CustomBatchLogger): + """ + Batches LiteLLM usage events and upserts them into a Faros graph. + + Environment Variables: + FAROS_API_KEY: Faros API key (required) + FAROS_API_URL: Faros API base url (default: https://prod.api.faros.ai) + FAROS_GRAPH: Faros graph to write to (default: "default") + FAROS_ORIGIN: Origin recorded on every row (default: "litellm") + FAROS_USER_SOURCE: vcs_User.source for LiteLLM users (default: "LiteLLM") + FAROS_TOOL_CATEGORY: vcs_UserTool.tool.category (default: "LiteLLM") + """ + + def __init__( + self, + api_key: Optional[str] = None, + api_url: Optional[str] = None, + graph: Optional[str] = None, + origin: Optional[str] = None, + user_source: Optional[str] = None, + tool_category: Optional[str] = None, + async_httpx_client: Optional[AsyncHTTPHandler] = None, + **kwargs, + ): + self.api_key = api_key or os.getenv("FAROS_API_KEY") + if not self.api_key: + raise ValueError( + "FAROS_API_KEY is not set. Set it in your environment or pass api_key to FarosLogger." + ) + resolved_api_url = ( + api_url or os.getenv("FAROS_API_URL") or FAROS_DEFAULT_API_URL + ).rstrip("/") + self.graph = graph or os.getenv("FAROS_GRAPH") or FAROS_DEFAULT_GRAPH + self.graphql_endpoint = f"{resolved_api_url}/graphs/{self.graph}/graphql" + self.origin = origin or os.getenv("FAROS_ORIGIN") or FAROS_DEFAULT_ORIGIN + self.user_source = ( + user_source or os.getenv("FAROS_USER_SOURCE") or FAROS_DEFAULT_USER_SOURCE + ) + self.tool_category = ( + tool_category + or os.getenv("FAROS_TOOL_CATEGORY") + or FAROS_DEFAULT_TOOL_CATEGORY + ) + self.async_httpx_client = async_httpx_client or get_async_httpx_client( + llm_provider=httpxSpecialProvider.LoggingCallback + ) + super().__init__(**kwargs, flush_lock=asyncio.Lock()) + try: + asyncio.create_task(self.periodic_flush()) + except RuntimeError: + verbose_logger.debug( + "FarosLogger: no running event loop; relying on batch_size flushes" + ) + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + payload: Optional[StandardLoggingPayload] = kwargs.get( + "standard_logging_object" + ) + if payload is None: + verbose_logger.warning( + "FarosLogger: standard_logging_object missing, skipping event" + ) + return + self.log_queue.append(self._usage_record(payload)) + if len(self.log_queue) >= self.batch_size: + await self.flush_queue() + + def _usage_record(self, payload: StandardLoggingPayload) -> Dict[str, str]: + metadata = payload.get("metadata") or {} + key_user_id = metadata.get("user_api_key_user_id") + if key_user_id == LITELLM_PROXY_ADMIN_NAME: + # placeholder identity for master key / admin requests, not a real user + key_user_id = None + user_uid = ( + metadata.get("user_api_key_user_email") + or key_user_id + or payload.get("end_user") + or "unknown" + ) + used_at = datetime.fromtimestamp( + payload["startTime"], tz=timezone.utc + ).isoformat(timespec="milliseconds") + return {"user_uid": user_uid, "used_at": used_at} + + def _usage_insert_input(self, record: Dict[str, str]) -> Dict[str, Any]: + return { + "usedAt": record["used_at"], + "origin": self.origin, + "userTool": { + "data": { + "tool": {"category": self.tool_category}, + "origin": self.origin, + "user": { + "data": { + "uid": record["user_uid"], + "source": self.user_source, + "origin": self.origin, + }, + "on_conflict": { + "constraint": "vcs_User_pkey", + "update_columns": ["refreshedAt"], + }, + }, + }, + "on_conflict": { + "constraint": "vcs_UserTool_pkey", + "update_columns": ["refreshedAt"], + }, + }, + } + + async def async_send_batch(self): + if not self.log_queue: + return + # (user, usedAt) is the row's primary key; a single upsert statement + # cannot touch the same row twice + unique_records = { + (record["user_uid"], record["used_at"]): record for record in self.log_queue + } + usages: List[Dict[str, Any]] = [ + self._usage_insert_input(record) for record in unique_records.values() + ] + response = await self.async_httpx_client.post( + url=self.graphql_endpoint, + json={"query": USAGE_MUTATION, "variables": {"usages": usages}}, + headers={"authorization": self.api_key, "content-type": "application/json"}, + ) + response.raise_for_status() + errors = response.json().get("errors") + if errors: + raise ValueError(f"Faros GraphQL request failed: {errors}") + verbose_logger.debug( + "FarosLogger: sent %s usage records to Faros graph %s", + len(usages), + self.graph, + ) diff --git a/litellm/litellm_core_utils/custom_logger_registry.py b/litellm/litellm_core_utils/custom_logger_registry.py index fd402b90d88..ffad81ce726 100644 --- a/litellm/litellm_core_utils/custom_logger_registry.py +++ b/litellm/litellm_core_utils/custom_logger_registry.py @@ -24,6 +24,7 @@ from litellm.integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger from litellm.integrations.datadog.datadog_metrics import DatadogMetricsLogger from litellm.integrations.deepeval import DeepEvalLogger from litellm.integrations.dotprompt import DotpromptManager +from litellm.integrations.faros.faros_logger import FarosLogger from litellm.integrations.focus.focus_logger import FocusLogger from litellm.integrations.vantage.vantage_logger import VantageLogger from litellm.integrations.galileo import GalileoObserve @@ -101,6 +102,7 @@ class CustomLoggerRegistry: "bitbucket": BitBucketPromptManager, "gitlab": GitLabPromptManager, "cloudzero": CloudZeroLogger, + "faros": FarosLogger, "focus": FocusLogger, "vantage": VantageLogger, "posthog": PostHogLogger, diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index dbfcf55d75d..460f2edbd14 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4113,6 +4113,15 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 cloudzero_logger = CloudZeroLogger() _in_memory_loggers.append(cloudzero_logger) return cloudzero_logger # type: ignore + elif logging_integration == "faros": + from litellm.integrations.faros.faros_logger import FarosLogger + + for callback in _in_memory_loggers: + if isinstance(callback, FarosLogger): + return callback # type: ignore + faros_logger = FarosLogger() + _in_memory_loggers.append(faros_logger) + return faros_logger # type: ignore elif logging_integration == "focus": from litellm.integrations.focus.focus_logger import FocusLogger @@ -4542,6 +4551,12 @@ def get_custom_logger_compatible_class( # noqa: PLR0915 for callback in _in_memory_loggers: if isinstance(callback, CloudZeroLogger): return callback + elif logging_integration == "faros": + from litellm.integrations.faros.faros_logger import FarosLogger + + for callback in _in_memory_loggers: + if isinstance(callback, FarosLogger): + return callback elif logging_integration == "focus": from litellm.integrations.focus.focus_logger import FocusLogger diff --git a/tests/test_litellm/integrations/faros/test_faros_logger.py b/tests/test_litellm/integrations/faros/test_faros_logger.py new file mode 100644 index 00000000000..29e20da5fa8 --- /dev/null +++ b/tests/test_litellm/integrations/faros/test_faros_logger.py @@ -0,0 +1,203 @@ +import json +from unittest.mock import AsyncMock, MagicMock + +import pytest + +import litellm +from litellm.integrations.faros.faros_logger import USAGE_MUTATION, FarosLogger +from litellm.litellm_core_utils.custom_logger_registry import CustomLoggerRegistry + + +def make_logger(mock_post=None, **kwargs) -> FarosLogger: + client = MagicMock() + client.post = mock_post or AsyncMock() + return FarosLogger(api_key="test-key", async_httpx_client=client, **kwargs) + + +def make_response(body=None): + response = MagicMock() + response.json.return_value = body if body is not None else {"data": {}} + return response + + +def make_kwargs( + startTime: float = 1765000000.5, + user_email=None, + user_id=None, + end_user=None, +) -> dict: + return { + "standard_logging_object": { + "startTime": startTime, + "end_user": end_user, + "metadata": { + "user_api_key_user_email": user_email, + "user_api_key_user_id": user_id, + }, + } + } + + +@pytest.mark.asyncio +async def test_init_requires_api_key(monkeypatch): + monkeypatch.delenv("FAROS_API_KEY", raising=False) + with pytest.raises(ValueError, match="FAROS_API_KEY"): + FarosLogger() + + +@pytest.mark.asyncio +async def test_user_uid_resolution_priority(): + logger = make_logger() + + await logger.async_log_success_event( + make_kwargs(user_email="dev@example.com", user_id="user-1", end_user="cust"), + None, + None, + None, + ) + await logger.async_log_success_event( + make_kwargs(user_id="user-1", end_user="cust"), None, None, None + ) + await logger.async_log_success_event(make_kwargs(end_user="cust"), None, None, None) + await logger.async_log_success_event(make_kwargs(), None, None, None) + + assert [record["user_uid"] for record in logger.log_queue] == [ + "dev@example.com", + "user-1", + "cust", + "unknown", + ] + + +@pytest.mark.asyncio +async def test_proxy_admin_placeholder_user_id_is_ignored(): + logger = make_logger() + await logger.async_log_success_event( + make_kwargs(user_id="default_user_id", end_user="dev@example.com"), + None, + None, + None, + ) + assert logger.log_queue[0]["user_uid"] == "dev@example.com" + + +@pytest.mark.asyncio +async def test_send_batch_posts_usage_mutation(): + mock_post = AsyncMock(return_value=make_response()) + logger = make_logger(mock_post=mock_post, graph="my-graph") + + await logger.async_log_success_event( + make_kwargs(startTime=1765000000.5, user_email="dev@example.com"), + None, + None, + None, + ) + await logger.async_send_batch() + + mock_post.assert_awaited_once() + call = mock_post.call_args + assert call.kwargs["url"] == "https://prod.api.faros.ai/graphs/my-graph/graphql" + assert call.kwargs["headers"]["authorization"] == "test-key" + + body = call.kwargs["json"] + assert body["query"] == USAGE_MUTATION + assert "insert_vcs_UserToolUsage" in body["query"] + usages = body["variables"]["usages"] + assert len(usages) == 1 + usage = usages[0] + assert usage["usedAt"] == "2025-12-06T05:46:40.500+00:00" + assert usage["origin"] == "litellm" + user_tool = usage["userTool"]["data"] + assert user_tool["tool"] == {"category": "LiteLLM"} + assert user_tool["user"]["data"] == { + "uid": "dev@example.com", + "source": "LiteLLM", + "origin": "litellm", + } + assert user_tool["user"]["on_conflict"]["constraint"] == "vcs_User_pkey" + assert usage["userTool"]["on_conflict"]["constraint"] == "vcs_UserTool_pkey" + json.dumps(body) + + +@pytest.mark.asyncio +async def test_send_batch_dedupes_rows_with_same_primary_key(): + mock_post = AsyncMock(return_value=make_response()) + logger = make_logger(mock_post=mock_post) + + for _ in range(2): + await logger.async_log_success_event( + make_kwargs(startTime=1765000000.5, user_id="user-1"), None, None, None + ) + await logger.async_log_success_event( + make_kwargs(startTime=1765000000.5, user_id="user-2"), None, None, None + ) + await logger.async_send_batch() + + usages = mock_post.call_args.kwargs["json"]["variables"]["usages"] + assert len(usages) == 2 + assert {u["userTool"]["data"]["user"]["data"]["uid"] for u in usages} == { + "user-1", + "user-2", + } + + +@pytest.mark.asyncio +async def test_graphql_errors_raise_and_preserve_queue(): + mock_post = AsyncMock( + return_value=make_response({"errors": [{"message": "unknown field"}]}) + ) + logger = make_logger(mock_post=mock_post) + + await logger.async_log_success_event( + make_kwargs(user_id="user-1"), None, None, None + ) + with pytest.raises(ValueError, match="unknown field"): + await logger.async_send_batch() + + await logger.flush_queue() + assert len(logger.log_queue) == 1 + + +@pytest.mark.asyncio +async def test_batch_size_triggers_flush(): + mock_post = AsyncMock(return_value=make_response()) + logger = make_logger(mock_post=mock_post, batch_size=2) + + await logger.async_log_success_event( + make_kwargs(startTime=1765000000.5, user_id="user-1"), None, None, None + ) + mock_post.assert_not_awaited() + await logger.async_log_success_event( + make_kwargs(startTime=1765000001.5, user_id="user-1"), None, None, None + ) + mock_post.assert_awaited_once() + assert logger.log_queue == [] + + +@pytest.mark.asyncio +async def test_event_without_standard_logging_object_is_skipped(): + logger = make_logger() + await logger.async_log_success_event({}, None, None, None) + assert logger.log_queue == [] + + +def test_faros_is_a_registered_callback(): + assert CustomLoggerRegistry.CALLBACK_CLASS_STR_TO_CLASS_TYPE["faros"] is FarosLogger + assert "faros" in litellm._known_custom_logger_compatible_callbacks + + +@pytest.mark.asyncio +async def test_init_custom_logger_compatible_class_returns_singleton(monkeypatch): + from litellm.litellm_core_utils.litellm_logging import ( + _init_custom_logger_compatible_class, + _in_memory_loggers, + ) + + monkeypatch.setenv("FAROS_API_KEY", "test-key") + created = _init_custom_logger_compatible_class("faros", None, None) + try: + assert isinstance(created, FarosLogger) + again = _init_custom_logger_compatible_class("faros", None, None) + assert again is created + finally: + _in_memory_loggers.remove(created)