diff --git a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py index 14a0a406dff..a6fbe33a866 100644 --- a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py @@ -133,22 +133,25 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict[str, Any]]: """ - Transform search request for Vertex AI RAG API + Transform a search request for the Vertex AI Search (Discovery Engine) API. + + Per-request params pass through to the engine: max_num_results maps to + pageSize, and extra_body is merged in raw and takes precedence, so callers + can send native Discovery Engine fields such as dataStoreSpecs, filter, + boostSpec, or contentSearchSpec. """ - # Convert query to string if it's a list if isinstance(query, list): query = " ".join(query) - # Vertex AI RAG API endpoint for retrieving contexts url = f"{api_base}:search" - # Construct full rag corpus path - # Build the request body for Vertex AI Search API - request_body = {"query": query, "pageSize": 10} + request_body: Dict[str, Any] = {"query": query, "pageSize": 10} + max_num_results = vector_store_search_optional_params.get("max_num_results") + if max_num_results is not None: + request_body["pageSize"] = max_num_results + if isinstance(extra_body, dict): + request_body.update(extra_body) - ######################################################### - # Update logging object with details of the request - ######################################################### litellm_logging_obj.model_call_details["query"] = query return url, request_body diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py index 5ca71dc08c3..83d9fd703d8 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py @@ -1,3 +1,5 @@ +from types import SimpleNamespace + import pytest from litellm.llms.vertex_ai.vector_stores.search_api.transformation import ( @@ -126,3 +128,69 @@ def test_should_raise_when_neither_engine_id_nor_vector_store_id_provided(): "vertex_location": "global", }, ) + + +_ENGINE_BASE = ( + "https://discoveryengine.googleapis.com/v1/projects/p/locations/global/" + "collections/default_collection/engines/app-2/servingConfigs/default_serving_config" +) + + +def _search_request(**overrides): + kwargs = dict( + vector_store_id="vs", + query="hello", + vector_store_search_optional_params={}, + api_base=_ENGINE_BASE, + litellm_logging_obj=SimpleNamespace(model_call_details={}), + litellm_params={}, + ) + kwargs.update(overrides) + return VertexSearchAPIVectorStoreConfig().transform_search_vector_store_request( + **kwargs + ) + + +def test_search_request_defaults_to_query_and_pagesize_10(): + url, body = _search_request() + + assert url == _ENGINE_BASE + ":search" + assert body == {"query": "hello", "pageSize": 10} + + +def test_search_request_maps_max_num_results_to_pagesize(): + _, body = _search_request( + vector_store_search_optional_params={"max_num_results": 25} + ) + + assert body["pageSize"] == 25 + + +def test_search_request_passes_datastorespecs_through_extra_body(): + specs = [ + { + "dataStore": "projects/p/locations/global/collections/default_collection/dataStores/ds-beta" + } + ] + + _, body = _search_request(extra_body={"dataStoreSpecs": specs}) + + assert body["dataStoreSpecs"] == specs + assert body["query"] == "hello" + assert body["pageSize"] == 10 + + +def test_search_request_extra_body_takes_precedence_over_defaults(): + _, body = _search_request( + vector_store_search_optional_params={"max_num_results": 5}, + extra_body={"pageSize": 50, "filter": 'category: ANY("docs")'}, + ) + + assert body["pageSize"] == 50 + assert body["filter"] == 'category: ANY("docs")' + + +def test_search_request_joins_list_query(): + _, body = _search_request(query=["foo", "bar"]) + + assert body["query"] == "foo bar"