prom - track spend for custom_llm_provider

This commit is contained in:
Ishaan Jaff 2024-11-19 16:25:20 -08:00
parent c2994f3a32
commit e22d995ba8
2 changed files with 10 additions and 1 deletions

View file

@ -15,6 +15,7 @@ import requests # type: ignore
import litellm import litellm
from litellm._logging import print_verbose, verbose_logger from litellm._logging import print_verbose, verbose_logger
from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadAccessors
from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.integrations.prometheus import * from litellm.types.integrations.prometheus import *
from litellm.types.utils import StandardLoggingPayload from litellm.types.utils import StandardLoggingPayload
@ -327,6 +328,7 @@ class PrometheusLogger(CustomLogger):
"team", "team",
"team_alias", "team_alias",
"user", "user",
CUSTOM_LLM_PROVIDER,
], ],
) )
@ -369,6 +371,9 @@ class PrometheusLogger(CustomLogger):
output_tokens = standard_logging_payload["completion_tokens"] output_tokens = standard_logging_payload["completion_tokens"]
tokens_used = standard_logging_payload["total_tokens"] tokens_used = standard_logging_payload["total_tokens"]
response_cost = standard_logging_payload["response_cost"] response_cost = standard_logging_payload["response_cost"]
custom_llm_provider = StandardLoggingPayloadAccessors.get_custom_llm_provider_from_standard_logging_payload(
standard_logging_payload
)
print_verbose( 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}" 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 +398,7 @@ class PrometheusLogger(CustomLogger):
user_api_team_alias=user_api_team_alias, user_api_team_alias=user_api_team_alias,
user_id=user_id, user_id=user_id,
response_cost=response_cost, response_cost=response_cost,
custom_llm_provider=custom_llm_provider,
) )
# input, output, total token metrics # input, output, total token metrics
@ -535,6 +541,7 @@ class PrometheusLogger(CustomLogger):
user_api_team_alias: Optional[str], user_api_team_alias: Optional[str],
user_id: Optional[str], user_id: Optional[str],
response_cost: float, response_cost: float,
custom_llm_provider: Optional[str],
): ):
self.litellm_requests_metric.labels( self.litellm_requests_metric.labels(
end_user_id, end_user_id,
@ -544,6 +551,7 @@ class PrometheusLogger(CustomLogger):
user_api_team, user_api_team,
user_api_team_alias, user_api_team_alias,
user_id, user_id,
custom_llm_provider,
).inc() ).inc()
self.litellm_spend_metric.labels( self.litellm_spend_metric.labels(
end_user_id, end_user_id,
@ -553,6 +561,7 @@ class PrometheusLogger(CustomLogger):
user_api_team, user_api_team,
user_api_team_alias, user_api_team_alias,
user_id, user_id,
custom_llm_provider,
).inc(response_cost) ).inc(response_cost)
def _set_virtual_key_rate_limit_metrics( def _set_virtual_key_rate_limit_metrics(

View file

@ -2730,7 +2730,7 @@ class StandardLoggingPayloadAccessors:
@staticmethod @staticmethod
def get_custom_llm_provider_from_standard_logging_payload( def get_custom_llm_provider_from_standard_logging_payload(
standard_logging_payload: Optional[StandardLoggingPayload], standard_logging_payload: Optional[Union[StandardLoggingPayload, dict]],
) -> Optional[str]: ) -> Optional[str]:
""" """
Accessor method to safely get custom_llm_provider from standard_logging_payload Accessor method to safely get custom_llm_provider from standard_logging_payload