diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index cbeb4d336..0b6cd8b3b 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -208,7 +208,7 @@ class PrometheusLogger(CustomLogger): "LLM Deployment Analytics - remaining requests for model, returned from LLM API Provider", labelnames=[ "model_group", - "api_provider", + CUSTOM_LLM_PROVIDER, "api_base", "litellm_model_name", "hashed_api_key", @@ -221,7 +221,7 @@ class PrometheusLogger(CustomLogger): "remaining tokens for model, returned from LLM API Provider", labelnames=[ "model_group", - "api_provider", + CUSTOM_LLM_PROVIDER, "api_base", "litellm_model_name", "hashed_api_key", @@ -233,7 +233,7 @@ class PrometheusLogger(CustomLogger): "litellm_model_name", "model_id", "api_base", - "api_provider", + CUSTOM_LLM_PROVIDER, ] team_and_key_labels = [ "hashed_api_key", @@ -327,6 +327,7 @@ class PrometheusLogger(CustomLogger): "team", "team_alias", "user", + CUSTOM_LLM_PROVIDER, ], ) @@ -336,6 +337,9 @@ class PrometheusLogger(CustomLogger): async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): # Define prometheus client + from litellm.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadAccessors, + ) from litellm.types.utils import StandardLoggingPayload verbose_logger.debug( @@ -369,6 +373,9 @@ class PrometheusLogger(CustomLogger): output_tokens = standard_logging_payload["completion_tokens"] tokens_used = standard_logging_payload["total_tokens"] response_cost = standard_logging_payload["response_cost"] + custom_llm_provider = StandardLoggingPayloadAccessors.get_custom_llm_provider_from_standard_logging_payload( + standard_logging_payload + ) print_verbose( f"inside track_prometheus_metrics, model {model}, response_cost {response_cost}, tokens_used {tokens_used}, end_user_id {end_user_id}, user_api_key {user_api_key}" @@ -393,6 +400,7 @@ class PrometheusLogger(CustomLogger): user_api_team_alias=user_api_team_alias, user_id=user_id, response_cost=response_cost, + custom_llm_provider=custom_llm_provider, ) # input, output, total token metrics @@ -535,6 +543,7 @@ class PrometheusLogger(CustomLogger): user_api_team_alias: Optional[str], user_id: Optional[str], response_cost: float, + custom_llm_provider: Optional[str], ): self.litellm_requests_metric.labels( end_user_id, @@ -544,6 +553,7 @@ class PrometheusLogger(CustomLogger): user_api_team, user_api_team_alias, user_id, + custom_llm_provider, ).inc() self.litellm_spend_metric.labels( end_user_id, @@ -553,6 +563,7 @@ class PrometheusLogger(CustomLogger): user_api_team, user_api_team_alias, user_id, + custom_llm_provider, ).inc(response_cost) def _set_virtual_key_rate_limit_metrics( @@ -790,7 +801,7 @@ class PrometheusLogger(CustomLogger): """ log these labels - ["litellm_model_name", "model_id", "api_base", "api_provider"] + ["litellm_model_name", "model_id", "api_base", CUSTOM_LLM_PROVIDER] """ self.set_deployment_partial_outage( litellm_model_name=litellm_model_name, @@ -882,7 +893,7 @@ class PrometheusLogger(CustomLogger): if remaining_requests: """ "model_group", - "api_provider", + CUSTOM_LLM_PROVIDER, "api_base", "litellm_model_name" """ @@ -907,7 +918,7 @@ class PrometheusLogger(CustomLogger): """ log these labels - ["litellm_model_name", "requested_model", model_id", "api_base", "api_provider"] + ["litellm_model_name", "requested_model", model_id", "api_base", CUSTOM_LLM_PROVIDER] """ self.set_deployment_healthy( litellm_model_name=litellm_model_name, diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 69d6adca4..99834a0d5 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2359,6 +2359,7 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 _in_memory_loggers.append(_mlflow_logger) return _mlflow_logger # type: ignore + def get_custom_logger_compatible_class( logging_integration: litellm._custom_logger_compatible_callbacks_literal, ) -> Optional[CustomLogger]: @@ -2719,6 +2720,31 @@ class StandardLoggingPayloadSetup: return clean_hidden_params +class StandardLoggingPayloadAccessors: + """ + Accessor methods for StandardLoggingPayload + + Class that allows easily reading fields from StandardLoggingPayload + + """ + + @staticmethod + def get_custom_llm_provider_from_standard_logging_payload( + standard_logging_payload: Optional[Union[StandardLoggingPayload, dict]], + ) -> Optional[str]: + """ + Accessor method to safely get custom_llm_provider from standard_logging_payload + """ + if standard_logging_payload is None: + return None + model_map_information = ( + standard_logging_payload.get("model_map_information") or {} + ) + model_map_value = model_map_information.get("model_map_value") or {} + custom_llm_provider = model_map_value.get("litellm_provider") + return custom_llm_provider + + def get_standard_logging_object_payload( kwargs: Optional[dict], init_response_obj: Union[Any, BaseModel, dict], diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index f5851ded9..09fe3c2fa 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1773,6 +1773,7 @@ class SpendLogsPayload(TypedDict): model_id: Optional[str] model_group: Optional[str] api_base: str + custom_llm_provider: Optional[str] user: str metadata: str # json str cache_hit: str diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 64045999c..7f717e6c8 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -192,6 +192,7 @@ model LiteLLM_SpendLogs { model_id String? @default("") // the model id stored in proxy model db model_group String? @default("") // public model_name / model_group api_base String? @default("") + custom_llm_provider String? @default("") // openai, vertex_ai etc user String? @default("") metadata Json? @default("{}") cache_hit String? @default("") diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 48924d521..e6a79b05e 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -10,8 +10,10 @@ from pydantic import BaseModel import litellm from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadAccessors from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload from litellm.proxy.utils import PrismaClient, hash_token +from litellm.types.utils import StandardLoggingPayload def _is_master_key(api_key: str, _master_key: Optional[str]) -> bool: @@ -49,6 +51,10 @@ def get_logging_payload( response_obj = {} # standardize this function to be used across, s3, dynamoDB, langfuse logging litellm_params = kwargs.get("litellm_params", {}) + standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get( + "standard_logging_object", None + ) + metadata = ( litellm_params.get("metadata", {}) or {} ) # if litellm_params['metadata'] == None @@ -150,6 +156,9 @@ def get_logging_payload( request_tags=request_tags, end_user=end_user_id or "", api_base=litellm_params.get("api_base", ""), + custom_llm_provider=StandardLoggingPayloadAccessors.get_custom_llm_provider_from_standard_logging_payload( + standard_logging_payload + ), model_group=_model_group, model_id=_model_id, requester_ip_address=clean_metadata.get("requester_ip_address", None), diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index d09ed9670..dc77340f9 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -1,4 +1,5 @@ REQUESTED_MODEL = "requested_model" +CUSTOM_LLM_PROVIDER = "api_provider" EXCEPTION_STATUS = "exception_status" EXCEPTION_CLASS = "exception_class" EXCEPTION_LABELS = [EXCEPTION_STATUS, EXCEPTION_CLASS] diff --git a/schema.prisma b/schema.prisma index 64045999c..7f717e6c8 100644 --- a/schema.prisma +++ b/schema.prisma @@ -192,6 +192,7 @@ model LiteLLM_SpendLogs { model_id String? @default("") // the model id stored in proxy model db model_group String? @default("") // public model_name / model_group api_base String? @default("") + custom_llm_provider String? @default("") // openai, vertex_ai etc user String? @default("") metadata Json? @default("{}") cache_hit String? @default("") diff --git a/tests/local_testing/test_spend_logs.py b/tests/local_testing/test_spend_logs.py index 926f4b5ad..fe292c12d 100644 --- a/tests/local_testing/test_spend_logs.py +++ b/tests/local_testing/test_spend_logs.py @@ -28,12 +28,25 @@ import litellm from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload from litellm.proxy.utils import SpendLogsMetadata, SpendLogsPayload # noqa: E402 +from litellm.types.utils import ( + StandardLoggingPayload, + StandardLoggingModelInformation, + StandardLoggingMetadata, + StandardLoggingHiddenParams, + ModelInfo, +) + def test_spend_logs_payload(): """ Ensure only expected values are logged in spend logs payload. """ + standard_logging_payload = _create_standard_logging_payload() + standard_logging_payload["model_map_information"]["model_map_value"][ + "litellm_provider" + ] = "very-obscure-name" + input_args: dict = { "kwargs": { "model": "chatgpt-v-2", @@ -47,6 +60,7 @@ def test_spend_logs_payload(): "user": "116544810872468347480", "extra_body": {}, }, + "standard_logging_object": standard_logging_payload, "litellm_params": { "acompletion": True, "api_key": "23c217a5b59f41b6b7a198017f4792f2", @@ -205,6 +219,9 @@ def test_spend_logs_payload(): assert ( payload["request_tags"] == '["model-anthropic-claude-v2.1", "app-ishaan-prod"]' ) + print("payload['custom_llm_provider']", payload["custom_llm_provider"]) + # Ensures custom llm provider is logged + read from standard logging payload + assert payload["custom_llm_provider"] == "very-obscure-name" def test_spend_logs_payload_whisper(): @@ -292,3 +309,61 @@ def test_spend_logs_payload_whisper(): assert payload["call_type"] == "atranscription" assert payload["spend"] == 0.00023398580000000003 + + +def _create_standard_logging_payload() -> StandardLoggingPayload: + """ + helper function that creates a standard logging payload for testing + + in the test you can edit the values in SLP that you would need + """ + return StandardLoggingPayload( + id="test_id", + type="test_id", + call_type="completion", + response_cost=0.1, + response_cost_failure_debug_info=None, + status="success", + total_tokens=30, + prompt_tokens=20, + completion_tokens=10, + startTime=1234567890.0, + endTime=1234567891.0, + completionStartTime=1234567890.5, + model_map_information=StandardLoggingModelInformation( + model_map_key="gpt-3.5-turbo", + model_map_value=ModelInfo(litellm_provider="azure"), + ), + model="gpt-3.5-turbo", + model_id="model-123", + model_group="openai-gpt", + api_base="https://api.openai.com", + metadata=StandardLoggingMetadata( + user_api_key_hash="test_hash", + user_api_key_org_id=None, + user_api_key_alias="test_alias", + user_api_key_team_id="test_team", + user_api_key_user_id="test_user", + user_api_key_team_alias="test_team_alias", + spend_logs_metadata=None, + requester_ip_address="127.0.0.1", + requester_metadata=None, + ), + cache_hit=False, + cache_key=None, + saved_cache_cost=0.0, + request_tags=[], + end_user=None, + requester_ip_address="127.0.0.1", + messages=[{"role": "user", "content": "Hello, world!"}], + response={"choices": [{"message": {"content": "Hi there!"}}]}, + error_str=None, + model_parameters={"stream": True}, + hidden_params=StandardLoggingHiddenParams( + model_id="model-123", + cache_key=None, + api_base="https://api.openai.com", + response_cost="0.1", + additional_headers=None, + ), + ) diff --git a/tests/logging_callback_tests/test_prometheus_unit_tests.py b/tests/logging_callback_tests/test_prometheus_unit_tests.py index 494f83a65..e3bd6510a 100644 --- a/tests/logging_callback_tests/test_prometheus_unit_tests.py +++ b/tests/logging_callback_tests/test_prometheus_unit_tests.py @@ -1,3 +1,7 @@ +""" +Unit tests for prometheus logging callback +""" + import io import os import sys @@ -333,15 +337,30 @@ def test_increment_top_level_request_and_spend_metrics(prometheus_logger): user_api_team_alias="team_alias1", user_id="user1", response_cost=0.1, + custom_llm_provider="openai", ) prometheus_logger.litellm_requests_metric.labels.assert_called_once_with( - "user1", "key1", "alias1", "gpt-3.5-turbo", "team1", "team_alias1", "user1" + "user1", + "key1", + "alias1", + "gpt-3.5-turbo", + "team1", + "team_alias1", + "user1", + "openai", ) prometheus_logger.litellm_requests_metric.labels().inc.assert_called_once() prometheus_logger.litellm_spend_metric.labels.assert_called_once_with( - "user1", "key1", "alias1", "gpt-3.5-turbo", "team1", "team_alias1", "user1" + "user1", + "key1", + "alias1", + "gpt-3.5-turbo", + "team1", + "team_alias1", + "user1", + "openai", ) prometheus_logger.litellm_spend_metric.labels().inc.assert_called_once_with(0.1)