diff --git a/litellm/router.py b/litellm/router.py index 1a76eb5d59b..d75dd4ac284 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -10242,7 +10242,7 @@ class Router: return litellm.token_counter(messages=cast(list, input_messages)) # cast-ok: transformed chat messages raise ValueError("Either messages or input must be provided to count tokens") - def _pre_call_checks( + def _pre_call_checks_sync( self, model: str, healthy_deployments: list, @@ -10419,6 +10419,29 @@ class Router: return _returned_deployments + async def _pre_call_checks_async( + self, + model: str, + healthy_deployments: list, + messages: list[dict[str, str]] | None = None, + input: str | list | None = None, + request_kwargs: dict | None = None, + ) -> list[dict]: + """ + Async wrapper around _pre_call_checks_sync. + + Offloads sync work to a worker thread so async callers can await + without blocking the event loop. + """ + return await asyncio.to_thread( + self._pre_call_checks_sync, + model=model, + healthy_deployments=healthy_deployments, + messages=messages, + input=input, + request_kwargs=request_kwargs, + ) + def _get_model_from_alias(self, model: str) -> str | None: """ Get the model from the alias. @@ -10768,7 +10791,7 @@ class Router: ) if self.enable_pre_call_checks and (messages is not None or input is not None): - healthy_deployments = self._pre_call_checks( + healthy_deployments = await self._pre_call_checks_async( model=model, healthy_deployments=cast(list[dict], healthy_deployments), messages=messages, @@ -11352,7 +11375,7 @@ class Router: # filter pre-call checks if self.enable_pre_call_checks and (messages is not None or input is not None): - healthy_deployments = self._pre_call_checks( + healthy_deployments = self._pre_call_checks_sync( model=model, healthy_deployments=healthy_deployments, messages=messages, @@ -11510,7 +11533,7 @@ class Router: # 5. Apply pre-call checks (if enabled) if self.enable_pre_call_checks and (messages is not None or input is not None): - pass_through_deployments = self._pre_call_checks( + pass_through_deployments = self._pre_call_checks_sync( model=model, healthy_deployments=pass_through_deployments, messages=messages, diff --git a/tests/local_testing/test_router.py b/tests/local_testing/test_router.py index 2ee1aee9710..57cabc37284 100644 --- a/tests/local_testing/test_router.py +++ b/tests/local_testing/test_router.py @@ -634,7 +634,7 @@ def test_router_rpm_pre_call_check(): router = Router(model_list=model_list, set_verbose=True, enable_pre_call_checks=True, num_retries=0) # type: ignore try: - router._pre_call_checks( + router._pre_call_checks_sync( model="fake-openai-endpoint", healthy_deployments=model_list, messages=[{"role": "user", "content": "Hey, how's it going?"}], @@ -831,7 +831,7 @@ def test_filter_invalid_params_pre_call_check(): router = Router(model_list=model_list, set_verbose=True, enable_pre_call_checks=True, num_retries=0) # type: ignore - filtered_deployments = router._pre_call_checks( + filtered_deployments = router._pre_call_checks_sync( model="gpt-3.5-turbo", healthy_deployments=model_list, messages=[{"role": "user", "content": "Hey, how's it going?"}], @@ -874,7 +874,7 @@ def test_router_region_pre_call_check(allowed_model_region): router = Router(model_list=model_list, enable_pre_call_checks=True) - _healthy_deployments = router._pre_call_checks( + _healthy_deployments = router._pre_call_checks_sync( model="gpt-3.5-turbo", healthy_deployments=model_list, messages=[{"role": "user", "content": "Hey!"}], @@ -1895,7 +1895,7 @@ def test_router_context_window_pre_call_check(model, base_model, llm_provider): litellm.token_counter.side_effect = token_counter_side_effect try: - updated_list = router._pre_call_checks( + updated_list = router._pre_call_checks_sync( model="gpt-4", healthy_deployments=model_list, messages=[{"role": "user", "content": "Hey, how's it going?"}], diff --git a/tests/router_unit_tests/test_pre_call_checks_optimization.py b/tests/router_unit_tests/test_pre_call_checks_optimization.py index 54d11d482a7..bf70f06848f 100644 --- a/tests/router_unit_tests/test_pre_call_checks_optimization.py +++ b/tests/router_unit_tests/test_pre_call_checks_optimization.py @@ -61,7 +61,7 @@ class TestPreCallChecksOptimization: snapshot = copy.deepcopy(deployments) # Call the function under test - router._pre_call_checks( + router._pre_call_checks_sync( model="gpt-5-mini", healthy_deployments=deployments, messages=[{"role": "user", "content": "test"}], @@ -113,7 +113,7 @@ class TestPreCallChecksOptimization: original_large_deployment = deployments[1] # max_input_tokens=10000 # Send a long message (100 words) that exceeds 50 tokens but fits in 10000 tokens - filtered = router._pre_call_checks( + filtered = router._pre_call_checks_sync( model="test", healthy_deployments=deployments, messages=[{"role": "user", "content": " ".join(["word"] * 100)}], diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index da2d78edb73..2e81c9cecc7 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -3131,7 +3131,7 @@ def test_pre_call_checks_skips_token_count_without_max_input_tokens(monkeypatch) {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d2"}}, ] - result = router._pre_call_checks( + result = router._pre_call_checks_sync( model="m", healthy_deployments=deployments, messages=[{"role": "user", "content": "hi"}], @@ -3167,7 +3167,7 @@ def test_pre_call_checks_counts_once_and_filters_on_max_input_tokens(monkeypatch {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d2"}}, ] with pytest.raises(litellm.ContextWindowExceededError): - router._pre_call_checks( + router._pre_call_checks_sync( model="m", healthy_deployments=deployments, messages=[{"role": "user", "content": "hi"}], @@ -3196,7 +3196,7 @@ def test_pre_call_checks_counts_tokens_from_responses_input_string(monkeypatch): {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, ] with pytest.raises(litellm.ContextWindowExceededError): - router._pre_call_checks( + router._pre_call_checks_sync( model="m", healthy_deployments=deployments, input="a very long prompt that exceeds the tiny context window", @@ -3223,7 +3223,7 @@ def test_pre_call_checks_counts_tokens_from_responses_input_list(monkeypatch): {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, ] with pytest.raises(litellm.ContextWindowExceededError): - router._pre_call_checks( + router._pre_call_checks_sync( model="m", healthy_deployments=deployments, input=[ @@ -3263,7 +3263,7 @@ def test_pre_call_checks_counts_responses_instructions_tokens(monkeypatch): router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": input_only_tokens} ) with pytest.raises(litellm.ContextWindowExceededError): - router._pre_call_checks( + router._pre_call_checks_sync( model="m", healthy_deployments=deployments, input=short_input, @@ -3324,7 +3324,7 @@ def test_pre_call_checks_no_messages_or_input_does_not_crash(monkeypatch): deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, ] - result = router._pre_call_checks(model="m", healthy_deployments=deployments) + result = router._pre_call_checks_sync(model="m", healthy_deployments=deployments) assert len(result) == 1 assert counted == [] # token counting skipped entirely, so no misleading error is logged @@ -7319,7 +7319,7 @@ def test_pre_call_checks_uses_deployment_model_when_model_info_lookup_raises(mon "model_info": {"id": "d1"}, } ] - result = router._pre_call_checks( + result = router._pre_call_checks_sync( model="custom-alias", healthy_deployments=deployments, messages=[{"role": "user", "content": "hi"}], @@ -7363,7 +7363,7 @@ def test_pre_call_checks_keeps_deployment_when_provider_is_unresolvable(monkeypa "model_info": {"id": "d1"}, } ] - result = router._pre_call_checks( + result = router._pre_call_checks_sync( model="custom-alias", healthy_deployments=deployments, messages=[{"role": "user", "content": "hi"}],