From eb5730957b783a5ed2bf97d39405f5b431b1a2fc Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 27 Jul 2026 11:20:04 +0000 Subject: [PATCH] fix(responses): search emulated file_search stores with their registered provider --- litellm/proxy/rag_endpoints/endpoints.py | 14 +- .../responses/file_search/emulated_handler.py | 72 ++++++-- .../llms/test_file_search_responses.py | 157 ++++++++++++++++++ .../proxy/rag_endpoints/test_rag_endpoints.py | 79 +++++++++ 4 files changed, 304 insertions(+), 18 deletions(-) diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index 27ffc49901b..b95312f2ddc 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -208,11 +208,17 @@ async def _save_vector_store_to_db_from_rag_ingest( # Extract provider-specific params from vector_store_config to save as litellm_params # This ensures params like aws_region_name, embedding_model, etc. are available for search - provider_specific_params = {} excluded_keys = {"custom_llm_provider", "vector_store_id"} - for key, value in vector_store_config.items(): - if key not in excluded_keys and value is not None: - provider_specific_params[key] = value + provider_specific_params = { + key: value for key, value in vector_store_config.items() if key not in excluded_keys and value is not None + } + + ingest_embedding_model = (ingest_options.get("embedding") or {}).get("model") + if ingest_embedding_model and "embedding_model" not in provider_specific_params: + provider_specific_params = { + **provider_specific_params, + "embedding_model": ingest_embedding_model, + } # Build file metadata entry using helper file_entry = _build_file_metadata_entry( diff --git a/litellm/responses/file_search/emulated_handler.py b/litellm/responses/file_search/emulated_handler.py index 6e27ad66f54..c77971d448a 100644 --- a/litellm/responses/file_search/emulated_handler.py +++ b/litellm/responses/file_search/emulated_handler.py @@ -125,10 +125,44 @@ def _replace_file_search_tools( # --------------------------------------------------------------------------- +async def _resolve_search_params_for_vector_store(vector_store_id: str) -> dict[str, Any]: + """ + Resolve the provider and provider-specific params a registered vector store needs at search time. + + Without this, `asearch` falls back to `custom_llm_provider="openai"` and drops params such as + `vector_bucket_name`, `index_name` and `embedding_model`, so a store backed by any other + provider (s3_vectors, bedrock, ...) is never actually searched. + """ + import litellm + + registry = litellm.vector_store_registry + if registry is None: + return {} + + try: + from litellm.proxy.proxy_server import prisma_client + except ImportError: + prisma_client = None + + vector_store = await registry.get_litellm_managed_vector_store_from_registry_or_db( + vector_store_id=vector_store_id, + prisma_client=prisma_client, + ) + if vector_store is None: + return {} + + litellm_params: dict[str, Any] = vector_store.get("litellm_params") or {} + search_params = {k: v for k, v in litellm_params.items() if k not in ("vector_store_id", "custom_llm_provider")} + custom_llm_provider = litellm_params.get("custom_llm_provider") or vector_store.get("custom_llm_provider") + if custom_llm_provider is None: + return search_params + return {**search_params, "custom_llm_provider": custom_llm_provider} + + async def _run_vector_searches( queries: List[str], vector_store_ids: List[str], -) -> Tuple[List[str], List[VectorStoreSearchResult]]: +) -> tuple[list[str], list[VectorStoreSearchResult], tuple[str, ...]]: """ Run `asearch` against all vector stores for all queries and collect results. @@ -137,12 +171,14 @@ async def _run_vector_searches( vector_store_ids: Vector store IDs to search Returns: - (queries_list, combined_results) + (queries_list, combined_results, search_errors) """ import litellm.vector_stores.main as vs_main all_results: List[VectorStoreSearchResult] = [] + errors: list[str] = [] ids_to_search = vector_store_ids + search_params_by_id = {vs_id: await _resolve_search_params_for_vector_store(vs_id) for vs_id in ids_to_search} # Execute each query against all vector stores for query in queries: @@ -151,19 +187,21 @@ async def _run_vector_searches( response = await vs_main.asearch( vector_store_id=vs_id, query=query, + **search_params_by_id[vs_id], ) results_data = response.get("data") if isinstance(response, dict) else getattr(response, "data", None) if results_data: all_results.extend(results_data) except Exception as exc: - verbose_logger.warning( + errors.append(f"vector_store_id={vs_id}: {exc}") + verbose_logger.exception( "file_search emulated: search failed for query='%s', vector_store_id='%s': %s", query, vs_id, exc, ) - return queries, all_results + return queries, all_results, tuple(errors) # --------------------------------------------------------------------------- @@ -242,6 +280,7 @@ def _build_file_search_call_output( queries: List[str], results: Optional[List[VectorStoreSearchResult]] = None, include_search_results: bool = False, + search_errors: tuple[str, ...] = (), ) -> Dict[str, Any]: """Build the file_search_call output item (mirrors OpenAI's format). @@ -249,18 +288,20 @@ def _build_file_search_call_output( call_id: Unique ID for this file_search call. queries: List of search queries used. results: The raw search results (used when include_search_results=True). - include_search_results: Populate search_results when the caller passed + include_search_results: Populate results/search_results when the caller passed ``include=["file_search_call.results"]``. + search_errors: Errors raised while searching; a call that produced no results and + only errors is reported as failed instead of completed. """ - search_results = None - if include_search_results and results: - search_results = _build_search_results_for_include(results) + formatted_results = _build_search_results_for_include(results) if include_search_results and results else None + status = "failed" if search_errors and not results else "completed" return { "type": "file_search_call", "id": call_id, - "status": "completed", + "status": status, "queries": queries, - "search_results": search_results, + "results": formatted_results, + "search_results": formatted_results, } @@ -426,11 +467,12 @@ async def _execute_file_search_tool_calls( all_vs_ids: List[str], input: Any, file_search_call_id: str, -) -> Tuple[List[Dict[str, Any]], List[str], List[VectorStoreSearchResult]]: +) -> tuple[list[dict[str, Any]], list[str], list[VectorStoreSearchResult], tuple[str, ...]]: """Run the vector search for each file_search tool_call and collect results.""" tool_results: List[Dict[str, Any]] = [] all_queries: List[str] = [] all_results: List[VectorStoreSearchResult] = [] + all_errors: list[str] = [] for tool_call in file_search_calls: call_id, raw_args = _extract_tool_call_fields(tool_call, fallback_call_id=file_search_call_id) @@ -445,12 +487,13 @@ async def _execute_file_search_tool_calls( vs_id_arg = args.get("vector_store_id") vs_ids_for_call = [vs_id_arg] if vs_id_arg else all_vs_ids - queries, results = await _run_vector_searches( + queries, results, errors = await _run_vector_searches( queries=queries_from_call, vector_store_ids=vs_ids_for_call, ) all_queries.extend(queries) all_results.extend(results) + all_errors.extend(errors) tool_results.append( { @@ -460,7 +503,7 @@ async def _execute_file_search_tool_calls( } ) - return tool_results, all_queries, all_results + return tool_results, all_queries, all_results, tuple(all_errors) def _build_follow_up_input( @@ -560,7 +603,7 @@ async def aresponses_with_emulated_file_search( # 4. Execute each file_search tool call file_search_call_id = f"fs_{uuid.uuid4().hex[:24]}" - tool_results, all_queries, all_results = await _execute_file_search_tool_calls( + tool_results, all_queries, all_results, search_errors = await _execute_file_search_tool_calls( file_search_calls=file_search_calls, all_vs_ids=all_vs_ids, input=input, @@ -600,6 +643,7 @@ async def aresponses_with_emulated_file_search( queries=all_queries or [str(input)], results=all_results, include_search_results=_include_search_results, + search_errors=search_errors, ), message_output=_build_message_output(response_text, all_results), first_response=first_response, diff --git a/tests/test_litellm/llms/test_file_search_responses.py b/tests/test_litellm/llms/test_file_search_responses.py index 2f7ad3874fa..79865fa5f2b 100644 --- a/tests/test_litellm/llms/test_file_search_responses.py +++ b/tests/test_litellm/llms/test_file_search_responses.py @@ -935,3 +935,160 @@ class TestEmulatedFileSearchHandler: f"Sub-call {i} must run with is_internal_call=True to suppress " "billing callbacks in wrapper_async" ) + + +# --------------------------------------------------------------------------- +# Regression: emulated file_search must honour the vector store registry entry +# (https://github.com/BerriAI/litellm/issues/34768) +# --------------------------------------------------------------------------- + + +class TestEmulatedFileSearchRegistryResolution: + """The emulated handler must search a registered store with its own provider + params.""" + + @staticmethod + def _s3_vectors_registry(vector_store_id: str = "my-bucket:my-index"): + from litellm.types.vector_stores import LiteLLM_ManagedVectorStore + from litellm.vector_stores.vector_store_registry import VectorStoreRegistry + + return VectorStoreRegistry( + vector_stores=[ + LiteLLM_ManagedVectorStore( + vector_store_id=vector_store_id, + custom_llm_provider="s3_vectors", + litellm_params={ + "custom_llm_provider": "s3_vectors", + "vector_store_id": vector_store_id, + "vector_bucket_name": "my-bucket", + "index_name": "my-index", + "embedding_model": "bedrock/amazon.titan-embed-text-v2:0", + }, + ) + ] + ) + + @pytest.mark.asyncio + async def test_search_uses_registry_provider_and_params(self): + from litellm.responses.file_search.emulated_handler import _run_vector_searches + + search_result = MagicMock() + search_result.file_id = "file-1" + search_result.filename = "probe.txt" + search_result.score = 0.55 + search_result.content = [{"type": "text", "text": "Premium plan costs 25 dollars"}] + search_response = MagicMock() + search_response.data = [search_result] + + mock_asearch = AsyncMock(return_value=search_response) + with ( + patch("litellm.vector_store_registry", self._s3_vectors_registry()), + patch("litellm.vector_stores.main.asearch", new=mock_asearch), + ): + _queries, results, errors = await _run_vector_searches( + queries=["premium wifi plan price"], + vector_store_ids=["my-bucket:my-index"], + ) + + assert errors == () + assert results == [search_result] + call_kwargs = mock_asearch.call_args.kwargs + assert call_kwargs["vector_store_id"] == "my-bucket:my-index" + assert call_kwargs["custom_llm_provider"] == "s3_vectors" + assert call_kwargs["vector_bucket_name"] == "my-bucket" + assert call_kwargs["index_name"] == "my-index" + assert call_kwargs["embedding_model"] == "bedrock/amazon.titan-embed-text-v2:0" + + @pytest.mark.asyncio + async def test_unregistered_store_keeps_default_provider(self): + from litellm.responses.file_search.emulated_handler import _run_vector_searches + + search_response = MagicMock() + search_response.data = [] + mock_asearch = AsyncMock(return_value=search_response) + with ( + patch("litellm.vector_store_registry", None), + patch("litellm.vector_stores.main.asearch", new=mock_asearch), + ): + await _run_vector_searches(queries=["q"], vector_store_ids=["vs_openai"]) + + assert "custom_llm_provider" not in mock_asearch.call_args.kwargs + + @pytest.mark.asyncio + async def test_search_failure_is_surfaced_as_failed_file_search_call(self): + from litellm.responses.file_search.emulated_handler import ( + aresponses_with_emulated_file_search, + ) + + first_resp = MagicMock() + first_resp.output = [ + { + "type": "function_call", + "name": "litellm_file_search", + "call_id": "call_fail", + "arguments": '{"queries": ["premium wifi plan price"]}', + } + ] + first_resp.id = "resp_fail" + first_resp.created_at = 1700000000 + first_resp.model = "claude-haiku-4-5" + first_resp.usage = None + + final_resp = MagicMock() + final_resp.output = [ + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "I could not find that."}], + } + ] + final_resp.id = "resp_fail_2" + final_resp.created_at = 1700000000 + final_resp.model = "claude-haiku-4-5" + final_resp.usage = None + + with ( + patch("litellm.vector_store_registry", self._s3_vectors_registry()), + patch( + "litellm.responses.file_search.emulated_handler._call_aresponses", + new=AsyncMock(side_effect=[first_resp, final_resp]), + ), + patch( + "litellm.vector_stores.main.asearch", + new=AsyncMock(side_effect=Exception("AccessDeniedException")), + ), + ): + result = await aresponses_with_emulated_file_search( + input="how much is the premium wifi plan?", + model="anthropic/claude-haiku-4-5", + tools=[{"type": "file_search", "vector_store_ids": ["my-bucket:my-index"]}], + include=["file_search_call.results"], + ) + + file_search_call = result.output[0] + status = file_search_call["status"] if isinstance(file_search_call, dict) else file_search_call.status + assert status == "failed" + + @pytest.mark.asyncio + async def test_include_populates_openai_results_field(self): + from litellm.responses.file_search.emulated_handler import ( + _build_file_search_call_output, + ) + + result = MagicMock() + result.file_id = "file-1" + result.filename = "probe.txt" + result.score = 0.55 + result.attributes = {} + result.content = [{"type": "text", "text": "Premium plan costs 25 dollars"}] + + output = _build_file_search_call_output( + call_id="fs_1", + queries=["premium wifi plan price"], + results=[result], + include_search_results=True, + ) + + assert output["status"] == "completed" + assert output["results"] == output["search_results"] + assert output["results"][0]["text"] == "Premium plan costs 25 dollars" + assert output["results"][0]["score"] == 0.55 diff --git a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py index 15a117bd6fc..70e4bb2eeb3 100644 --- a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py @@ -327,3 +327,82 @@ def test_rag_query_stream_returns_event_stream(client_internal_user): assert response.headers.get("content-type", "").startswith("text/event-stream") assert '"object":"chat.completion.chunk"' in response.text assert "data: [DONE]" in response.text + + +@pytest.mark.asyncio +async def test_rag_ingest_persists_embedding_model_on_registry_entry(): + """ + /rag/ingest must record the embedding model it used on the new registry entry, otherwise + search-time query embedding falls back to text-embedding-3-small + (https://github.com/BerriAI/litellm/issues/34768). + """ + from litellm.proxy.rag_endpoints.endpoints import ( + _save_vector_store_to_db_from_rag_ingest, + ) + + repo = MagicMock() + repo.return_value.table.find_unique = AsyncMock(return_value=None) + create_in_db = AsyncMock() + + with ( + patch("litellm.proxy.rag_endpoints.endpoints.ManagedVectorStoresRepository", repo), + patch( + "litellm.proxy.vector_store_endpoints.management_endpoints.create_vector_store_in_db", + new=create_in_db, + ), + ): + await _save_vector_store_to_db_from_rag_ingest( + response={"status": "completed", "vector_store_id": "my-bucket:my-index", "file_id": "file-1"}, + ingest_options={ + "embedding": {"model": "bedrock/amazon.titan-embed-text-v2:0"}, + "vector_store": { + "custom_llm_provider": "s3_vectors", + "vector_bucket_name": "my-bucket", + "index_name": "my-index", + }, + }, + prisma_client=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(user_id="u1"), + file_data=("probe.txt", b"hello", "text/plain"), + ) + + litellm_params = create_in_db.call_args.kwargs["litellm_params"] + assert litellm_params["embedding_model"] == "bedrock/amazon.titan-embed-text-v2:0" + assert litellm_params["vector_bucket_name"] == "my-bucket" + assert litellm_params["index_name"] == "my-index" + + +@pytest.mark.asyncio +async def test_rag_ingest_keeps_explicit_vector_store_embedding_model(): + """An embedding model set on the vector store config wins over ingest_options.embedding.""" + from litellm.proxy.rag_endpoints.endpoints import ( + _save_vector_store_to_db_from_rag_ingest, + ) + + repo = MagicMock() + repo.return_value.table.find_unique = AsyncMock(return_value=None) + create_in_db = AsyncMock() + + with ( + patch("litellm.proxy.rag_endpoints.endpoints.ManagedVectorStoresRepository", repo), + patch( + "litellm.proxy.vector_store_endpoints.management_endpoints.create_vector_store_in_db", + new=create_in_db, + ), + ): + await _save_vector_store_to_db_from_rag_ingest( + response={"status": "completed", "vector_store_id": "my-bucket:my-index"}, + ingest_options={ + "embedding": {"model": "text-embedding-3-small"}, + "vector_store": { + "custom_llm_provider": "s3_vectors", + "vector_bucket_name": "my-bucket", + "embedding_model": "bedrock/amazon.titan-embed-text-v2:0", + }, + }, + prisma_client=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(user_id="u1"), + ) + + litellm_params = create_in_db.call_args.kwargs["litellm_params"] + assert litellm_params["embedding_model"] == "bedrock/amazon.titan-embed-text-v2:0"