scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

Merge pull request #39972 from BerriAI/litellm_lit_7027_emulated_file_search_scope

Browse source 
fix(file_search): scope emulated file_search to the request's vector stores
This commit is contained in:
Mateo Wang 2026-09-05 17:07:04 -07:00 • committed by GitHub
commit 0aa346cba5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 119 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
litellm/responses/file_search/emulated_handler.py
View file

@ -459,7 +459,14 @@ async def _execute_file_search_tool_calls(
queries_from_call = _resolve_queries_from_args(args, input)
vs_id_arg = args.get("vector_store_id")
vs_ids_for_call = [cast(str, vs_id_arg)] if vs_id_arg else all_vs_ids # cast-ok: model-supplied, as today
if vs_id_arg is not None and vs_id_arg not in all_vs_ids:
verbose_logger.warning(
"file_search emulated: model picked vector_store_id=%r outside the request's vector_store_ids %s; "
"searching the request's stores instead",
vs_id_arg,
all_vs_ids,
)
vs_ids_for_call = [cast(str, vs_id_arg)] if vs_id_arg in all_vs_ids else all_vs_ids # cast-ok: request id
queries, results = await _run_vector_searches(
queries=queries_from_call,

112
tests/test_litellm/llms/test_file_search_responses.py
View file

@ -8,10 +8,11 @@ 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
H1-H17 emulated_handler unit tests
"""
import base64
import logging
from typing import Any, Dict, List, Optional
from importlib import import_module
from unittest.mock import AsyncMock, MagicMock, patch
@ -936,3 +937,112 @@ class TestEmulatedFileSearchHandler:
f"Sub-call {i} must run with is_internal_call=True to suppress "
"billing callbacks in wrapper_async"
)
@pytest.mark.asyncio
async def test_H16_model_chosen_id_outside_request_is_not_searched(self, caplog):
"""A vector_store_id the model returns that the request did not list is never
searched; the request's own stores are searched instead, with a warning."""
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_unlisted",
"arguments": '{"queries": ["launch codeword"], "vector_store_id": "vs_unlisted"}',
}
]
first_resp.id = "resp_unlisted"
first_resp.created_at = 1700000000
first_resp.model = "claude-3-5-sonnet"
first_resp.usage = None
final_resp = self._make_mock_responses_api_response(text="done")
search_result = MagicMock()
search_result.file_id = "file-allowed"
search_result.filename = "allowed.txt"
search_result.score = 0.9
search_result.content = [{"type": "text", "text": "allowed context"}]
mock_search_response = MagicMock()
mock_search_response.data = [search_result]
mock_asearch = AsyncMock(return_value=mock_search_response)
with (
patch.object(
import_module("litellm.responses.file_search.emulated_handler"),
"_call_aresponses",
new=AsyncMock(side_effect=[first_resp, final_resp]),
),
patch("litellm.vector_stores.main.asearch", new=mock_asearch), # test-quality-ok: asserts store searched
caplog.at_level(logging.WARNING, logger="LiteLLM"),
):
await aresponses_with_emulated_file_search(
input="What is the launch codeword?",
model="anthropic/claude-3-5-sonnet",
tools=[{"type": "file_search", "vector_store_ids": ["vs_allowed"]}],
)
searched_ids = [c.kwargs["vector_store_id"] for c in mock_asearch.call_args_list]
assert searched_ids, "Expected the vector store to be searched at least once"
assert "vs_unlisted" not in searched_ids, "Handler searched a store the request did not list"
assert set(searched_ids) == {"vs_allowed"}
dropped_id_warnings = [r for r in caplog.records if "vs_unlisted" in r.getMessage()]
assert len(dropped_id_warnings) == 1, "Expected one warning naming the dropped model-picked id"
assert dropped_id_warnings[0].levelno == logging.WARNING
assert "vs_allowed" in dropped_id_warnings[0].getMessage()
@pytest.mark.asyncio
async def test_H17_model_chosen_id_within_request_narrows_search(self):
"""A vector_store_id the model returns that IS one of the request's stores is honored:
only that store is searched, not every store in the request."""
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_narrow",
"arguments": '{"queries": ["q"], "vector_store_id": "vs_two"}',
}
]
first_resp.id = "resp_narrow"
first_resp.created_at = 1700000000
first_resp.model = "claude-3-5-sonnet"
first_resp.usage = None
final_resp = self._make_mock_responses_api_response(text="done")
search_result = MagicMock()
search_result.file_id = "file-two"
search_result.filename = "two.txt"
search_result.score = 0.9
search_result.content = [{"type": "text", "text": "context"}]
mock_search_response = MagicMock()
mock_search_response.data = [search_result]
mock_asearch = AsyncMock(return_value=mock_search_response)
with (
patch.object(
import_module("litellm.responses.file_search.emulated_handler"),
"_call_aresponses",
new=AsyncMock(side_effect=[first_resp, final_resp]),
),
patch("litellm.vector_stores.main.asearch", new=mock_asearch), # test-quality-ok: asserts store searched
):
await aresponses_with_emulated_file_search(
input="q",
model="anthropic/claude-3-5-sonnet",
tools=[{"type": "file_search", "vector_store_ids": ["vs_one", "vs_two"]}],
)
searched_ids = [c.kwargs["vector_store_id"] for c in mock_asearch.call_args_list]
assert set(searched_ids) == {"vs_two"}, (
"A request-listed id the model picks should narrow the search to that store only"
)