diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index cd39aa3931f..552fd833e79 100644 --- a/basedpyright-code-budget.json +++ b/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 diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index b27221c9beb..b2dc0a52c8f 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/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: diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 2aba8cabe17..d7d06387d85 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/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": diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 0f46f1b718c..39fa8776578 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/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 diff --git a/litellm/proxy/_lazy_openapi_snapshot.py b/litellm/proxy/_lazy_openapi_snapshot.py index 49d277cd3d1..92578aa43b9 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.py +++ b/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): diff --git a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py index 31dca5a7de2..c8284fac440 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py +++ b/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") diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 149c44ed083..0418eeaac8f 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -9,7 +9,7 @@ "limit": 827 }, "ANN201": { - "limit": 2012 + "limit": 2011 }, "ANN202": { "limit": 847 diff --git a/tests/test_litellm/rag/test_main.py b/tests/test_litellm/rag/test_main.py index 584124ba06a..2d1b460513f 100644 --- a/tests/test_litellm/rag/test_main.py +++ b/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