mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-13 05:06:09 +00:00
[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
This commit is contained in:
Ishaan Jaff
@@ -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:
|
||||
|
||||
@@ -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],
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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]]
|
||||
|
||||
@@ -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
|
||||
|
||||
+43
-1
@@ -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]
|
||||
|
||||
|
||||
@@ -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]]
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user