phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 03:04:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

(refactor) - migrate router.deployment_callback_on_success to use StandardLoggingPayload (#7015)

Browse source 
* migrate deployment_callback_on_success to use SLP

* test_deployment_callback_on_success
This commit is contained in:
Ishaan Jaff 2024-12-03 20:29:25 -08:00 • committed by GitHub
parent 93c419868e
commit 54407b4592
3 changed files with 145 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

17
litellm/router.py
View file

@ -136,6 +136,7 @@ from litellm.types.router import (
from litellm.types.services import ServiceLoggerPayload, ServiceTypes
from litellm.types.utils import OPENAI_RESPONSE_HEADERS
from litellm.types.utils import ModelInfo as ModelMapInfo
from litellm.types.utils import StandardLoggingPayload
from litellm.utils import (
CustomStreamWrapper,
ModelResponse,
@ -3297,26 +3298,26 @@ class Router:
Track remaining tpm/rpm quota for model in model_list
"""
try:
standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get(
"standard_logging_object", None
)
if standard_logging_object is None:
raise ValueError("standard_logging_object is None")
if kwargs["litellm_params"].get("metadata") is None:
pass
else:
deployment_name = kwargs["litellm_params"]["metadata"].get(
"deployment", None
) # stable name - works for wildcard routes as well
model_group = kwargs["litellm_params"]["metadata"].get(
"model_group", None
)
model_info = kwargs["litellm_params"].get("model_info", {}) or {}
id = model_info.get("id", None)
model_group = standard_logging_object.get("model_group", None)
id = standard_logging_object.get("model_id", None)
if model_group is None or id is None:
return
elif isinstance(id, int):
id = str(id)
parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs)
_usage_obj = completion_response.get("usage")
total_tokens = _usage_obj.get("total_tokens", 0) if _usage_obj else 0
total_tokens: float = standard_logging_object.get("total_tokens", 0)
# ------------
# Setup values