scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

test(prompt-management): fake the router boundary instead of patching litellm.vector_stores

Browse source 
Drop the redundant call-site comment Greptile flagged, and replace
patch("litellm.vector_stores.asearch") with a fake router exposing
avector_store_search, matching the DI pattern already used in
test_vector_store_pre_call_hook.py. Fixes the TQ008 test-quality-budget
CI failure (patching a litellm internal).

Signed-off-by: Pujitha Paladugu <10557236+pujitha24@users.noreply.github.com>
This commit is contained in:
Pujitha Paladugu 2026-09-18 07:37:07 -07:00
parent e9fc70e766
commit 748e06d29b
No known key found for this signature in database
1 changed files with 26 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

48
tests/unit/litellm_core_utils/test_litellm_logging.py
View file

@ -7290,7 +7290,6 @@ async def test_prompt_hooks_compose_vector_search_with_anthropic_cache_control(l
)
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
# main.py:504 seeds this key at runtime before the prompt-management hook is selected.
AnthropicCacheControlHook.maybe_seed_default_injection_points(
non_default_params=params, messages=messages, model=model, custom_llm_provider="bedrock"
)
@ -7301,35 +7300,40 @@ async def test_prompt_hooks_compose_vector_search_with_anthropic_cache_control(l
)
assert isinstance(selected, AnthropicCacheControlHook)
search: Final = AsyncMock(
return_value=VectorStoreSearchResponse(
object="vector_store.search_results.page",
search_query="What does the handbook say about refunds?",
data=[
VectorStoreSearchResult(
score=1.0,
content=[VectorStoreResultContent(text="Refunds take seven days", type="text")],
)
],
)
)
class _RecordingRouter:
def __init__(self) -> None:
self.calls: list[dict[str, object]] = []
async def avector_store_search(self, **kwargs: object) -> VectorStoreSearchResponse:
self.calls.append(kwargs)
return VectorStoreSearchResponse(
object="vector_store.search_results.page",
search_query="What does the handbook say about refunds?",
data=[
VectorStoreSearchResult(
score=1.0,
content=[VectorStoreResultContent(text="Refunds take seven days", type="text")],
)
],
)
router = _RecordingRouter()
runtime = MagicMock()
runtime.llm_router.return_value = None
runtime.llm_router.return_value = router
runtime.prisma_client.return_value = None
vector_store_hook = VectorStorePreCallHook(proxy_runtime=runtime)
previous_loggers = tuple(logging_module._in_memory_loggers)
logging_module._in_memory_loggers.clear()
logging_module._in_memory_loggers.append(vector_store_hook)
try:
with patch("litellm.vector_stores.asearch", search):
_, result_messages, remaining_params = await logging_obj.async_get_chat_completion_prompt(
model=model,
messages=messages,
non_default_params=params,
prompt_variables=None,
)
_, result_messages, remaining_params = await logging_obj.async_get_chat_completion_prompt(
model=model,
messages=messages,
non_default_params=params,
prompt_variables=None,
)
search.assert_awaited_once()
assert len(router.calls) == 1
assert result_messages[0]["content"] == "Context:\n\nRefunds take seven days\n\n"
assert result_messages[1]["content"] == "What does the handbook say about refunds?"
assert result_messages[1]["cache_control"] == {"type": "ephemeral"}