Merge 96bd1ea9cc into e1d16f51d1

litellm/router.py

@@ -2912,6 +2912,7 @@ class Router:
         cast(_HiddenParamsHost, fallback_item)._hidden_params = {
             **item_hidden_params,
             **fallback_hidden_params,
+            "response_cost": item_hidden_params.get("response_cost"),
             "additional_headers": {**item_headers, **fallback_headers},
         }
 

tests/integration/_support/wire.py

@@ -28,6 +28,7 @@ class Reply:
     chunks: tuple[bytes, ...] | None = None
     abort_after: int | None = None
     gate_after_first: threading.Event | None = None
+    gate_timeout_seconds: float = 5.0
     pause_between_chunks: float = 0
     headers: Mapping[str, str] = MappingProxyType({})
     drop_connection: bool = False
@@ -109,7 +110,9 @@ def wire_server(
                         self.wfile.write(b"%x\r\n%s\r\n" % (len(chunk), chunk))
                         self.wfile.flush()
                         if index == 0 and reply.gate_after_first is not None:
-                            assert reply.gate_after_first.wait(timeout=5), "Stream barrier was never released"
+                            assert reply.gate_after_first.wait(timeout=reply.gate_timeout_seconds), (
+                                "Stream barrier was never released"
+                            )
                         if reply.pause_between_chunks and index + 1 < len(reply.chunks):
                             time.sleep(reply.pause_between_chunks)
                     else:

tests/integration/streaming/fallback_cost_recorder.py

@@ -0,0 +1,44 @@
+import json
+import os
+from collections.abc import AsyncGenerator
+from itertools import count
+from typing import Final
+
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.types.utils import ModelResponseStream
+
+
+class FallbackCostRecorder(CustomLogger):
+    async def async_post_call_streaming_iterator_hook(
+        self,
+        user_api_key_dict: UserAPIKeyAuth,
+        response: AsyncGenerator[ModelResponseStream, None],
+        request_data: dict[str, object],
+    ) -> AsyncGenerator[ModelResponseStream, None]:
+        log_path: Final = os.environ["LITELLM_FALLBACK_COST_LOG"]
+        chunk_indices: Final = count()
+        async for item in response:
+            hidden_params: Final = getattr(item, "_hidden_params", None)
+            usage: Final = getattr(item, "usage", None)
+            event_type: Final = getattr(item, "type", None)
+            record: Final = {
+                "request_model": request_data.get("model"),
+                "item_type": type(item).__name__,
+                "worker_pid": os.getpid(),
+                "response_cost": hidden_params.get("response_cost") if isinstance(hidden_params, dict) else None,
+                "usage_cost": getattr(
+                    usage,
+                    "cost",
+                    getattr(getattr(getattr(item, "response", None), "usage", None), "cost", None),
+                ),
+                "event_type": event_type,
+                "chunk_index": next(chunk_indices) if event_type is None else None,
+                "has_hidden_params": hasattr(item, "_hidden_params"),
+            }
+            with open(log_path, "a", encoding="utf-8") as log:
+                log.write(json.dumps(record, default=str) + "\n")
+            yield item
+
+
+proxy_handler_instance = FallbackCostRecorder()

tests/unit/test_router/test_router.py

@@ -3035,6 +3035,36 @@ def test_completion_streaming_iterator_preserves_response_headers():
     assert result._hidden_params["additional_headers"]["llm_provider-x-request-id"] == "req-provider-sync"
 
 
+@pytest.mark.parametrize("response_cost", (None, 0.00015))
+def test_apply_fallback_hidden_params_to_item_keeps_the_chunks_own_response_cost(
+    response_cost: float | None,
+) -> None:
+    chunk: Final = litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "chunk"}}])
+    chunk._hidden_params = {
+        "response_cost": response_cost,
+        "model_id": "failed-deployment",
+        "additional_headers": {"x-a": "1"},
+    }
+
+    Router._apply_fallback_hidden_params_to_item(
+        chunk,
+        (
+            {
+                "response_cost": 0.0,
+                "model_id": "fallback-deployment",
+                "additional_headers": {"x-b": "2"},
+            },
+            {"x-b": "2"},
+        ),
+    )
+
+    assert chunk._hidden_params == {
+        "response_cost": response_cost,
+        "model_id": "fallback-deployment",
+        "additional_headers": {"x-a": "1", "x-b": "2"},
+    }
+
+
 def test_adopt_fallback_response_headers_replaces_rather_than_merges():
     """LIT-6767: direct unit for FallbackAwareStreamWrapper.adopt_fallback_response_headers.