From d8327b4740846c2f44c36f0e25c663fa4680c21d Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 15 Jul 2025 21:33:49 -0700
Subject: [PATCH] [Bug Fix] [Bug]: Knowledge Base Call returning error (#12628)

* bug fix using vector stores as tools

* test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call_with_tools
---
 litellm/litellm_core_utils/litellm_logging.py |  9 ++++---
 .../test_bedrock_knowledgebase_hook.py        | 24 +++++++++++++++++++
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py
index 54b8ec9fa9c..7caf726f1ef 100644
--- a/litellm/litellm_core_utils/litellm_logging.py
+++ b/litellm/litellm_core_utils/litellm_logging.py
@@ -600,7 +600,9 @@ class Logging(LiteLLMLoggingBaseClass):
         custom_logger = (
             prompt_management_logger
             or self.get_custom_logger_for_prompt_management(
-                model=model, non_default_params=non_default_params
+                model=model, 
+                tools=tools,
+                non_default_params=non_default_params
             )
         )
 
@@ -625,7 +627,7 @@ class Logging(LiteLLMLoggingBaseClass):
         return model, messages, non_default_params
 
     def get_custom_logger_for_prompt_management(
-        self, model: str, non_default_params: Dict
+        self, model: str, non_default_params: Dict, tools: Optional[List[Dict]] = None
     ) -> Optional[CustomLogger]:
         """
         Get a custom logger for prompt management based on model name or available callbacks.
@@ -673,7 +675,8 @@ class Logging(LiteLLMLoggingBaseClass):
         #########################################################
         if litellm.vector_store_registry is not None:
             if vector_store_to_run := litellm.vector_store_registry.get_vector_store_to_run(
-                non_default_params=non_default_params
+                non_default_params=non_default_params,
+                tools=tools
             ):
                 vector_store_custom_logger = (
                     litellm.ProviderConfigManager.get_provider_vector_store_config(
diff --git a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py
index 4c75d485277..5943c27b788 100644
--- a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py
+++ b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py
@@ -132,6 +132,30 @@ async def test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call(setup_vecto
     assert response is not None
 
 
+
+
+@pytest.mark.asyncio
+async def test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call_with_tools(setup_vector_store_registry):
+    """
+    Test that the Bedrock Knowledge Base Hook works when making a real llm api call
+    """
+    
+    # Init client
+    litellm._turn_on_debug()
+    litellm.callbacks = [BedrockVectorStore(aws_region_name="us-west-2")]
+    response = await litellm.acompletion(
+        model="anthropic/claude-3-5-haiku-latest",
+        messages=[{"role": "user", "content": "what is litellm?"}],
+        max_tokens=10,
+        tools=[
+            {
+                "type": "file_search",
+                "vector_store_ids": ["T37J8R4WTM"]
+            }
+        ],
+    )
+    assert response is not None
+
 @pytest.mark.asyncio
 async def test_openai_with_knowledge_base_mock_openai(setup_vector_store_registry):
     """