diff --git a/litellm/__init__.py b/litellm/__init__.py index 25da650440..4f0b0a16be 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -756,6 +756,7 @@ from .utils import ( create_pretrained_tokenizer, create_tokenizer, supports_function_calling, + supports_web_search, supports_response_schema, supports_parallel_function_calling, supports_vision, diff --git a/litellm/utils.py b/litellm/utils.py index 03e69acf4e..052857cd94 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1975,7 +1975,7 @@ def supports_system_messages(model: str, custom_llm_provider: Optional[str]) -> ) -def supports_web_search(model: str, custom_llm_provider: Optional[str]) -> bool: +def supports_web_search(model: str, custom_llm_provider: Optional[str] = None) -> bool: """ Check if the given model supports web search and return a boolean value. diff --git a/tests/litellm_utils_tests/test_utils.py b/tests/litellm_utils_tests/test_utils.py index 42df7c495c..fea225e4a3 100644 --- a/tests/litellm_utils_tests/test_utils.py +++ b/tests/litellm_utils_tests/test_utils.py @@ -477,6 +477,25 @@ def test_supports_function_calling(model, expected_bool): pytest.fail(f"Error occurred: {e}") +@pytest.mark.parametrize( + "model, expected_bool", + [ + ("gpt-4o-mini-search-preview", True), + ("openai/gpt-4o-mini-search-preview", True), + ("gpt-4o-search-preview", True), + ("openai/gpt-4o-search-preview", True), + ("groq/deepseek-r1-distill-llama-70b", False), + ("groq/llama-3.3-70b-versatile", False), + ("codestral/codestral-latest", False), + ], +) +def test_supports_web_search(model, expected_bool): + try: + assert litellm.supports_web_search(model=model) == expected_bool + except Exception as e: + pytest.fail(f"Error occurred: {e}") + + def test_get_max_token_unit_test(): """ More complete testing in `test_completion_cost.py`