From 6d0d67f3f25da5a26e574a51520c8f1fea6b6aff Mon Sep 17 00:00:00 2001 From: Rohit Kanithi Date: Sun, 13 Sep 2026 02:08:42 -0500 Subject: [PATCH 1/3] fix(search): raise exceptions on non-2xx HTTP status codes and plumb client parameter --- litellm/llms/custom_httpx/llm_http_handler.py | 4 +- litellm/search/main.py | 2 + .../search/test_base_search_transformation.py | 179 +++++++++++++++--- 3 files changed, 152 insertions(+), 33 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 2fe4130a310..a4b54946093 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1918,7 +1918,6 @@ class BaseLLMHTTPHandler: url=complete_url, headers=signed_headers, ) - response.raise_for_status() else: # A signed body must be sent verbatim, re-serializing it would break the signature response = client.post( @@ -1928,6 +1927,7 @@ class BaseLLMHTTPHandler: json=data if signed_json_body is None else None, timeout=timeout, ) + response.raise_for_status() except httpx.HTTPStatusError as e: raise provider_config.get_http_error_class(e) except Exception as e: @@ -2022,7 +2022,6 @@ class BaseLLMHTTPHandler: url=complete_url, headers=signed_headers, ) - response.raise_for_status() else: # A signed body must be sent verbatim, re-serializing it would break the signature response = await async_httpx_client.post( @@ -2032,6 +2031,7 @@ class BaseLLMHTTPHandler: json=data if signed_json_body is None else None, timeout=timeout, ) + response.raise_for_status() except httpx.HTTPStatusError as e: raise provider_config.get_http_error_class(e) except Exception as e: diff --git a/litellm/search/main.py b/litellm/search/main.py index b2dd51799a1..25379332b05 100644 --- a/litellm/search/main.py +++ b/litellm/search/main.py @@ -230,6 +230,7 @@ def search( try: litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.pop("litellm_logging_obj") litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id", None) + client: Final = kwargs.get("client", None) _is_async: Final = kwargs.pop("asearch", False) is True # Validate query parameter @@ -303,6 +304,7 @@ def search( api_key=api_key, api_base=complete_url, custom_llm_provider=search_provider, + client=client, asearch=_is_async, headers=headers, provider_config=search_provider_config, diff --git a/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py b/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py index e6aad7688d1..787acf42c14 100644 --- a/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py +++ b/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py @@ -10,18 +10,23 @@ leaving keyless providers and legitimate operator overrides untouched. """ from typing import Dict, Tuple, Type -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest import litellm from litellm.llms.apiserpent.search.transformation import APISerpentSearchConfig from litellm.llms.azure.search.transformation import BingGroundingSearchConfig +from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.search.transformation import ( BaseSearchConfig, + SearchResponse, _is_trusted_search_api_base, ) from litellm.llms.brave.search.transformation import BraveSearchConfig +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.llms.dataforseo.search.transformation import DataForSEOSearchConfig from litellm.llms.exa_ai.search.transformation import ExaAISearchConfig from litellm.llms.fastcrw.search.transformation import FastCRWSearchConfig @@ -107,9 +112,7 @@ PROVIDERS: Tuple[ProviderSpec, ...] = ( _IDS = tuple(spec[0].__name__ for spec in PROVIDERS) -@pytest.mark.parametrize( - "config_cls, server_env, caller_key, extra_env", PROVIDERS, ids=_IDS -) +@pytest.mark.parametrize("config_cls, server_env, caller_key, extra_env", PROVIDERS, ids=_IDS) def test_server_secret_refused_for_caller_api_base( config_cls: Type[BaseSearchConfig], server_env: Dict[str, str], @@ -124,9 +127,7 @@ def test_server_secret_refused_for_caller_api_base( config_cls().validate_environment(headers={}, api_base=ATTACKER_BASE) -@pytest.mark.parametrize( - "config_cls, server_env, caller_key, extra_env", PROVIDERS, ids=_IDS -) +@pytest.mark.parametrize("config_cls, server_env, caller_key, extra_env", PROVIDERS, ids=_IDS) def test_caller_supplied_key_is_honored_for_custom_api_base( config_cls: Type[BaseSearchConfig], server_env: Dict[str, str], @@ -139,14 +140,10 @@ def test_caller_supplied_key_is_honored_for_custom_api_base( # An explicit caller key is the caller's own credential, so pointing it at # the caller's own host must be allowed. - config_cls().validate_environment( - headers={}, api_key=caller_key, api_base=ATTACKER_BASE - ) + config_cls().validate_environment(headers={}, api_key=caller_key, api_base=ATTACKER_BASE) -@pytest.mark.parametrize( - "config_cls, server_env, caller_key, extra_env", PROVIDERS, ids=_IDS -) +@pytest.mark.parametrize("config_cls, server_env, caller_key, extra_env", PROVIDERS, ids=_IDS) def test_server_secret_used_without_caller_api_base( config_cls: Type[BaseSearchConfig], server_env: Dict[str, str], @@ -167,9 +164,7 @@ def test_keyless_provider_allows_caller_api_base( ) -> None: monkeypatch.delenv("SEARXNG_API_KEY", raising=False) - headers = SearXNGSearchConfig().validate_environment( - headers={}, api_base="https://my-searxng.internal" - ) + headers = SearXNGSearchConfig().validate_environment(headers={}, api_base="https://my-searxng.internal") assert "Authorization" not in headers @@ -182,9 +177,7 @@ def test_operator_env_base_override_is_trusted( # Mirrors the second validate_environment call in the search handler, which # receives the already-resolved operator base as api_base. - headers = SerperSearchConfig().validate_environment( - headers={}, api_base="https://serper.internal.corp/search" - ) + headers = SerperSearchConfig().validate_environment(headers={}, api_base="https://serper.internal.corp/search") assert headers["X-API-KEY"] == "srv" @@ -200,9 +193,7 @@ class TestResolveServerApiKey: ) assert result == "mine" - def test_returns_none_when_no_server_secret( - self, monkeypatch: pytest.MonkeyPatch - ) -> None: + def test_returns_none_when_no_server_secret(self, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.delenv("SEARXNG_API_KEY", raising=False) result = BaseSearchConfig().resolve_server_api_key( caller_api_key=None, @@ -228,14 +219,10 @@ class TestResolveServerApiKey: class TestIsTrustedSearchApiBase: def test_matches_default_host(self) -> None: - assert _is_trusted_search_api_base( - "https://google.serper.dev/search", "https://google.serper.dev", None - ) + assert _is_trusted_search_api_base("https://google.serper.dev/search", "https://google.serper.dev", None) def test_foreign_host_untrusted(self) -> None: - assert not _is_trusted_search_api_base( - ATTACKER_BASE, "https://google.serper.dev", None - ) + assert not _is_trusted_search_api_base(ATTACKER_BASE, "https://google.serper.dev", None) def test_env_override_host_trusted(self, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("SERPER_API_BASE", "https://serper.internal.corp") @@ -248,9 +235,7 @@ class TestIsTrustedSearchApiBase: def test_schemeless_candidate_untrusted(self) -> None: # Without a scheme urlsplit puts the value in the path, leaving an empty # netloc; an unparseable host must never be treated as trusted. - assert not _is_trusted_search_api_base( - "attacker.example.com", "https://google.serper.dev", None - ) + assert not _is_trusted_search_api_base("attacker.example.com", "https://google.serper.dev", None) @pytest.mark.asyncio @@ -333,3 +318,135 @@ async def test_query_param_key_not_leaked_with_dummy_caller_key( assert captured["url"], "expected an outbound request to be attempted" assert server_key not in captured["url"] assert "sk-CALLER-DUMMY" in captured["url"] + + +class TestSearchHTTPErrorHandling: + """Tests verifying that non-2xx HTTP responses from search providers raise proper exceptions and custom clients are forwarded.""" + + @pytest.mark.parametrize( + ("status_code", "expected_exception"), + [ + (401, litellm.AuthenticationError), + (429, litellm.RateLimitError), + (500, litellm.InternalServerError), + (503, litellm.ServiceUnavailableError), + ], + ) + def test_sync_search_http_error_raising( + self, + status_code: int, + expected_exception: type[BaseLLMException], + ) -> None: + """Verify that synchronous search raises mapped exceptions on error status codes.""" + mock_response = httpx.Response( + status_code=status_code, + request=httpx.Request("POST", "https://api.tavily.com/search"), + json={"error": "Test error message"}, + ) + + with patch.object(HTTPHandler, "post", return_value=mock_response): + with pytest.raises(expected_exception): + litellm.search( + query="test query", + search_provider="tavily", + api_key="tvly-testkey", + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("status_code", "expected_exception"), + [ + (401, litellm.AuthenticationError), + (429, litellm.RateLimitError), + (500, litellm.InternalServerError), + (503, litellm.ServiceUnavailableError), + ], + ) + async def test_async_search_http_error_raising( + self, + status_code: int, + expected_exception: type[BaseLLMException], + ) -> None: + """Verify that asynchronous search raises mapped exceptions on error status codes.""" + mock_response = httpx.Response( + status_code=status_code, + request=httpx.Request("POST", "https://api.tavily.com/search"), + json={"error": "Async test error message"}, + ) + + with patch.object(AsyncHTTPHandler, "post", new_callable=AsyncMock) as mock_post: + mock_post.return_value = mock_response + with pytest.raises(expected_exception): + await litellm.asearch( + query="test query", + search_provider="tavily", + api_key="tvly-testkey", + ) + + def test_sync_search_plumbs_custom_client(self) -> None: + """Verify that a custom HTTPHandler passed to litellm.search is forwarded to the underlying handler.""" + custom_client = HTTPHandler() + mock_response = httpx.Response( + status_code=200, + request=httpx.Request("POST", "https://api.tavily.com/search"), + json={"results": [{"title": "Test", "url": "https://example.com", "content": "Sample"}]}, + ) + + with ( + patch.object(custom_client, "post", return_value=mock_response) as mock_post, + patch.object(BaseLLMHTTPHandler, "search", wraps=BaseLLMHTTPHandler().search) as mock_handler_search, + ): + response = litellm.search( + query="test query", + search_provider="tavily", + api_key="tvly-testkey", + client=custom_client, + ) + assert isinstance(response, SearchResponse) + assert mock_handler_search.call_count == 1 + assert mock_handler_search.call_args.kwargs["client"] is custom_client + assert mock_post.call_count == 1 + + @pytest.mark.asyncio + async def test_async_search_plumbs_custom_client(self) -> None: + """Verify that a custom AsyncHTTPHandler passed to litellm.asearch is forwarded to the underlying handler.""" + custom_client = AsyncHTTPHandler() + mock_response = httpx.Response( + status_code=200, + request=httpx.Request("POST", "https://api.tavily.com/search"), + json={"results": [{"title": "Async Test", "url": "https://example.com", "content": "Async Sample"}]}, + ) + + with ( + patch.object(custom_client, "post", new_callable=AsyncMock) as mock_post, + patch.object( + BaseLLMHTTPHandler, "async_search", wraps=BaseLLMHTTPHandler().async_search + ) as mock_handler_asearch, + ): + mock_post.return_value = mock_response + response = await litellm.asearch( + query="test query", + search_provider="tavily", + api_key="tvly-testkey", + client=custom_client, + ) + assert isinstance(response, SearchResponse) + assert mock_handler_asearch.call_count == 1 + assert mock_handler_asearch.call_args.kwargs["client"] is custom_client + assert mock_post.call_count == 1 + + def test_sync_search_handles_empty_error_response_body(self) -> None: + """Verify that 500 responses with empty bodies still raise InternalServerError properly.""" + mock_response = httpx.Response( + status_code=500, + request=httpx.Request("POST", "https://api.tavily.com/search"), + text="", + ) + + with patch.object(HTTPHandler, "post", return_value=mock_response): + with pytest.raises(litellm.InternalServerError): + litellm.search( + query="test query", + search_provider="tavily", + api_key="tvly-testkey", + ) From 30bf2b360cfad2a9c3740916ec31f67a830cafaa Mon Sep 17 00:00:00 2001 From: Rohit Kanithi Date: Sun, 13 Sep 2026 02:14:52 -0500 Subject: [PATCH 2/3] fix(search): validate client type for search and asearch modes --- litellm/search/main.py | 12 ++++++++++ .../search/test_base_search_transformation.py | 23 +++++++++++++++++++ 2 files changed, 35 insertions(+) diff --git a/litellm/search/main.py b/litellm/search/main.py index 25379332b05..148bba31b48 100644 --- a/litellm/search/main.py +++ b/litellm/search/main.py @@ -15,6 +15,7 @@ from litellm._logging import verbose_logger from litellm.constants import request_timeout from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.search.transformation import BaseSearchConfig, SearchResponse +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.types.utils import SearchProviders from litellm.utils import ProviderConfigManager, client, filter_out_litellm_params @@ -233,6 +234,17 @@ def search( client: Final = kwargs.get("client", None) _is_async: Final = kwargs.pop("asearch", False) is True + if _is_async: + if client is not None and not isinstance(client, AsyncHTTPHandler): + raise ValueError( + f"client must be an instance of AsyncHTTPHandler for asynchronous search, got {type(client)}" + ) + else: + if client is not None and not isinstance(client, HTTPHandler): + raise ValueError( + f"client must be an instance of HTTPHandler for synchronous search, got {type(client)}" + ) + # Validate query parameter if not isinstance(query, (str, list)): raise ValueError(f"query must be a string or list of strings, got {type(query)}") diff --git a/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py b/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py index 787acf42c14..cd8b78e7910 100644 --- a/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py +++ b/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py @@ -450,3 +450,26 @@ class TestSearchHTTPErrorHandling: search_provider="tavily", api_key="tvly-testkey", ) + + def test_sync_search_rejects_async_client(self) -> None: + """Verify that passing an AsyncHTTPHandler to synchronous search raises a clear exception.""" + async_client = AsyncHTTPHandler() + with pytest.raises(Exception, match="client must be an instance of HTTPHandler"): + litellm.search( + query="test query", + search_provider="tavily", + api_key="tvly-testkey", + client=async_client, + ) + + @pytest.mark.asyncio + async def test_async_search_rejects_sync_client(self) -> None: + """Verify that passing a sync HTTPHandler to asynchronous search raises a clear exception.""" + sync_client = HTTPHandler() + with pytest.raises(Exception, match="client must be an instance of AsyncHTTPHandler"): + await litellm.asearch( + query="test query", + search_provider="tavily", + api_key="tvly-testkey", + client=sync_client, + ) From f344a20e6c8a87301bf6ec72a59fa364fa43c6a1 Mon Sep 17 00:00:00 2001 From: Rohit Kanithi Date: Sun, 13 Sep 2026 02:32:47 -0500 Subject: [PATCH 3/3] test(search): use MockTransport for search error testing to satisfy TQ gate --- .../search/test_base_search_transformation.py | 161 +++++++++--------- 1 file changed, 83 insertions(+), 78 deletions(-) diff --git a/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py b/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py index cd8b78e7910..91c2c685b45 100644 --- a/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py +++ b/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py @@ -338,19 +338,22 @@ class TestSearchHTTPErrorHandling: expected_exception: type[BaseLLMException], ) -> None: """Verify that synchronous search raises mapped exceptions on error status codes.""" - mock_response = httpx.Response( - status_code=status_code, - request=httpx.Request("POST", "https://api.tavily.com/search"), - json={"error": "Test error message"}, - ) - with patch.object(HTTPHandler, "post", return_value=mock_response): - with pytest.raises(expected_exception): - litellm.search( - query="test query", - search_provider="tavily", - api_key="tvly-testkey", - ) + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + status_code=status_code, + request=request, + json={"error": "Test error message"}, + ) + + client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler))) + with pytest.raises(expected_exception): + litellm.search( + query="test query", + search_provider="tavily", + api_key="tvly-testkey", + client=client, + ) @pytest.mark.asyncio @pytest.mark.parametrize( @@ -368,88 +371,90 @@ class TestSearchHTTPErrorHandling: expected_exception: type[BaseLLMException], ) -> None: """Verify that asynchronous search raises mapped exceptions on error status codes.""" - mock_response = httpx.Response( - status_code=status_code, - request=httpx.Request("POST", "https://api.tavily.com/search"), - json={"error": "Async test error message"}, - ) - with patch.object(AsyncHTTPHandler, "post", new_callable=AsyncMock) as mock_post: - mock_post.return_value = mock_response - with pytest.raises(expected_exception): - await litellm.asearch( - query="test query", - search_provider="tavily", - api_key="tvly-testkey", - ) + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + status_code=status_code, + request=request, + json={"error": "Async test error message"}, + ) - def test_sync_search_plumbs_custom_client(self) -> None: - """Verify that a custom HTTPHandler passed to litellm.search is forwarded to the underlying handler.""" - custom_client = HTTPHandler() - mock_response = httpx.Response( - status_code=200, - request=httpx.Request("POST", "https://api.tavily.com/search"), - json={"results": [{"title": "Test", "url": "https://example.com", "content": "Sample"}]}, - ) - - with ( - patch.object(custom_client, "post", return_value=mock_response) as mock_post, - patch.object(BaseLLMHTTPHandler, "search", wraps=BaseLLMHTTPHandler().search) as mock_handler_search, - ): - response = litellm.search( + client = AsyncHTTPHandler() + client._client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + with pytest.raises(expected_exception): + await litellm.asearch( query="test query", search_provider="tavily", api_key="tvly-testkey", - client=custom_client, + client=client, ) - assert isinstance(response, SearchResponse) - assert mock_handler_search.call_count == 1 - assert mock_handler_search.call_args.kwargs["client"] is custom_client - assert mock_post.call_count == 1 + + def test_sync_search_plumbs_custom_client(self) -> None: + """Verify that a custom HTTPHandler passed to litellm.search is forwarded to the underlying handler.""" + called = False + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal called + called = True + return httpx.Response( + status_code=200, + request=request, + json={"results": [{"title": "Test", "url": "https://example.com", "content": "Sample"}]}, + ) + + custom_client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler))) + response = litellm.search( + query="test query", + search_provider="tavily", + api_key="tvly-testkey", + client=custom_client, + ) + assert isinstance(response, SearchResponse) + assert called @pytest.mark.asyncio async def test_async_search_plumbs_custom_client(self) -> None: """Verify that a custom AsyncHTTPHandler passed to litellm.asearch is forwarded to the underlying handler.""" - custom_client = AsyncHTTPHandler() - mock_response = httpx.Response( - status_code=200, - request=httpx.Request("POST", "https://api.tavily.com/search"), - json={"results": [{"title": "Async Test", "url": "https://example.com", "content": "Async Sample"}]}, - ) + called = False - with ( - patch.object(custom_client, "post", new_callable=AsyncMock) as mock_post, - patch.object( - BaseLLMHTTPHandler, "async_search", wraps=BaseLLMHTTPHandler().async_search - ) as mock_handler_asearch, - ): - mock_post.return_value = mock_response - response = await litellm.asearch( - query="test query", - search_provider="tavily", - api_key="tvly-testkey", - client=custom_client, + def handler(request: httpx.Request) -> httpx.Response: + nonlocal called + called = True + return httpx.Response( + status_code=200, + request=request, + json={"results": [{"title": "Async Test", "url": "https://example.com", "content": "Async Sample"}]}, ) - assert isinstance(response, SearchResponse) - assert mock_handler_asearch.call_count == 1 - assert mock_handler_asearch.call_args.kwargs["client"] is custom_client - assert mock_post.call_count == 1 + + custom_client = AsyncHTTPHandler() + custom_client._client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + response = await litellm.asearch( + query="test query", + search_provider="tavily", + api_key="tvly-testkey", + client=custom_client, + ) + assert isinstance(response, SearchResponse) + assert called def test_sync_search_handles_empty_error_response_body(self) -> None: """Verify that 500 responses with empty bodies still raise InternalServerError properly.""" - mock_response = httpx.Response( - status_code=500, - request=httpx.Request("POST", "https://api.tavily.com/search"), - text="", - ) - with patch.object(HTTPHandler, "post", return_value=mock_response): - with pytest.raises(litellm.InternalServerError): - litellm.search( - query="test query", - search_provider="tavily", - api_key="tvly-testkey", - ) + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + status_code=500, + request=request, + text="", + ) + + client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler))) + with pytest.raises(litellm.InternalServerError): + litellm.search( + query="test query", + search_provider="tavily", + api_key="tvly-testkey", + client=client, + ) def test_sync_search_rejects_async_client(self) -> None: """Verify that passing an AsyncHTTPHandler to synchronous search raises a clear exception."""