From 507485de67b2c5c38496b6ed1e734b6574094f31 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 1 Jul 2026 04:15:24 +0000 Subject: [PATCH] refactor: extract helpers to stay within C901 complexity budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../spend_tracking/spend_tracking_utils.py | 36 +++++++++------ litellm/proxy/utils.py | 44 ++++++++----------- 2 files changed, 40 insertions(+), 40 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index f55358c8665..1ec65c5456b 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -224,6 +224,20 @@ def _extract_usage_for_ocr_call(response_obj: Any, response_obj_dict: dict) -> d return {} +def _extract_standard_logging_payload_fields( + slp: Optional[StandardLoggingPayload], +) -> tuple[str, str, str, str, str]: + if slp is None: + return ("", "", "", "", "") + return ( + slp.get("call_type", "") or "", + slp.get("model_id", "") or "", + slp.get("model_group", "") or "", + slp.get("api_base", "") or "", + slp.get("custom_llm_provider", "") or "", + ) + + def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogsPayload: if kwargs is None: kwargs = {} @@ -267,8 +281,10 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs usage = _combined_usage.model_dump() standard_logging_payload = cast(Optional[StandardLoggingPayload], kwargs.get("standard_logging_object", None)) - if not call_type and standard_logging_payload is not None: - call_type = standard_logging_payload.get("call_type", "") + _slp_call_type, _slp_model_id, _slp_model_group, _slp_api_base, _slp_custom_llm_provider = ( + _extract_standard_logging_payload_fields(standard_logging_payload) + ) + call_type = call_type or _slp_call_type id = get_spend_logs_id(call_type or "acompletion", response_obj_dict, kwargs) @@ -301,13 +317,8 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs ): # use 'tags' from standard logging payload instead request_tags = safe_dumps(standard_logging_payload["request_tags"]) - _model_id = metadata.get("model_info", {}).get("id", "") - _model_group = metadata.get("model_group", "") - if standard_logging_payload is not None: - if not _model_id: - _model_id = standard_logging_payload.get("model_id", "") or "" - if not _model_group: - _model_group = standard_logging_payload.get("model_group", "") or "" + _model_id = metadata.get("model_info", {}).get("id", "") or _slp_model_id + _model_group = metadata.get("model_group", "") or _slp_model_group # Extract overhead from hidden_params if available litellm_overhead_time_ms = None @@ -391,9 +402,7 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs # Extract agent_id for A2A requests (set directly on model_call_details) agent_id: Optional[str] = kwargs.get("agent_id") or metadata.get("agent_id") - custom_llm_provider = kwargs.get("custom_llm_provider") - if not custom_llm_provider and standard_logging_payload is not None: - custom_llm_provider = standard_logging_payload.get("custom_llm_provider", "") or "" + custom_llm_provider = kwargs.get("custom_llm_provider") or _slp_custom_llm_provider raw_model = cast(str, kwargs.get("model") or "") model_name = reconstruct_model_name(raw_model, custom_llm_provider, metadata or {}) @@ -418,8 +427,7 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs completion_tokens=usage.get("completion_tokens", standard_logging_completion_tokens), request_tags=request_tags, end_user=end_user_id or "", - api_base=litellm_params.get("api_base", "") - or (standard_logging_payload.get("api_base", "") if standard_logging_payload is not None else ""), + api_base=litellm_params.get("api_base", "") or _slp_api_base, model_group=_model_group, model_id=_model_id, mcp_namespaced_tool_name=mcp_namespaced_tool_name, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 407eb5e8484..037c3257290 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -396,6 +396,23 @@ class _CallbackCapabilities: resolved_callbacks: Tuple[Any, ...] = field(default_factory=tuple) +def _lift_logging_obj_fields_to_request_data(request_data: dict) -> None: + _logging_obj = request_data.get("litellm_logging_obj") + if _logging_obj is None: + return + _model_call_details = getattr(_logging_obj, "model_call_details", {}) + _first_handoff = _model_call_details.get("first_api_call_start_time") + if _first_handoff is not None: + request_data["first_api_call_start_time"] = _first_handoff + _recovered_usage = _model_call_details.get("combined_usage_object") + if _recovered_usage is not None: + request_data["combined_usage_object"] = _recovered_usage + request_data["response_cost"] = _model_call_details.get("response_cost") + _slo = _model_call_details.get("standard_logging_object") + if _slo is not None and not request_data.get("standard_logging_object"): + request_data["standard_logging_object"] = _slo + + class ProxyLogging: """ Logging/Custom Handlers for proxy. @@ -2038,32 +2055,7 @@ class ProxyLogging: original_exception=original_exception, ) - # Lift the first-handoff instant onto request_data (top-level - # internal key, not metadata) so failure-path callbacks can still - # compute preprocessing latency after the logging object is popped. - _logging_obj = request_data.get("litellm_logging_obj") - if _logging_obj is not None: - _model_call_details = getattr(_logging_obj, "model_call_details", {}) - _first_handoff = _model_call_details.get("first_api_call_start_time") - if _first_handoff is not None: - request_data["first_api_call_start_time"] = _first_handoff - - # A stream that broke mid-flight still billed the provider for the - # chunks already delivered; the streaming handler stashes that - # recovered usage and cost here. Lift them onto request_data so the - # failure-path spend callbacks (which run after the logging object - # is popped) record the real partial spend instead of zero. - _recovered_usage = _model_call_details.get("combined_usage_object") - if _recovered_usage is not None: - request_data["combined_usage_object"] = _recovered_usage - request_data["response_cost"] = _model_call_details.get("response_cost") - - # Lift standard_logging_object so failure-path spend tracking can - # read model_id, model_group, and other fields that are only - # available on the logging object's model_call_details. - _slo = _model_call_details.get("standard_logging_object") - if _slo is not None and not request_data.get("standard_logging_object"): - request_data["standard_logging_object"] = _slo + _lift_logging_obj_fields_to_request_data(request_data) # Remove before callbacks iterate — not serialisable request_data.pop("litellm_logging_obj", None)