litellm/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py

import os



import litellm
import litellm.vector_stores.main
import json
from typing import Optional

import pytest

import litellm
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.utils import (
    StandardLoggingPayload,
)


class MockCustomLogger(CustomLogger):
    def __init__(self):
        self.standard_logging_payload: Optional[StandardLoggingPayload] = None
        self.completion_logging_payload: Optional[StandardLoggingPayload] = None
        super().__init__()

    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
        payload = kwargs.get("standard_logging_object")
        # Store the payload - completion calls have call_type='acompletion'
        if payload and payload.get("call_type") == "acompletion":
            self.completion_logging_payload = payload
        self.standard_logging_payload = payload
        pass


@pytest.fixture(autouse=True)
def add_aws_region_to_env(monkeypatch):
    monkeypatch.setenv("AWS_REGION", "us-west-2")


@pytest.fixture
def setup_vector_store_registry():
    from litellm.vector_stores.vector_store_registry import (
        VectorStoreRegistry,
        LiteLLM_ManagedVectorStore,
    )

    # Init vector store registry
    litellm.vector_store_registry = VectorStoreRegistry(
        vector_stores=[
            LiteLLM_ManagedVectorStore(
                vector_store_id="T37J8R4WTM", custom_llm_provider="bedrock"
            )
        ]
    )


@pytest.mark.asyncio
async def test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call(
    setup_vector_store_registry,
):
    """
    Test that the Bedrock Knowledge Base Hook works when making a real llm api call and returns citations.
    """

    # Init client
    litellm.turn_on_debug()
    async_client = AsyncHTTPHandler()
    response = await litellm.acompletion(
        model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
        messages=[{"role": "user", "content": "what is litellm?"}],
        vector_store_ids=["T37J8R4WTM"],
        client=async_client,
    )
    print("OPENAI RESPONSE:", json.dumps(dict(response), indent=4, default=str))
    assert response is not None

    # Check that search_results are present in provider_specific_fields
    assert hasattr(response.choices[0].message, "provider_specific_fields")
    provider_fields = response.choices[0].message.provider_specific_fields
    assert provider_fields is not None
    assert "search_results" in provider_fields
    search_results = provider_fields["search_results"]
    assert search_results is not None
    assert len(search_results) > 0

    # Check search result structure (OpenAI-compatible format)
    first_search_result = search_results[0]
    assert "object" in first_search_result
    assert first_search_result["object"] == "vector_store.search_results.page"
    assert "data" in first_search_result
    assert len(first_search_result["data"]) > 0

    # Check individual result structure
    first_result = first_search_result["data"][0]
    assert "score" in first_result
    assert "content" in first_result
    print(f"Search results returned: {len(search_results)}")
    print(f"First search result has {len(first_search_result['data'])} items")


@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()
    response = await litellm.acompletion(
        model=f"anthropic/{os.environ.get('CI_CD_DEFAULT_ANTHROPIC_MODEL', 'claude-haiku-4-5-20251001')}",
        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_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call_with_tools_and_filters(
    setup_vector_store_registry,
):
    """
    Test that filters from file_search tools are properly passed through to vector store search.
    This test verifies the entire flow: tool parsing -> filter extraction -> vector store API call.

    In this case we filter for a non-existent user_id, which should return no results.
    """
    litellm.turn_on_debug()

    response = await litellm.acompletion(
        model=f"anthropic/{os.environ.get('CI_CD_DEFAULT_ANTHROPIC_MODEL', 'claude-haiku-4-5-20251001')}",
        messages=[{"role": "user", "content": "what is litellm?"}],
        max_tokens=10,
        tools=[
            {
                "type": "file_search",
                "vector_store_ids": ["T37J8R4WTM"],
                "filters": {
                    "key": "user_id",
                    "value": "fake-user-id",
                    "operator": "eq",
                },
            }
        ],
    )

    # Verify response is not None
    assert response is not None

    # Verify search results were added to the response (this proves the search was called)
    assert hasattr(response.choices[0].message, "provider_specific_fields")
    provider_fields = response.choices[0].message.provider_specific_fields
    assert provider_fields is not None
    assert (
        "search_results" in provider_fields
    ), "search_results not in provider_specific_fields"

    search_results = provider_fields["search_results"]
    assert (
        search_results is not None and len(search_results) > 0
    ), "No search results found"

    # The search was performed - this confirms filters were passed through
    # The logs above show:  litellm.asearch(... filters={'key': 'user_id', 'value': 'fake-user-id', 'operator': 'eq'})
    # And the Bedrock API request contains: {'filter': {'equals': {'key': 'user_id', 'value': 'fake-user-id'}}}

    print("✅ Filters were successfully passed through to vector store search")
    print(f"   Search was performed and {len(search_results)} result(s) returned")


# @pytest.mark.asyncio
# async def test_logging_with_knowledge_base_hook(setup_vector_store_registry):
#     """
#     Test that the knowledge base request was logged in standard logging payload
#     """
#     test_custom_logger = MockCustomLogger()
#     litellm.set_verbose = True
#     await litellm.acompletion(
#         model="gpt-5.5",
#         messages=[{"role": "user", "content": "what is litellm?"}],
#         vector_store_ids = [
#             "T37J8R4WTM"
#         ],
#     )

#     # sleep for 1 second to allow the logging callback to run
#     await asyncio.sleep(1)

#     # assert that the knowledge base request was logged in the standard logging payload
#     standard_logging_payload: Optional[StandardLoggingPayload] = test_custom_logger.standard_logging_payload
#     assert standard_logging_payload is not None


#     metadata = standard_logging_payload["metadata"]
#     standard_logging_vector_store_request_metadata: Optional[List[StandardLoggingVectorStoreRequest]] = metadata["vector_store_request_metadata"]

#     print("standard_logging_vector_store_request_metadata:", json.dumps(standard_logging_vector_store_request_metadata, indent=4, default=str))

#     # 1 vector store request was made, expect 1 vector store request metadata object
#     assert len(standard_logging_vector_store_request_metadata) == 1

#     # expect the vector store request metadata object to have the correct values
#     vector_store_request_metadata = standard_logging_vector_store_request_metadata[0]
#     assert vector_store_request_metadata.get("vector_store_id") == "T37J8R4WTM"
#     assert vector_store_request_metadata.get("query") == "what is litellm?"
#     assert vector_store_request_metadata.get("custom_llm_provider") == "bedrock"


#     vector_store_search_response: VectorStoreSearchResponse = vector_store_request_metadata.get("vector_store_search_response")
#     assert vector_store_search_response is not None
#     assert vector_store_search_response.get("search_query") == "what is litellm?"
#     assert len(vector_store_search_response.get("data", [])) >=0
#     for item in vector_store_search_response.get("data", []):
#         assert item.get("score") is not None
#         assert item.get("content") is not None
#         assert len(item.get("content", [])) >= 0
#         for content_item in item.get("content", []):
#             text_content = content_item.get("text")
#             assert text_content is not None
#             assert len(text_content) > 0