From 28cb7cc0eded2b2019f421d3f578f42fbdcf01e1 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 2 May 2025 13:43:20 -0700 Subject: [PATCH] [Feat] KnowledgeBase/Vector Store - Log `StandardLoggingVectorStoreRequest` for requests made when a vector store is used (#10509) * ensure vector store results are logged in SLP * fix tests * fix tests with vector_store_request_metadata * fix linting --- litellm/integrations/custom_logger.py | 4 + .../langfuse/langfuse_prompt_management.py | 5 +- .../rag_hooks/bedrock_knowledgebase.py | 92 +++++++++++++++++-- litellm/litellm_core_utils/litellm_logging.py | 10 ++ litellm/proxy/_types.py | 2 + .../spend_tracking/spend_tracking_utils.py | 13 +++ litellm/types/utils.py | 44 ++++++++- litellm/types/vector_stores.py | 24 +++++ .../test_spend_management_endpoints.py | 6 +- .../gcs_pub_sub_body/spend_logs_payload.json | 2 +- .../standard_logging_payload.json | 3 +- .../test_bedrock_knowledgebase_hook.py | 66 +++++++++++++ .../test_otel_logging.py | 1 + 13 files changed, 257 insertions(+), 15 deletions(-) diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index c5295ada19..1fe25d7623 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -30,9 +30,12 @@ from litellm.types.utils import ( if TYPE_CHECKING: from opentelemetry.trace import Span as _Span + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + Span = Union[_Span, Any] else: Span = Any + LiteLLMLoggingObj = Any class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callback#callback-class @@ -80,6 +83,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac prompt_id: Optional[str], prompt_variables: Optional[dict], dynamic_callback_params: StandardCallbackDynamicParams, + litellm_logging_obj: LiteLLMLoggingObj, ) -> Tuple[str, List[AllMessageValues], dict]: """ Returns: diff --git a/litellm/integrations/langfuse/langfuse_prompt_management.py b/litellm/integrations/langfuse/langfuse_prompt_management.py index dcd3d9933a..8de38a0154 100644 --- a/litellm/integrations/langfuse/langfuse_prompt_management.py +++ b/litellm/integrations/langfuse/langfuse_prompt_management.py @@ -26,13 +26,15 @@ if TYPE_CHECKING: from langfuse import Langfuse from langfuse.client import ChatPromptClient, TextPromptClient + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + LangfuseClass: TypeAlias = Langfuse PROMPT_CLIENT = Union[TextPromptClient, ChatPromptClient] else: PROMPT_CLIENT = Any LangfuseClass = Any - + LiteLLMLoggingObj = Any in_memory_dynamic_logger_cache = DynamicLoggingCache() @@ -172,6 +174,7 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge prompt_id: Optional[str], prompt_variables: Optional[dict], dynamic_callback_params: StandardCallbackDynamicParams, + litellm_logging_obj: LiteLLMLoggingObj, ) -> Tuple[ str, List[AllMessageValues], diff --git a/litellm/integrations/rag_hooks/bedrock_knowledgebase.py b/litellm/integrations/rag_hooks/bedrock_knowledgebase.py index 4bd45e6b2e..7307a57c0e 100644 --- a/litellm/integrations/rag_hooks/bedrock_knowledgebase.py +++ b/litellm/integrations/rag_hooks/bedrock_knowledgebase.py @@ -26,6 +26,17 @@ from litellm.types.integrations.rag.bedrock_knowledgebase import ( BedrockKBRetrievalResult, ) from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage +from litellm.types.utils import StandardLoggingVectorStoreRequest +from litellm.types.vector_stores import ( + VectorStoreResultContent, + VectorStoreSearchResult, + VectorStorSearchResponse, +) + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +else: + LiteLLMLoggingObj = Any if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import StandardCallbackDynamicParams @@ -35,6 +46,7 @@ else: class BedrockKnowledgeBaseHook(CustomPromptManagement, BaseAWSLLM): CONTENT_PREFIX_STRING = "Context: \n\n" + CUSTOM_LLM_PROVIDER = "bedrock" def __init__( self, @@ -58,26 +70,90 @@ class BedrockKnowledgeBaseHook(CustomPromptManagement, BaseAWSLLM): prompt_id: Optional[str], prompt_variables: Optional[dict], dynamic_callback_params: StandardCallbackDynamicParams, + litellm_logging_obj: LiteLLMLoggingObj, ) -> Tuple[str, List[AllMessageValues], dict]: """ Retrieves the context from the Bedrock Knowledge Base and appends it to the messages. """ vector_store_ids = non_default_params.pop("vector_store_ids", None) + vector_store_request_metadata: List[StandardLoggingVectorStoreRequest] = [] if vector_store_ids: for vector_store_id in vector_store_ids: - response = await self.make_bedrock_kb_retrieve_request( + query = self._get_kb_query_from_messages(messages) + bedrock_kb_response = await self.make_bedrock_kb_retrieve_request( knowledge_base_id=vector_store_id, - query=self._get_kb_query_from_messages(messages), + query=query, + ) + verbose_logger.debug( + f"Bedrock Knowledge Base Response: {bedrock_kb_response}" ) - verbose_logger.debug(f"Bedrock Knowledge Base Response: {response}") - context_message = ( - self.get_chat_completion_message_from_bedrock_kb_response(response) + context_message, context_string = ( + self.get_chat_completion_message_from_bedrock_kb_response( + bedrock_kb_response + ) ) if context_message is not None: messages.append(context_message) + + ################################################################################################# + ########## LOGGING for Standard Logging Payload, Langfuse, s3, LiteLLM DB etc. ################## + ################################################################################################# + vector_store_search_response: VectorStorSearchResponse = ( + self.transform_bedrock_kb_response_to_vector_store_search_response( + bedrock_kb_response=bedrock_kb_response, query=query + ) + ) + vector_store_request_metadata.append( + StandardLoggingVectorStoreRequest( + vector_store_id=vector_store_id, + query=query, + vector_store_search_response=vector_store_search_response, + custom_llm_provider=self.CUSTOM_LLM_PROVIDER, + ) + ) + + litellm_logging_obj.model_call_details["vector_store_request_metadata"] = ( + vector_store_request_metadata + ) + return model, messages, non_default_params + def transform_bedrock_kb_response_to_vector_store_search_response( + self, + bedrock_kb_response: BedrockKBResponse, + query: str, + ) -> VectorStorSearchResponse: + """ + Transform a BedrockKBResponse to a VectorStorSearchResponse + """ + retrieval_results: Optional[List[BedrockKBRetrievalResult]] = ( + bedrock_kb_response.get("retrievalResults", None) + ) + vector_store_search_response: VectorStorSearchResponse = ( + VectorStorSearchResponse(search_query=query, data=[]) + ) + if retrieval_results is None: + return vector_store_search_response + + vector_search_response_data: List[VectorStoreSearchResult] = [] + for retrieval_result in retrieval_results: + content: Optional[BedrockKBContent] = retrieval_result.get("content", None) + if content is None: + continue + content_text: Optional[str] = content.get("text", None) + if content_text is None: + continue + vector_store_search_result: VectorStoreSearchResult = ( + VectorStoreSearchResult( + score=retrieval_result.get("score", None), + content=[VectorStoreResultContent(text=content_text, type="text")], + ) + ) + vector_search_response_data.append(vector_store_search_result) + vector_store_search_response["data"] = vector_search_response_data + return vector_store_search_response + def _get_kb_query_from_messages(self, messages: List[AllMessageValues]) -> str: """ Uses the text `content` field of the last message in the list of messages @@ -256,7 +332,7 @@ class BedrockKnowledgeBaseHook(CustomPromptManagement, BaseAWSLLM): @staticmethod def get_chat_completion_message_from_bedrock_kb_response( response: BedrockKBResponse, - ) -> Optional[ChatCompletionUserMessage]: + ) -> Tuple[Optional[ChatCompletionUserMessage], str]: """ Retrieves the context from the Bedrock Knowledge Base response and returns a ChatCompletionUserMessage object. """ @@ -264,7 +340,7 @@ class BedrockKnowledgeBaseHook(CustomPromptManagement, BaseAWSLLM): "retrievalResults", None ) if retrieval_results is None: - return None + return None, "" # string to combine the context from the knowledge base context_string: str = BedrockKnowledgeBaseHook.CONTENT_PREFIX_STRING @@ -284,4 +360,4 @@ class BedrockKnowledgeBaseHook(CustomPromptManagement, BaseAWSLLM): role="user", content=context_string, ) - return message + return message, context_string diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 61d4f65c87..4f7f43c528 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -91,6 +91,7 @@ from litellm.types.utils import ( StandardLoggingPayloadErrorInformation, StandardLoggingPayloadStatus, StandardLoggingPromptManagementMetadata, + StandardLoggingVectorStoreRequest, TextCompletionResponse, TranscriptionResponse, Usage, @@ -551,6 +552,7 @@ class Logging(LiteLLMLoggingBaseClass): prompt_id=prompt_id, prompt_variables=prompt_variables, dynamic_callback_params=self.standard_callback_dynamic_params, + litellm_logging_obj=self, ) self.messages = messages return model, messages, non_default_params @@ -3286,6 +3288,9 @@ class StandardLoggingPayloadSetup: prompt_integration: Optional[str] = None, applied_guardrails: Optional[List[str]] = None, mcp_tool_call_metadata: Optional[StandardLoggingMCPToolCall] = None, + vector_store_request_metadata: Optional[ + List[StandardLoggingVectorStoreRequest] + ] = None, usage_object: Optional[dict] = None, ) -> StandardLoggingMetadata: """ @@ -3334,6 +3339,7 @@ class StandardLoggingPayloadSetup: prompt_management_metadata=prompt_management_metadata, applied_guardrails=applied_guardrails, mcp_tool_call_metadata=mcp_tool_call_metadata, + vector_store_request_metadata=vector_store_request_metadata, usage_object=usage_object, ) if isinstance(metadata, dict): @@ -3694,6 +3700,9 @@ def get_standard_logging_object_payload( prompt_integration=kwargs.get("prompt_integration", None), applied_guardrails=kwargs.get("applied_guardrails", None), mcp_tool_call_metadata=kwargs.get("mcp_tool_call_metadata", None), + vector_store_request_metadata=kwargs.get( + "vector_store_request_metadata", None + ), usage_object=usage.model_dump(), ) @@ -3838,6 +3847,7 @@ def get_standard_logging_metadata( prompt_management_metadata=None, applied_guardrails=None, mcp_tool_call_metadata=None, + vector_store_request_metadata=None, usage_object=None, ) if isinstance(metadata, dict): diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index beabe11434..cdc408a5de 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -31,6 +31,7 @@ from litellm.types.utils import ( StandardLoggingModelInformation, StandardLoggingPayloadErrorInformation, StandardLoggingPayloadStatus, + StandardLoggingVectorStoreRequest, StandardPassThroughResponseObject, TextCompletionResponse, ) @@ -1990,6 +1991,7 @@ class SpendLogsMetadata(TypedDict): requester_ip_address: Optional[str] applied_guardrails: Optional[List[str]] mcp_tool_call_metadata: Optional[StandardLoggingMCPToolCall] + vector_store_request_metadata: Optional[List[StandardLoggingVectorStoreRequest]] status: StandardLoggingPayloadStatus proxy_server_request: Optional[str] batch_models: Optional[List[str]] diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index cde2c43559..a6308a73f8 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -17,6 +17,7 @@ from litellm.types.utils import ( StandardLoggingMCPToolCall, StandardLoggingModelInformation, StandardLoggingPayload, + StandardLoggingVectorStoreRequest, ) from litellm.utils import get_end_user_id_for_cost_tracking @@ -43,6 +44,9 @@ def _get_spend_logs_metadata( applied_guardrails: Optional[List[str]] = None, batch_models: Optional[List[str]] = None, mcp_tool_call_metadata: Optional[StandardLoggingMCPToolCall] = None, + vector_store_request_metadata: Optional[ + List[StandardLoggingVectorStoreRequest] + ] = None, usage_object: Optional[dict] = None, model_map_information: Optional[StandardLoggingModelInformation] = None, ) -> SpendLogsMetadata: @@ -63,6 +67,7 @@ def _get_spend_logs_metadata( proxy_server_request=None, batch_models=None, mcp_tool_call_metadata=None, + vector_store_request_metadata=None, model_map_information=None, usage_object=None, ) @@ -82,6 +87,7 @@ def _get_spend_logs_metadata( clean_metadata["applied_guardrails"] = applied_guardrails clean_metadata["batch_models"] = batch_models clean_metadata["mcp_tool_call_metadata"] = mcp_tool_call_metadata + clean_metadata["vector_store_request_metadata"] = vector_store_request_metadata clean_metadata["usage_object"] = usage_object clean_metadata["model_map_information"] = model_map_information return clean_metadata @@ -226,6 +232,13 @@ def get_logging_payload( # noqa: PLR0915 if standard_logging_payload is not None else None ), + vector_store_request_metadata=( + standard_logging_payload["metadata"].get( + "vector_store_request_metadata", None + ) + if standard_logging_payload is not None + else None + ), usage_object=( standard_logging_payload["metadata"].get("usage_object", None) if standard_logging_payload is not None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 5bd8dfc50d..1cee451192 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2,7 +2,17 @@ import json import time import uuid from enum import Enum -from typing import Any, Dict, List, Literal, Mapping, Optional, Tuple, Union +from typing import ( + TYPE_CHECKING, + Any, + Dict, + List, + Literal, + Mapping, + Optional, + Tuple, + Union, +) from aiohttp import FormData from openai._models import BaseModel as OpenAIObject @@ -41,6 +51,11 @@ from .llms.openai import ( ) from .rerank import RerankResponse +if TYPE_CHECKING: + from .vector_stores import VectorStorSearchResponse +else: + VectorStorSearchResponse = Any + def _generate_id(): # private helper function return "chatcmpl-" + str(uuid.uuid4()) @@ -1705,6 +1720,32 @@ class StandardLoggingMCPToolCall(TypedDict, total=False): """ +class StandardLoggingVectorStoreRequest(TypedDict, total=False): + """ + Logging information for a vector store request/payload + """ + + vector_store_id: Optional[str] + """ + ID of the vector store + """ + + custom_llm_provider: Optional[str] + """ + Custom LLM provider the vector store is associated with eg. bedrock, openai, anthropic, etc. + """ + + query: Optional[str] + """ + Query to the vector store + """ + + vector_store_search_response: Optional[VectorStorSearchResponse] + """ + OpenAI format vector store search response + """ + + class StandardBuiltInToolsParams(TypedDict, total=False): """ Standard built-in OpenAItools parameters @@ -1736,6 +1777,7 @@ class StandardLoggingMetadata(StandardLoggingUserAPIKeyMetadata): requester_metadata: Optional[dict] prompt_management_metadata: Optional[StandardLoggingPromptManagementMetadata] mcp_tool_call_metadata: Optional[StandardLoggingMCPToolCall] + vector_store_request_metadata: Optional[List[StandardLoggingVectorStoreRequest]] applied_guardrails: Optional[List[str]] usage_object: Optional[dict] diff --git a/litellm/types/vector_stores.py b/litellm/types/vector_stores.py index 3841971775..25128930cf 100644 --- a/litellm/types/vector_stores.py +++ b/litellm/types/vector_stores.py @@ -61,3 +61,27 @@ class VectorStoreDeleteRequest(BaseModel): class VectorStoreInfoRequest(BaseModel): vector_store_id: str + + +class VectorStoreResultContent(TypedDict, total=False): + """Content of a vector store result""" + + text: Optional[str] + type: Optional[str] + + +class VectorStoreSearchResult(TypedDict, total=False): + """Result of a vector store search""" + + score: Optional[float] + content: Optional[List[VectorStoreResultContent]] + + +class VectorStorSearchResponse(TypedDict, total=False): + """Response after searching a vector store""" + + object: Literal[ + "vector_store.search_results.page" + ] # Always "vector_store.search_results.page" + search_query: Optional[str] + data: Optional[List[VectorStoreSearchResult]] diff --git a/tests/litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 96eb254ed1..3f04bfe154 100644 --- a/tests/litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -469,7 +469,7 @@ class TestSpendLogsPayload: "model": "gpt-4o", "user": "", "team_id": "", - "metadata": '{"applied_guardrails": [], "batch_models": null, "mcp_tool_call_metadata": null, "usage_object": {"completion_tokens": 20, "prompt_tokens": 10, "total_tokens": 30, "completion_tokens_details": null, "prompt_tokens_details": null}, "model_map_information": {"model_map_key": "gpt-4o", "model_map_value": {"key": "gpt-4o", "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, "input_cost_per_token": 2.5e-06, "cache_creation_input_token_cost": null, "cache_read_input_token_cost": 1.25e-06, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": 1.25e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token": 1e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_reasoning_token": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "openai", "mode": "chat", "supports_system_messages": true, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": false, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": false, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": true, "supports_reasoning": false, "search_context_cost_per_query": {"search_context_size_low": 0.03, "search_context_size_medium": 0.035, "search_context_size_high": 0.05}, "tpm": null, "rpm": null, "supported_openai_params": ["frequency_penalty", "logit_bias", "logprobs", "top_logprobs", "max_tokens", "max_completion_tokens", "modalities", "prediction", "n", "presence_penalty", "seed", "stop", "stream", "stream_options", "temperature", "top_p", "tools", "tool_choice", "function_call", "functions", "max_retries", "extra_headers", "parallel_tool_calls", "audio", "response_format", "user"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": null}}', + "metadata": '{"applied_guardrails": [], "batch_models": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "usage_object": {"completion_tokens": 20, "prompt_tokens": 10, "total_tokens": 30, "completion_tokens_details": null, "prompt_tokens_details": null}, "model_map_information": {"model_map_key": "gpt-4o", "model_map_value": {"key": "gpt-4o", "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, "input_cost_per_token": 2.5e-06, "cache_creation_input_token_cost": null, "cache_read_input_token_cost": 1.25e-06, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": 1.25e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token": 1e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_reasoning_token": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "openai", "mode": "chat", "supports_system_messages": true, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": false, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": false, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": true, "supports_reasoning": false, "search_context_cost_per_query": {"search_context_size_low": 0.03, "search_context_size_medium": 0.035, "search_context_size_high": 0.05}, "tpm": null, "rpm": null, "supported_openai_params": ["frequency_penalty", "logit_bias", "logprobs", "top_logprobs", "max_tokens", "max_completion_tokens", "modalities", "prediction", "n", "presence_penalty", "seed", "stop", "stream", "stream_options", "temperature", "top_p", "tools", "tool_choice", "function_call", "functions", "max_retries", "extra_headers", "parallel_tool_calls", "audio", "response_format", "user"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": null}}', "cache_key": "Cache OFF", "spend": 0.00022500000000000002, "total_tokens": 30, @@ -560,7 +560,7 @@ class TestSpendLogsPayload: "model": "claude-3-7-sonnet-20250219", "user": "", "team_id": "", - "metadata": '{"applied_guardrails": [], "batch_models": null, "mcp_tool_call_metadata": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-3-7-sonnet-20250219", "model_map_value": {"key": "claude-3-7-sonnet-20250219", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}', + "metadata": '{"applied_guardrails": [], "batch_models": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-3-7-sonnet-20250219", "model_map_value": {"key": "claude-3-7-sonnet-20250219", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}', "cache_key": "Cache OFF", "spend": 0.01383, "total_tokens": 2598, @@ -649,7 +649,7 @@ class TestSpendLogsPayload: "model": "claude-3-7-sonnet-20250219", "user": "", "team_id": "", - "metadata": '{"applied_guardrails": [], "batch_models": null, "mcp_tool_call_metadata": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-3-7-sonnet-20250219", "model_map_value": {"key": "claude-3-7-sonnet-20250219", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}', + "metadata": '{"applied_guardrails": [], "batch_models": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-3-7-sonnet-20250219", "model_map_value": {"key": "claude-3-7-sonnet-20250219", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}', "cache_key": "Cache OFF", "spend": 0.01383, "total_tokens": 2598, diff --git a/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json b/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json index 8d0815ff69..bdb5b21978 100644 --- a/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json +++ b/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json @@ -10,7 +10,7 @@ "model": "gpt-4o", "user": "", "team_id": "", - "metadata": "{\"applied_guardrails\": [], \"batch_models\": null, \"mcp_tool_call_metadata\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}}", + "metadata": "{\"applied_guardrails\": [], \"batch_models\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}}", "cache_key": "Cache OFF", "spend": 0.00022500000000000002, "total_tokens": 30, diff --git a/tests/logging_callback_tests/gcs_pub_sub_body/standard_logging_payload.json b/tests/logging_callback_tests/gcs_pub_sub_body/standard_logging_payload.json index 1dc72b704f..26ab12bad9 100644 --- a/tests/logging_callback_tests/gcs_pub_sub_body/standard_logging_payload.json +++ b/tests/logging_callback_tests/gcs_pub_sub_body/standard_logging_payload.json @@ -26,7 +26,8 @@ "user_api_key_end_user_id": null, "prompt_management_metadata": null, "applied_guardrails": [], - "mcp_tool_call_metadata": null + "mcp_tool_call_metadata": null, + "vector_store_request_metadata": null }, "cache_key": null, "response_cost": 0.00022500000000000002, diff --git a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py index 4ef9263d23..e41c6924a6 100644 --- a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py +++ b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py @@ -11,6 +11,7 @@ import gzip import json import logging import time +from typing import Optional, List from unittest.mock import AsyncMock, patch, Mock import pytest @@ -20,6 +21,18 @@ from litellm import completion from litellm._logging import verbose_logger from litellm.integrations.rag_hooks.bedrock_knowledgebase import BedrockKnowledgeBaseHook from litellm.llms.custom_httpx.http_handler import HTTPHandler, AsyncHTTPHandler +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.utils import StandardLoggingPayload, StandardLoggingVectorStoreRequest +from litellm.types.vector_stores import VectorStorSearchResponse + +class TestCustomLogger(CustomLogger): + def __init__(self): + self.standard_logging_payload: Optional[StandardLoggingPayload] = None + super().__init__() + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.standard_logging_payload = kwargs.get("standard_logging_object") + pass @pytest.mark.asyncio @@ -144,3 +157,56 @@ async def test_openai_with_knowledge_base_mock_openai(): assert BedrockKnowledgeBaseHook.CONTENT_PREFIX_STRING in messages[1]["content"] +@pytest.mark.asyncio +async def test_logging_with_knowledge_base_hook(): + """ + Test that the knowledge base request was logged in standard logging payload + """ + test_custom_logger = TestCustomLogger() + litellm.callbacks = [BedrockKnowledgeBaseHook(), test_custom_logger] + litellm.set_verbose = True + await litellm.acompletion( + model="gpt-4", + 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: VectorStorSearchResponse = 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 + + diff --git a/tests/logging_callback_tests/test_otel_logging.py b/tests/logging_callback_tests/test_otel_logging.py index ecbeef5d88..29f9d30517 100644 --- a/tests/logging_callback_tests/test_otel_logging.py +++ b/tests/logging_callback_tests/test_otel_logging.py @@ -276,6 +276,7 @@ def validate_redacted_message_span_attributes(span): "metadata.user_api_key_user_email", "metadata.applied_guardrails", "metadata.mcp_tool_call_metadata", + "metadata.vector_store_request_metadata", ] _all_attributes = set(