refactor: extract helpers to stay within C901 complexity budget

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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,

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)