From 76176f2a643a01af7544d68fdb7ad85d2c313bf3 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 18 Mar 2026 11:26:27 +0530 Subject: [PATCH] fix(file_search): restore should_use_emulated helper, fix dedup, extract DB helper, clean docstring - Re-add should_use_emulated_file_search() to emulated_handler.py so H5/H6/H7/H13 tests don't fail with ImportError - Remove per-file-id deduplication from _build_search_results_for_include so all chunks are returned (matching OpenAI native file_search behaviour); update test_H14 to assert 2 results - Extract raw prisma DB query in check_vector_store_ids_access into a static _fetch_managed_vector_stores_by_uuids helper so the hot request path uses a named, testable function instead of an inline prisma_client.db.* call - Remove developer-local path from test module docstring Made-with: Cursor --- .../proxy/hooks/managed_files.py | 23 +++++++++++++-- .../responses/file_search/emulated_handler.py | 28 +++++++++++++++---- .../llms/test_file_search_responses.py | 16 ++++++----- 3 files changed, 52 insertions(+), 15 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 351fe05755..eecaddade7 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -741,6 +741,23 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return vs_ids + @staticmethod + async def _fetch_managed_vector_stores_by_uuids( + uuids: List[str], + prisma_client: Any, + ) -> List[Any]: + """ + Fetch managed vector store rows by their internal UUIDs. + + Isolated here so callers on the hot request path use a named helper + rather than a raw prisma_client.db.* call inline, keeping the + critical-path code auditable and the DB query easy to stub in tests. + """ + return await prisma_client.db.litellm_managedvectorstorestable.find_many( + where={"vector_store_id": {"in": uuids}}, + take=len(uuids), + ) + async def check_vector_store_ids_access( self, vector_store_ids: List[str], @@ -771,9 +788,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if not uuid_to_unified: return - rows = await prisma_client.db.litellm_managedvectorstorestable.find_many( - where={"vector_store_id": {"in": list(uuid_to_unified.keys())}}, - take=len(uuid_to_unified), + rows = await self._fetch_managed_vector_stores_by_uuids( + uuids=list(uuid_to_unified.keys()), + prisma_client=prisma_client, ) found_uuids = {row.vector_store_id for row in rows} diff --git a/litellm/responses/file_search/emulated_handler.py b/litellm/responses/file_search/emulated_handler.py index 584aea4a6a..62452539d1 100644 --- a/litellm/responses/file_search/emulated_handler.py +++ b/litellm/responses/file_search/emulated_handler.py @@ -26,6 +26,25 @@ ToolParam = Any FILE_SEARCH_FUNCTION_NAME = "litellm_file_search" +# --------------------------------------------------------------------------- +# Detection +# --------------------------------------------------------------------------- + +def should_use_emulated_file_search( + tools: Optional[Iterable[ToolParam]], + provider_config: Any, # BaseResponsesAPIConfig +) -> bool: + """Return True when there is a file_search tool and the provider can't handle it natively.""" + if not tools: + return False + has_fs = any( + isinstance(t, dict) and t.get("type") == "file_search" for t in tools + ) + if not has_fs: + return False + return provider_config is None or not provider_config.supports_native_file_search() + + # --------------------------------------------------------------------------- # Tool conversion # --------------------------------------------------------------------------- @@ -195,15 +214,14 @@ def _build_search_results_for_include( """ Convert VectorStoreSearchResult objects to the format expected in file_search_call.search_results (mirrors OpenAI's include= format). + + All chunks are returned — no deduplication by file_id — matching the + behaviour of OpenAI's native file_search which surfaces every relevant + chunk even when multiple chunks originate from the same document. """ formatted: List[Dict[str, Any]] = [] - seen_file_ids: set = set() for result in results: file_id = _get_field(result, "file_id") or "" - if file_id and file_id in seen_file_ids: - continue - if file_id: - seen_file_ids.add(file_id) content_items = _get_field(result, "content") or [] text_chunks = [ c.get("text", "") if isinstance(c, dict) else getattr(c, "text", "") diff --git a/tests/test_litellm/llms/test_file_search_responses.py b/tests/test_litellm/llms/test_file_search_responses.py index 63599781c1..3a8fa95be5 100644 --- a/tests/test_litellm/llms/test_file_search_responses.py +++ b/tests/test_litellm/llms/test_file_search_responses.py @@ -1,7 +1,5 @@ """ -Unit tests for Phase 1: file_search / vector_store support in the Responses API. - -Test plan reference: ~/.gstack/projects/BerriAI-litellm/sameerkankute-res-test-plan-*.md +Unit tests for file_search / vector_store support in the Responses API. Coverage: A1-A7 _decode_vector_store_ids_in_tools() @@ -10,6 +8,7 @@ Coverage: E1-E4 file_search guard in responses/main.py F1-F6 ManagedFiles hook access control G1-G3 get_vector_store_ids_from_file_search_tools() + H1-H14 emulated_handler unit tests """ import base64 @@ -659,7 +658,9 @@ class TestEmulatedFileSearchHandler: annotations = _build_file_citation_annotations([r1, r2], "text") assert len(annotations) == 1 - def test_H14_include_search_results_dedupes_by_file_id(self): + def test_H14_include_search_results_returns_all_chunks(self): + """All chunks are returned even when they originate from the same file, + matching OpenAI native file_search behaviour.""" from litellm.responses.file_search.emulated_handler import ( _build_search_results_for_include, ) @@ -670,15 +671,16 @@ class TestEmulatedFileSearchHandler: r1.score = 0.9 r1.attributes = {} r1.content = [{"type": "text", "text": "first hit"}] - r2.file_id = "file-abc" # same file appears for a second query + r2.file_id = "file-abc" # same file, different chunk from a second query r2.filename = "doc.pdf" r2.score = 0.85 r2.attributes = {} r2.content = [{"type": "text", "text": "second hit"}] search_results = _build_search_results_for_include([r1, r2]) - assert len(search_results) == 1 - assert search_results[0]["file_id"] == "file-abc" + assert len(search_results) == 2, "Both chunks should be returned, not deduplicated" + assert search_results[0]["text"] == "first hit" + assert search_results[1]["text"] == "second hit" # --- End-to-end (mocked) ---