feat(prometheus.py): initial working commit of passing team/key metadata as prometheus metrics

Closes LIT-1006
This commit is contained in:
Krrish Dholakia 2025-09-30 17:02:40 -07:00
parent 6ca7752381
commit d6800ee706
6 changed files with 93 additions and 38 deletions

View file

@ -794,9 +794,16 @@ class PrometheusLogger(CustomLogger):
output_tokens = standard_logging_payload["completion_tokens"]
tokens_used = standard_logging_payload["total_tokens"]
response_cost = standard_logging_payload["response_cost"]
_requester_metadata = standard_logging_payload["metadata"].get(
_requester_metadata: Optional[dict] = standard_logging_payload["metadata"].get(
"requester_metadata"
)
user_api_key_auth_metadata: Optional[dict] = standard_logging_payload[
"metadata"
].get("user_api_key_auth_metadata")
combined_metadata: Dict[str, Any] = {
**(_requester_metadata if _requester_metadata else {}),
**(user_api_key_auth_metadata if user_api_key_auth_metadata else {}),
}
if standard_logging_payload is not None and isinstance(
standard_logging_payload, dict
):
@ -828,8 +835,7 @@ class PrometheusLogger(CustomLogger):
exception_status=None,
exception_class=None,
custom_metadata_labels=get_custom_labels_from_metadata(
metadata=standard_logging_payload["metadata"].get("requester_metadata")
or {}
metadata=combined_metadata
),
route=standard_logging_payload["metadata"].get(
"user_api_key_request_route"

View file

@ -4019,6 +4019,7 @@ class StandardLoggingPayloadSetup:
usage_object=usage_object,
requester_custom_headers=None,
cold_storage_object_key=None,
user_api_key_auth_metadata=None,
)
if isinstance(metadata, dict):
# Filter the metadata dictionary to include only the specified keys
@ -4685,6 +4686,7 @@ def get_standard_logging_metadata(
requester_custom_headers=None,
user_api_key_request_route=None,
cold_storage_object_key=None,
user_api_key_auth_metadata=None,
)
if isinstance(metadata, dict):
# Update the clean_metadata with values from input metadata that match StandardLoggingMetadata fields

View file

@ -1,30 +1,9 @@
model_list:
- model_name: byok-fixed-gpt-4o-mini
- model_name: openai/gpt-4o
litellm_params:
model: openai/gpt-4o-mini
api_base: "https://webhook.site/2f385e05-00aa-402b-86d1-efc9261471a5"
api_key: dummy
- model_name: "byok-wildcard/*"
litellm_params:
model: openai/*
- model_name: xai-grok-3
litellm_params:
model: xai/grok-3
- model_name: hosted_vllm/whisper-v3
litellm_params:
model: hosted_vllm/whisper-v3
api_base: "https://webhook.site/2f385e05-00aa-402b-86d1-efc9261471a5"
api_key: dummy
mcp_servers:
github_mcp:
url: "https://api.githubcopilot.com/mcp"
auth_type: oauth2
authorization_url: https://github.com/login/oauth/authorize
token_url: https://github.com/login/oauth/access_token
client_id: os.environ/GITHUB_OAUTH_CLIENT_ID
client_secret: os.environ/GITHUB_OAUTH_CLIENT_SECRET
scopes: ["public_repo", "user:email"]
allowed_tools: ["list_tools"]
# disallowed_tools: ["repo_delete"]
model: openai/gpt-4o
api_key: os.environ/OPENAI_API_KEY
litellm_settings:
callbacks: ["prometheus"]
custom_prometheus_metadata_labels: ["metadata.initiative"]

View file

@ -3066,6 +3066,7 @@ LiteLLM_ManagementEndpoint_MetadataFields_Premium = [
"tags",
"team_member_key_duration",
"prompts",
"logging",
]

View file

@ -579,7 +579,12 @@ class LiteLLMProxyRequestSetup:
user_api_key_end_user_id=user_api_key_dict.end_user_id,
user_api_key_user_email=user_api_key_dict.user_email,
user_api_key_request_route=user_api_key_dict.request_route,
user_api_key_budget_reset_at=user_api_key_dict.budget_reset_at.isoformat() if user_api_key_dict.budget_reset_at else None,
user_api_key_budget_reset_at=(
user_api_key_dict.budget_reset_at.isoformat()
if user_api_key_dict.budget_reset_at
else None
),
user_api_key_auth_metadata=None,
)
return user_api_key_logged_metadata
@ -607,6 +612,35 @@ class LiteLLMProxyRequestSetup:
)
return data
@staticmethod
def add_management_endpoint_metadata_to_request_metadata(
data: dict,
management_endpoint_metadata: dict,
_metadata_variable_name: str,
) -> dict:
"""
Adds the `UserAPIKeyAuth` metadata to the request metadata.
ignore any sensitive fields like logging, api_key, etc.
"""
from litellm.proxy._types import (
LiteLLM_ManagementEndpoint_MetadataFields,
LiteLLM_ManagementEndpoint_MetadataFields_Premium,
)
# ignore any special fields
added_metadata = {}
for k, v in management_endpoint_metadata.items():
if k not in (
LiteLLM_ManagementEndpoint_MetadataFields_Premium
+ LiteLLM_ManagementEndpoint_MetadataFields
):
added_metadata[k] = v
data[_metadata_variable_name].setdefault(
"user_api_key_auth_metadata", {}
).update(added_metadata)
return data
@staticmethod
def add_key_level_controls(
key_metadata: Optional[dict], data: dict, _metadata_variable_name: str
@ -651,6 +685,13 @@ class LiteLLMProxyRequestSetup:
key_metadata["disable_fallbacks"], bool
):
data["disable_fallbacks"] = key_metadata["disable_fallbacks"]
## KEY-LEVEL METADATA
data = LiteLLMProxyRequestSetup.add_management_endpoint_metadata_to_request_metadata(
data=data,
management_endpoint_metadata=key_metadata,
_metadata_variable_name=_metadata_variable_name,
)
return data
@staticmethod
@ -889,6 +930,15 @@ async def add_litellm_data_to_request( # noqa: PLR0915
"spend_logs_metadata"
]
## TEAM-LEVEL METADATA
data = (
LiteLLMProxyRequestSetup.add_management_endpoint_metadata_to_request_metadata(
data=data,
management_endpoint_metadata=team_metadata,
_metadata_variable_name=_metadata_variable_name,
)
)
# Team spend, budget - used by prometheus.py
data[_metadata_variable_name][
"user_api_key_team_max_budget"

View file

@ -123,12 +123,18 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
max_output_tokens: Required[Optional[int]]
input_cost_per_token: Required[float]
input_cost_per_token_flex: Optional[float] # OpenAI flex service tier pricing
input_cost_per_token_priority: Optional[float] # OpenAI priority service tier pricing
input_cost_per_token_priority: Optional[
float
] # OpenAI priority service tier pricing
cache_creation_input_token_cost: Optional[float]
cache_creation_input_token_cost_above_1hr: Optional[float]
cache_read_input_token_cost: Optional[float]
cache_read_input_token_cost_flex: Optional[float] # OpenAI flex service tier pricing
cache_read_input_token_cost_priority: Optional[float] # OpenAI priority service tier pricing
cache_read_input_token_cost_flex: Optional[
float
] # OpenAI flex service tier pricing
cache_read_input_token_cost_priority: Optional[
float
] # OpenAI priority service tier pricing
input_cost_per_character: Optional[float] # only for vertex ai models
input_cost_per_audio_token: Optional[float]
input_cost_per_token_above_128k_tokens: Optional[float] # only for vertex ai models
@ -147,7 +153,9 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
output_cost_per_token_batches: Optional[float]
output_cost_per_token: Required[float]
output_cost_per_token_flex: Optional[float] # OpenAI flex service tier pricing
output_cost_per_token_priority: Optional[float] # OpenAI priority service tier pricing
output_cost_per_token_priority: Optional[
float
] # OpenAI priority service tier pricing
output_cost_per_character: Optional[float] # only for vertex ai models
output_cost_per_audio_token: Optional[float]
output_cost_per_token_above_128k_tokens: Optional[
@ -1856,6 +1864,7 @@ class StandardLoggingUserAPIKeyMetadata(TypedDict):
user_api_key_team_alias: Optional[str]
user_api_key_end_user_id: Optional[str]
user_api_key_request_route: Optional[str]
user_api_key_auth_metadata: Optional[Dict[str, str]]
class StandardLoggingMCPToolCall(TypedDict, total=False):
@ -2059,10 +2068,12 @@ class StandardLoggingGuardrailInformation(TypedDict, total=False):
StandardLoggingPayloadStatus = Literal["success", "failure"]
class CachingDetails(TypedDict):
"""
Track all caching related metrics, fields for a given request
"""
cache_hit: Optional[bool]
"""
Whether the request hit the cache
@ -2072,12 +2083,16 @@ class CachingDetails(TypedDict):
Duration for reading from cache
"""
class CostBreakdown(TypedDict):
"""
Detailed cost breakdown for a request
"""
input_cost: float # Cost of input/prompt tokens
output_cost: float # Cost of output/completion tokens (includes reasoning if applicable)
output_cost: (
float # Cost of output/completion tokens (includes reasoning if applicable)
)
total_cost: float # Total cost (input + output + tool usage)
tool_usage_cost: float # Cost of usage of built-in tools
@ -2616,6 +2631,7 @@ class SpecialEnums(Enum):
class ServiceTier(Enum):
"""Enum for service tier types used in cost calculations."""
FLEX = "flex"
PRIORITY = "priority"
@ -2662,13 +2678,14 @@ CostResponseTypes = Union[
class PriorityReservationSettings(BaseModel):
"""
Settings for priority-based rate limiting reservation.
Defines what priority to assign to keys without explicit priority metadata.
The priority_reservation mapping is configured separately via litellm.priority_reservation.
"""
default_priority: float = Field(
default=0.5,
description="Priority level to assign to API keys without explicit priority metadata. Should match a key in litellm.priority_reservation."
description="Priority level to assign to API keys without explicit priority metadata. Should match a key in litellm.priority_reservation.",
)
model_config = ConfigDict(protected_namespaces=())