diff --git a/tests/test_litellm/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py b/tests/test_litellm/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py index 9c0ed38f1a3..4dd97d22822 100644 --- a/tests/test_litellm/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py +++ b/tests/test_litellm/integrations/vector_store_integrations/test_vector_store_pre_call_hook.py @@ -1,6 +1,7 @@ import logging +from collections.abc import Iterator from dataclasses import dataclass, field -from typing import Any +from typing import Protocol import pytest @@ -9,6 +10,7 @@ from litellm._logging import verbose_logger from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( VectorStorePreCallHook, ) +from litellm.types.llms.openai import AllMessageValues from litellm.types.vector_stores import ( VectorStoreResultContent, VectorStoreSearchResponse, @@ -36,11 +38,11 @@ def _search_response(text: str) -> VectorStoreSearchResponse: @dataclass class RecordingRouter: failing_vector_store_ids: frozenset[str] = frozenset() - calls: list[dict[str, Any]] = field(default_factory=list) + calls: list[dict[str, object]] = field(default_factory=list) - async def avector_store_search(self, **kwargs: Any) -> VectorStoreSearchResponse: + async def avector_store_search(self, **kwargs: object) -> VectorStoreSearchResponse: self.calls.append(kwargs) - vector_store_id = kwargs["vector_store_id"] + vector_store_id = str(kwargs["vector_store_id"]) if vector_store_id in self.failing_vector_store_ids: raise litellm.BadRequestError( message=f"no healthy deployments for {vector_store_id}", @@ -70,8 +72,12 @@ class RecordingHandler(logging.Handler): self.records.append(record) +class RegisterStores(Protocol): + def __call__(self, *vector_store_ids: str, custom_llm_provider: str = "bedrock") -> None: ... + + @pytest.fixture -def registry_with(monkeypatch: pytest.MonkeyPatch): +def registry_with(monkeypatch: pytest.MonkeyPatch) -> RegisterStores: def _register(*vector_store_ids: str, custom_llm_provider: str = "bedrock") -> None: monkeypatch.setattr( litellm, @@ -88,7 +94,7 @@ def registry_with(monkeypatch: pytest.MonkeyPatch): @pytest.fixture -def warnings(): +def warnings() -> Iterator[list[logging.LogRecord]]: handler = RecordingHandler() verbose_logger.addHandler(handler) yield handler.records @@ -96,11 +102,15 @@ def warnings(): class FakeLoggingObj: - def __init__(self, metadata: dict[str, Any]) -> None: - self.model_call_details: dict[str, Any] = {"litellm_params": {"metadata": metadata}} + def __init__(self, metadata: dict[str, str]) -> None: + self.model_call_details: dict[str, object] = {"litellm_params": {"metadata": metadata}} -async def _run_hook(hook: VectorStorePreCallHook, vector_store_ids: list[str], logging_obj: FakeLoggingObj): +async def _run_hook( + hook: VectorStorePreCallHook, + vector_store_ids: list[str], + logging_obj: FakeLoggingObj, +) -> tuple[str, list[AllMessageValues], dict[str, object]]: return await hook.async_get_chat_completion_prompt( model="chat-model", messages=[{"role": "user", "content": "what is litellm?"}], @@ -113,7 +123,9 @@ async def _run_hook(hook: VectorStorePreCallHook, vector_store_ids: list[str], l @pytest.mark.asyncio -async def test_hook_searches_through_the_injected_router_with_the_request_metadata(registry_with): +async def test_hook_searches_through_the_injected_router_with_the_request_metadata( + registry_with: RegisterStores, +) -> None: """Regression (LIT-6752): the hook must reach the Router through its injected runtime, not a proxy_server import.""" registry_with("vs-router") router = RecordingRouter() @@ -137,7 +149,10 @@ async def test_hook_searches_through_the_injected_router_with_the_request_metada @pytest.mark.asyncio -async def test_hook_falls_back_to_the_sdk_when_the_runtime_has_no_router(registry_with, warnings): +async def test_hook_falls_back_to_the_sdk_when_the_runtime_has_no_router( + registry_with: RegisterStores, + warnings: list[logging.LogRecord], +) -> None: registry_with("vs-sdk", custom_llm_provider="lit6752-not-a-provider") _, messages, _ = await _run_hook( @@ -157,7 +172,7 @@ async def test_hook_falls_back_to_the_sdk_when_the_runtime_has_no_router(registr @pytest.mark.asyncio -async def test_every_healthy_vector_store_contributes_its_own_context(registry_with): +async def test_every_healthy_vector_store_contributes_its_own_context(registry_with: RegisterStores) -> None: """Regression (LIT-6752): each store appended its context to the original messages, so only the last one survived.""" registry_with("vs-one", "vs-two") router = RecordingRouter() @@ -176,7 +191,10 @@ async def test_every_healthy_vector_store_contributes_its_own_context(registry_w @pytest.mark.asyncio -async def test_a_failing_vector_store_warns_with_its_id_and_the_other_stores_still_answer(registry_with, warnings): +async def test_a_failing_vector_store_warns_with_its_id_and_the_other_stores_still_answer( + registry_with: RegisterStores, + warnings: list[logging.LogRecord], +) -> None: """Regression (LIT-6752): one unreachable store must not silently drop every other store's context.""" registry_with("vs-broken", "vs-healthy") router = RecordingRouter(failing_vector_store_ids=frozenset({"vs-broken"})) @@ -188,9 +206,12 @@ async def test_a_failing_vector_store_warns_with_its_id_and_the_other_stores_sti logging_obj, ) + search_results = logging_obj.model_call_details["search_results"] + assert [call["vector_store_id"] for call in router.calls] == ["vs-broken", "vs-healthy"] assert messages[0]["content"] == "Context:\n\ncontext from vs-healthy\n\n" - assert len(logging_obj.model_call_details["search_results"]) == 1 + assert isinstance(search_results, list) + assert len(search_results) == 1 assert [record.getMessage() for record in warnings] == [ "Vector store search failed for vector_store_id=vs-broken, continuing without its context: " "litellm.BadRequestError: no healthy deployments for vs-broken" @@ -198,7 +219,10 @@ async def test_a_failing_vector_store_warns_with_its_id_and_the_other_stores_sti @pytest.mark.asyncio -async def test_the_only_vector_store_failing_leaves_the_messages_untouched(registry_with, warnings): +async def test_the_only_vector_store_failing_leaves_the_messages_untouched( + registry_with: RegisterStores, + warnings: list[logging.LogRecord], +) -> None: registry_with("vs-broken") original_messages = [{"role": "user", "content": "what is litellm?"}]