test(vector-stores): type the pre-call hook regression tests without Any

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?"}]