merge: sync with PR base litellm_internal_staging (936e07b0)

basedpyright-code-budget.json

@@ -57,7 +57,7 @@
     "limit": 5659
   },
   "reportMissingTypeArgument": {
-    "limit": 15484
+    "limit": 15482
   },
   "reportMissingTypeStubs": {
     "limit": 40
@@ -105,13 +105,13 @@
     "limit": 109
   },
   "reportUnknownMemberType": {
-    "limit": 38782
+    "limit": 38779
   },
   "reportUnknownParameterType": {
-    "limit": 19829
+    "limit": 19827
   },
   "reportUnknownVariableType": {
-    "limit": 30349
+    "limit": 30348
   },
   "reportUnnecessaryCast": {
     "limit": 117

litellm-proxy-extras/litellm_proxy_extras/utils.py

@@ -470,7 +470,7 @@ class ProxyExtrasDBManager:
         ProxyExtrasDBManager._mark_migrations_applied(migrations_dir)
 
     @staticmethod
-    def _mark_migrations_applied(migrations_dir: str):
+    def _mark_migrations_applied(migrations_dir: str) -> None:
         migration_names = ProxyExtrasDBManager._get_migration_names(migrations_dir)
         logger.info(f"Resolving {len(migration_names)} migrations")
         for migration_name in migration_names:

litellm/integrations/SlackAlerting/slack_alerting.py

@@ -1447,10 +1447,8 @@ Model Info:
 
         from datetime import datetime
 
-        # Get the current timestamp
         current_time: Final = datetime.now().strftime("%H:%M:%S")
         _proxy_base_url: Final = os.getenv("PROXY_BASE_URL", None)
-        # Use .name if it's an enum, otherwise use as is
         alert_type_name: Final = getattr(alert_type, "name", alert_type)
         alert_type_formatted: Final = f"Alert type: `{alert_type_name}`"
         if alert_type == "daily_reports" or alert_type == "new_model_added":

litellm/litellm_core_utils/streaming_handler.py

@@ -819,7 +819,9 @@ class CustomStreamWrapper:
         except Exception as e:
             raise e
 
-    def model_response_creator(self, chunk: dict | None = None, hidden_params: Mapping[str, object] | None = None):
+    def model_response_creator(
+        self, chunk: dict | None = None, hidden_params: Mapping[str, object] | None = None
+    ) -> ModelResponseStream:
         _model: Final = self._cached_model_name
         _logging_obj_llm_provider: Final = self._cached_logging_llm_provider
 

litellm/proxy/_lazy_openapi_snapshot.py

@@ -126,7 +126,6 @@ def _feature_fragment(app: "FastAPI", feat: "LazyFeature", used_operation_ids: s
     full: Final = get_openapi(title=app.title, version=app.version, routes=feat_routes)
     paths: Final = full.get("paths", {})
     _normalize_operation_ids(paths)
-    # Group all of a feature's routes under one tag.
     for path_ops in paths.values():
         for method, op in path_ops.items():
             if isinstance(op, dict):

litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py

@@ -385,7 +385,7 @@ class CrowdStrikeAIDRHandler(CustomGuardrail):
         return [_extract_text_from_message(msg) for msg in tail]
 
     async def _call_or_fail_open(
-        self, payload: dict[str, Any], hook_name: str, request_data: dict
+        self, payload: dict[str, Any], hook_name: str, request_data: dict[str, object]
     ) -> _GuardChatCompletionsResult:
         start_time: Final = time.time()
         try:
@@ -421,7 +421,7 @@ class CrowdStrikeAIDRHandler(CustomGuardrail):
         structured_messages: list[AllMessageValues],
         guard_output: _GuardInput,
         sent_indices: tuple[int, ...],
-        request_data: dict,
+        request_data: dict[str, object],
     ) -> list[AllMessageValues] | None:
         if effective_skip_system_message_for_guardrail(self) or effective_skip_tool_message_for_guardrail(self):
             request_messages: Final = request_data.get("messages")

ruff-strict-budget.json

@@ -9,7 +9,7 @@
     "limit": 827
   },
   "ANN201": {
-    "limit": 2012
+    "limit": 2011
   },
   "ANN202": {
     "limit": 847

tests/test_litellm/rag/test_main.py

@@ -18,9 +18,24 @@ import pytest
 import litellm
 from litellm._internal_context import is_internal_call
 from litellm.integrations.custom_logger import CustomLogger
+from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
 from litellm.types.utils import CallTypes, ModelResponse
 
 
+async def _drain_logging_worker() -> None:
+    """Run every queued logging task to completion on the current event loop.
+
+    The success event is delivered through the fire-and-forget GLOBAL_LOGGING_WORKER
+    singleton, whose queue survives across tests. start() rebinds any tasks left over
+    from a previous test's event loop onto the current one, and flush() waits until
+    the queue is fully processed, so tests neither miss their own event nor observe
+    a neighbour's
+    """
+    await asyncio.sleep(0)
+    GLOBAL_LOGGING_WORKER.start()
+    await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0)
+
+
 class RecordingLogger(CustomLogger):
     def __init__(self):
         super().__init__()
@@ -39,6 +54,7 @@ async def test_aquery_single_billing_event_carries_completion_usage_and_cost(use
     not the vector store search response. The proxy always passes a router, so
     both the router and non-router completion branches are pinned.
     """
+    await _drain_logging_worker()
     recording_logger = RecordingLogger()
     original_callbacks = litellm.callbacks
     litellm.callbacks = [recording_logger]
@@ -66,11 +82,7 @@ async def test_aquery_single_billing_event_carries_completion_usage_and_cost(use
         assert isinstance(response, ModelResponse)
         assert is_internal_call.get() is False
 
-        for _ in range(50):
-            if recording_logger.success_events:
-                break
-            await asyncio.sleep(0.1)
-        await asyncio.sleep(0.5)
+        await _drain_logging_worker()
     finally:
         litellm.callbacks = original_callbacks
 
@@ -102,6 +114,8 @@ async def test_aquery_response_hidden_params_carry_completion_cost():
         mock_response="hi there",
     )
 
+    await _drain_logging_worker()
+
     assert isinstance(response, ModelResponse)
     response_cost = response._hidden_params.get("response_cost")
     assert response_cost is not None
@@ -115,6 +129,7 @@ async def test_aquery_billed_cost_includes_priced_vector_store_search():
     that cost must be folded into the aquery billing instead of being dropped
     with the suppressed sub-call event.
     """
+    await _drain_logging_worker()
     recording_logger = RecordingLogger()
     original_callbacks = litellm.callbacks
     litellm.callbacks = [recording_logger]
@@ -128,11 +143,7 @@ async def test_aquery_billed_cost_includes_priced_vector_store_search():
                 mock_response="hi there",
             )
 
-        for _ in range(50):
-            if recording_logger.success_events:
-                break
-            await asyncio.sleep(0.1)
-        await asyncio.sleep(0.5)
+        await _drain_logging_worker()
     finally:
         litellm.callbacks = original_callbacks
 
@@ -155,6 +166,7 @@ async def test_aquery_with_rerank_bills_once_and_folds_rerank_cost():
     """
     from litellm.types.rerank import RerankResponse
 
+    await _drain_logging_worker()
     recording_logger = RecordingLogger()
     original_callbacks = litellm.callbacks
     litellm.callbacks = [recording_logger]
@@ -176,11 +188,7 @@ async def test_aquery_with_rerank_bills_once_and_folds_rerank_cost():
                 mock_response="hi there",
             )
 
-        for _ in range(50):
-            if recording_logger.success_events:
-                break
-            await asyncio.sleep(0.1)
-        await asyncio.sleep(0.5)
+        await _drain_logging_worker()
     finally:
         litellm.callbacks = original_callbacks
 
@@ -210,6 +218,7 @@ async def test_aquery_streaming_bills_sub_call_costs_into_final_event():
     """
     from litellm.types.rerank import RerankResponse
 
+    await _drain_logging_worker()
     recording_logger = RecordingLogger()
     original_callbacks = litellm.callbacks
     litellm.callbacks = [recording_logger]
@@ -237,11 +246,7 @@ async def test_aquery_streaming_bills_sub_call_costs_into_final_event():
             async for _ in response:
                 pass
 
-        for _ in range(50):
-            if recording_logger.success_events:
-                break
-            await asyncio.sleep(0.1)
-        await asyncio.sleep(0.5)
+        await _drain_logging_worker()
     finally:
         litellm.callbacks = original_callbacks