mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
fix(router): route model-less sync vector store calls to the SDK
_generic_api_call_with_fallbacks requires a model, so sync vector_store_search and vector_store_create raised a TypeError whenever the call carried no model. Model-less calls now go directly to the SDK function, with the router injected for search, matching the async wrapper's behavior
This commit is contained in:
parent
cf4738c3b7
commit
3914de24ef
2 changed files with 67 additions and 7 deletions
|
|
@ -6299,8 +6299,6 @@ class Router:
|
|||
"responses",
|
||||
"generate_content",
|
||||
"generate_content_stream",
|
||||
"vector_store_search",
|
||||
"vector_store_create",
|
||||
"ocr",
|
||||
"search",
|
||||
"video_generation",
|
||||
|
|
@ -6324,6 +6322,8 @@ class Router:
|
|||
return sync_wrapper
|
||||
|
||||
if call_type in (
|
||||
"vector_store_search",
|
||||
"vector_store_create",
|
||||
"vector_store_retrieve",
|
||||
"vector_store_list",
|
||||
"vector_store_update",
|
||||
|
|
@ -6335,11 +6335,16 @@ class Router:
|
|||
client: object | None = None,
|
||||
**kwargs,
|
||||
):
|
||||
if custom_llm_provider and "custom_llm_provider" not in kwargs:
|
||||
kwargs["custom_llm_provider"] = custom_llm_provider
|
||||
if kwargs.get("model"):
|
||||
return self._generic_api_call_with_fallbacks(original_function=original_function, **kwargs)
|
||||
return original_function(**kwargs)
|
||||
provider_kwargs: Final = (
|
||||
MappingProxyType({**kwargs, "custom_llm_provider": custom_llm_provider})
|
||||
if custom_llm_provider and "custom_llm_provider" not in kwargs
|
||||
else MappingProxyType(kwargs)
|
||||
)
|
||||
if provider_kwargs.get("model"):
|
||||
return self._generic_api_call_with_fallbacks(original_function=original_function, **provider_kwargs)
|
||||
if call_type == "vector_store_search":
|
||||
return original_function(**MappingProxyType({**provider_kwargs, "router": self}))
|
||||
return original_function(**provider_kwargs)
|
||||
|
||||
return vector_store_sync_wrapper
|
||||
|
||||
|
|
|
|||
|
|
@ -7567,6 +7567,61 @@ async def test_avector_store_create_does_not_inject_router():
|
|||
assert "router" not in mock_acreate.await_args.kwargs
|
||||
|
||||
|
||||
def test_vector_store_search_injects_router():
|
||||
"""
|
||||
Sync parity for the router injection: router.vector_store_search must pass
|
||||
the router down to the SDK search call so provider transforms can resolve
|
||||
router-managed embedding models, same as avector_store_search.
|
||||
"""
|
||||
from litellm.types.vector_stores import VectorStoreSearchResponse
|
||||
|
||||
expected_response = VectorStoreSearchResponse(
|
||||
object="vector_store.search_results.page", search_query="q", data=[]
|
||||
)
|
||||
mock_search = MagicMock(return_value=expected_response)
|
||||
# Router.__init__ binds search via a local import, so patch the module
|
||||
# attribute before constructing the Router.
|
||||
with patch("litellm.vector_stores.main.search", new=mock_search): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {"model": "openai/gpt-3.5-turbo", "api_key": "test-key"},
|
||||
}
|
||||
]
|
||||
)
|
||||
search_response = router.vector_store_search(
|
||||
vector_store_id="v", query="q", custom_llm_provider="s3_vectors"
|
||||
)
|
||||
|
||||
assert search_response is expected_response
|
||||
mock_search.assert_called_once()
|
||||
assert mock_search.call_args.kwargs["router"] is router
|
||||
assert mock_search.call_args.kwargs["custom_llm_provider"] == "s3_vectors"
|
||||
|
||||
|
||||
def test_vector_store_create_does_not_inject_router():
|
||||
"""The sync create path must keep calling the SDK without a router kwarg."""
|
||||
expected_response = {"id": "vs_1", "object": "vector_store"}
|
||||
mock_create = MagicMock(return_value=expected_response)
|
||||
# Router.__init__ binds create via a local import, so patch the module
|
||||
# attribute before constructing the Router.
|
||||
with patch("litellm.vector_stores.main.create", new=mock_create): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {"model": "openai/gpt-3.5-turbo", "api_key": "test-key"},
|
||||
}
|
||||
]
|
||||
)
|
||||
create_response = router.vector_store_create(custom_llm_provider="openai")
|
||||
|
||||
assert create_response is expected_response
|
||||
mock_create.assert_called_once()
|
||||
assert "router" not in mock_create.call_args.kwargs
|
||||
|
||||
|
||||
class TestPreRoutingStrategyRegistryLifecycle:
|
||||
"""
|
||||
Regression tests: a deployment leaving the model_list must release the
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue