fix(vector-store): route pre-call searches through router

litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py

@@ -80,10 +80,15 @@ class VectorStorePreCallHook(CustomLogger):
 
             # Get prisma_client for database fallback
             prisma_client = None
+            llm_router = None
             try:
-                from litellm.proxy.proxy_server import prisma_client as _prisma_client
+                from litellm.proxy.proxy_server import (
+                    llm_router as _llm_router,
+                    prisma_client as _prisma_client,
+                )
 
                 prisma_client = _prisma_client
+                llm_router = _llm_router
             except ImportError:
                 pass
 
@@ -114,12 +119,23 @@ class VectorStorePreCallHook(CustomLogger):
                 vector_store_id = vector_store_to_run.get("vector_store_id", "")
                 custom_llm_provider = vector_store_to_run.get("custom_llm_provider")
                 litellm_params_for_vector_store = vector_store_to_run.get("litellm_params", {}) or {}
-                # Call litellm.vector_stores.search() with the required parameters
-                search_response = await litellm.vector_stores.asearch(
+                request_litellm_params: Final = (
+                    litellm_logging_obj.model_call_details.get("litellm_params", {})
+                    if litellm_logging_obj is not None
+                    else {}
+                )
+                request_metadata: Final = (
+                    request_litellm_params.get("metadata", {}) if isinstance(request_litellm_params, dict) else {}
+                )
+                search_function: Final = (
+                    llm_router.avector_store_search if llm_router is not None else litellm.vector_stores.asearch
+                )
+                search_response = await search_function(
                     **{
                         "vector_store_id": vector_store_id,
                         "query": query,
                         "custom_llm_provider": custom_llm_provider,
+                        "metadata": request_metadata,
                         **litellm_params_for_vector_store,
                     },
                 )

tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py

@@ -71,6 +71,48 @@ def setup_vector_store_registry():
     )
 
 
+@pytest.mark.asyncio
+async def test_vector_store_hook_routes_search_through_proxy_router(
+    setup_vector_store_registry,
+):
+    proxy_router = Mock()
+    proxy_router.avector_store_search = AsyncMock(
+        return_value=VectorStoreSearchResponse(
+            object="vector_store.search_results.page",
+            search_query="what is litellm?",
+            data=[
+                VectorStoreSearchResult(
+                    score=1.0,
+                    content=[VectorStoreResultContent(text="routed context", type="text")],
+                )
+            ],
+        )
+    )
+    logging_obj = Mock()
+    logging_obj.model_call_details = {
+        "litellm_params": {"metadata": {"user_api_key_team_id": "team-a"}}
+    }
+
+    with patch("litellm.proxy.proxy_server.llm_router", proxy_router):
+        _, messages, _ = await VectorStorePreCallHook().async_get_chat_completion_prompt(
+            model="chat-model",
+            messages=[{"role": "user", "content": "what is litellm?"}],
+            non_default_params={"vector_store_ids": ["T37J8R4WTM"]},
+            prompt_id=None,
+            prompt_variables=None,
+            dynamic_callback_params={},
+            litellm_logging_obj=logging_obj,
+        )
+
+    proxy_router.avector_store_search.assert_awaited_once_with(
+        vector_store_id="T37J8R4WTM",
+        query="what is litellm?",
+        custom_llm_provider="bedrock",
+        metadata={"user_api_key_team_id": "team-a"},
+    )
+    assert messages[0]["content"] == "Context:\n\nrouted context\n\n"
+
+
 @pytest.mark.asyncio
 async def test_e2e_bedrock_knowledgebase_retrieval_with_completion(
     setup_vector_store_registry,