tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-09 16:24:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

Responses API - support tags in metadata

Browse Source 
* fix(ui/): fix routing for custom server root path

* fix: fix eslint errors

* fix(vector_store_pre_call_hook.py): Fix https://github.com/BerriAI/litellm/issues/15724

* fix(responses/main.py): have 'tags' work across metadata + litellm_metadata

* fix: add unit testing
This commit is contained in:
Krish Dholakia
2025-10-25 12:07:54 -07:00
committed by GitHub
parent f8d6a6edb9
commit 1543891763
7 changed files with 309 additions and 169 deletions
Download Patch File Download Diff File Expand all files Collapse all files
enterprise/litellm_enterprise/integrations/custom_guardrail.py
+1 -2
View File
@@ -29,11 +29,10 @@ class EnterpriseCustomGuardrailHelper:
if event_hook is None or not isinstance(event_hook, Mode):
return None
metadata: dict = data.get("litellm_metadata") or data.get("metadata", {})
proxy_server_request = data.get("proxy_server_request", {})
request_tags = StandardLoggingPayloadSetup._get_request_tags(
metadata=metadata,
litellm_params=data,
proxy_server_request=proxy_server_request,
)
enterprise/litellm_enterprise/integrations/prometheus.py
+6 -4
View File
@@ -1226,8 +1226,8 @@ class PrometheusLogger(CustomLogger):
try:
_tags = StandardLoggingPayloadSetup._get_request_tags(
request_data.get("metadata", {}),
request_data.get("proxy_server_request", {}),
litellm_params=request_data,
proxy_server_request=request_data.get("proxy_server_request", {}),
)
enum_values = UserAPIKeyLabelValues(
end_user=user_api_key_dict.end_user_id,
@@ -1289,7 +1289,8 @@ class PrometheusLogger(CustomLogger):
status_code="200",
route=user_api_key_dict.request_route,
tags=StandardLoggingPayloadSetup._get_request_tags(
data.get("metadata", {}), data.get("proxy_server_request", {})
litellm_params=data,
proxy_server_request=data.get("proxy_server_request", {}),
),
)
_labels = prometheus_label_factory(
@@ -2212,8 +2213,9 @@ class PrometheusLogger(CustomLogger):
)
# Create metrics ASGI app
if 'PROMETHEUS_MULTIPROC_DIR' in os.environ:
if "PROMETHEUS_MULTIPROC_DIR" in os.environ:
from prometheus_client import CollectorRegistry, multiprocess
registry = CollectorRegistry()
multiprocess.MultiProcessCollector(registry)
metrics_app = make_asgi_app(registry)
litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py
+126 -80
View File
@@ -25,6 +25,7 @@ if TYPE_CHECKING:
else:
LiteLLMLoggingObj = None
class VectorStorePreCallHook(CustomLogger):
CONTENT_PREFIX_STRING = "Context:\n\n"
"""
@@ -54,7 +55,7 @@ class VectorStorePreCallHook(CustomLogger):
) -> Tuple[str, List[AllMessageValues], dict]:
"""
Perform vector store search and append results as context to messages.
Args:
model: The model name
messages: List of messages
@@ -64,7 +65,7 @@ class VectorStorePreCallHook(CustomLogger):
dynamic_callback_params: Optional dynamic callback parameters
prompt_label: Optional prompt label
prompt_version: Optional prompt version
Returns:
Tuple of (model, modified_messages, non_default_params)
"""
@@ -73,134 +74,155 @@ class VectorStorePreCallHook(CustomLogger):
if litellm.vector_store_registry is None:
return model, messages, non_default_params
vector_stores_to_run: List[LiteLLM_ManagedVectorStore] = litellm.vector_store_registry.pop_vector_stores_to_run(
non_default_params=non_default_params, tools=tools
vector_stores_to_run: List[LiteLLM_ManagedVectorStore] = (
litellm.vector_store_registry.pop_vector_stores_to_run(
non_default_params=non_default_params, tools=tools
)
)
if not vector_stores_to_run:
return model, messages, non_default_params
# Extract the query from the last user message
query = self._extract_query_from_messages(messages)
if not query:
verbose_logger.debug("No query found in messages for vector store search")
verbose_logger.debug(
"No query found in messages for vector store search"
)
return model, messages, non_default_params
modified_messages: List[AllMessageValues] = messages.copy()
all_search_results: List[VectorStoreSearchResponse] = []
for vector_store_to_run in vector_stores_to_run:
# Get vector store id from the vector store config
vector_store_id = vector_store_to_run.get("vector_store_id", "")
custom_llm_provider = vector_store_to_run.get("custom_llm_provider")
litellm_params_for_vector_store = vector_store_to_run.get("litellm_params", {}) or {}
litellm_params_for_vector_store = (
vector_store_to_run.get("litellm_params", {}) or {}
)
# Call litellm.vector_stores.search() with the required parameters
search_response = await litellm.vector_stores.asearch(
vector_store_id=vector_store_id,
query=query,
custom_llm_provider=custom_llm_provider,
**litellm_params_for_vector_store
**{
"vector_store_id": vector_store_id,
"query": query,
"custom_llm_provider": custom_llm_provider,
**litellm_params_for_vector_store,
},
)
verbose_logger.debug(f"search_response: {search_response}")
# Store search results for later use in citations
all_search_results.append(search_response)
# Process search results and append as context
modified_messages = self._append_search_results_to_messages(
messages=messages,
search_response=search_response
messages=messages, search_response=search_response
)
# Get the number of results for logging
num_results = 0
num_results = len(search_response.get("data", []) or [])
verbose_logger.debug(f"Vector store search completed. Added context from {num_results} results")
verbose_logger.debug(
f"Vector store search completed. Added context from {num_results} results"
)
# Store search results as-is (already in OpenAI-compatible format)
if litellm_logging_obj and all_search_results:
litellm_logging_obj.model_call_details["search_results"] = all_search_results
litellm_logging_obj.model_call_details["search_results"] = (
all_search_results
)
return model, modified_messages, non_default_params
except Exception as e:
verbose_logger.exception(f"Error in VectorStorePreCallHook: {str(e)}")
# Return original parameters on error
return model, messages, non_default_params
def _extract_query_from_messages(self, messages: List[AllMessageValues]) -> Optional[str]:
def _extract_query_from_messages(
self, messages: List[AllMessageValues]
) -> Optional[str]:
"""
Extract the query from the last user message.
Args:
messages: List of messages
Returns:
The extracted query string or None if not found
"""
if not messages or len(messages) == 0:
return None
last_message = messages[-1]
if not isinstance(last_message, dict) or "content" not in last_message:
return None
content = last_message["content"]
if isinstance(content, str):
return content
elif isinstance(content, list) and len(content) > 0:
# Handle list of content items, extract text from first text item
for item in content:
if isinstance(item, dict) and item.get("type") == "text" and "text" in item:
if (
isinstance(item, dict)
and item.get("type") == "text"
and "text" in item
):
return item["text"]
return None
def _append_search_results_to_messages(
self,
messages: List[AllMessageValues],
search_response: VectorStoreSearchResponse
self,
messages: List[AllMessageValues],
search_response: VectorStoreSearchResponse,
) -> List[AllMessageValues]:
"""
Append search results as context to the messages.
Args:
messages: Original list of messages
search_response: Response from vector store search
Returns:
Modified list of messages with context appended
"""
search_response_data: Optional[List[VectorStoreSearchResult]] = search_response.get("data")
search_response_data: Optional[List[VectorStoreSearchResult]] = (
search_response.get("data")
)
if not search_response_data:
return messages
context_content = self.CONTENT_PREFIX_STRING
for result in search_response_data:
result_content: Optional[List[VectorStoreResultContent]] = result.get("content")
result_content: Optional[List[VectorStoreResultContent]] = result.get(
"content"
)
if result_content:
for content_item in result_content:
content_text: Optional[str] = content_item.get("text")
if content_text:
context_content += content_text + "\n\n"
# Only add context if we found any content
if context_content != "Context:\n\n":
# Create a copy of messages to avoid modifying the original
modified_messages = messages.copy()
# Add context as a new message before the last user message
context_message: ChatCompletionUserMessage = {
"role": "user",
"content": context_content
"role": "user",
"content": context_content,
}
modified_messages.insert(-1, cast(AllMessageValues, context_message))
return modified_messages
return messages
async def async_post_call_success_deployment_hook(
@@ -211,52 +233,65 @@ class VectorStorePreCallHook(CustomLogger):
) -> Optional[Any]:
"""
Add search results to the response after successful LLM call.
This hook adds the vector store search results (already in OpenAI-compatible format)
to the response's provider_specific_fields.
"""
try:
verbose_logger.debug("VectorStorePreCallHook.async_post_call_success_deployment_hook called")
verbose_logger.debug(
"VectorStorePreCallHook.async_post_call_success_deployment_hook called"
)
# Get logging object from request_data
litellm_logging_obj = request_data.get("litellm_logging_obj")
if not litellm_logging_obj:
verbose_logger.debug("No litellm_logging_obj in request_data")
return None
verbose_logger.debug(f"model_call_details keys: {list(litellm_logging_obj.model_call_details.keys())}")
verbose_logger.debug(
f"model_call_details keys: {list(litellm_logging_obj.model_call_details.keys())}"
)
# Get search results from model_call_details (already in OpenAI format)
search_results: Optional[List[VectorStoreSearchResponse]] = (
litellm_logging_obj.model_call_details.get("search_results")
)
verbose_logger.debug(f"Search results found: {search_results is not None}")
if not search_results:
verbose_logger.debug("No search results found")
return None
# Add search results to response object
if hasattr(response, "choices") and response.choices:
for choice in response.choices:
if hasattr(choice, "message") and choice.message:
# Get existing provider_specific_fields or create new dict
provider_fields = getattr(choice.message, "provider_specific_fields", None) or {}
provider_fields = (
getattr(choice.message, "provider_specific_fields", None)
or {}
)
# Add search results (already in OpenAI-compatible format)
provider_fields["search_results"] = search_results
# Set the provider_specific_fields
setattr(choice.message, "provider_specific_fields", provider_fields)
verbose_logger.debug(f"Added {len(search_results)} search results to response")
setattr(
choice.message, "provider_specific_fields", provider_fields
)
verbose_logger.debug(
f"Added {len(search_results)} search results to response"
)
# Return modified response
return response
except Exception as e:
verbose_logger.exception(f"Error adding search results to response: {str(e)}")
verbose_logger.exception(
f"Error adding search results to response: {str(e)}"
)
# Don't fail the request if search results fail to be added
return None
@@ -268,43 +303,54 @@ class VectorStorePreCallHook(CustomLogger):
) -> Optional[Any]:
"""
Add search results to the final streaming chunk.
This hook is called for the final streaming chunk, allowing us to add
search results to the stream before it's returned to the user.
"""
try:
verbose_logger.debug("VectorStorePreCallHook.async_post_call_streaming_deployment_hook called")
verbose_logger.debug(
"VectorStorePreCallHook.async_post_call_streaming_deployment_hook called"
)
# Get search results from model_call_details (already in OpenAI format)
search_results: Optional[List[VectorStoreSearchResponse]] = (
request_data.get("search_results")
)
verbose_logger.debug(f"Search results found for streaming chunk: {search_results is not None}")
verbose_logger.debug(
f"Search results found for streaming chunk: {search_results is not None}"
)
if not search_results:
verbose_logger.debug("No search results found for streaming chunk")
return response_chunk
# Add search results to streaming chunk
if hasattr(response_chunk, "choices") and response_chunk.choices:
for choice in response_chunk.choices:
if hasattr(choice, "delta") and choice.delta:
# Get existing provider_specific_fields or create new dict
provider_fields = getattr(choice.delta, "provider_specific_fields", None) or {}
provider_fields = (
getattr(choice.delta, "provider_specific_fields", None)
or {}
)
# Add search results (already in OpenAI-compatible format)
provider_fields["search_results"] = search_results
# Set the provider_specific_fields
choice.delta.provider_specific_fields = provider_fields
verbose_logger.debug(f"Added {len(search_results)} search results to streaming chunk")
verbose_logger.debug(
f"Added {len(search_results)} search results to streaming chunk"
)
# Return modified chunk
return response_chunk
except Exception as e:
verbose_logger.exception(f"Error adding search results to streaming chunk: {str(e)}")
verbose_logger.exception(
f"Error adding search results to streaming chunk: {str(e)}"
)
# Don't fail the request if search results fail to be added
return response_chunk
litellm/litellm_core_utils/litellm_logging.py
+28 -17
View File
@@ -701,8 +701,13 @@ class Logging(LiteLLMLoggingBaseClass):
vector_store_custom_logger.__class__.__name__
)
# Add to global callbacks so post-call hooks are invoked
if vector_store_custom_logger and vector_store_custom_logger not in litellm.callbacks:
litellm.logging_callback_manager.add_litellm_callback(vector_store_custom_logger)
if (
vector_store_custom_logger
and vector_store_custom_logger not in litellm.callbacks
):
litellm.logging_callback_manager.add_litellm_callback(
vector_store_custom_logger
)
return vector_store_custom_logger
return None
@@ -4447,12 +4452,18 @@ class StandardLoggingPayloadSetup:
return header_tags if header_tags else None
@staticmethod
def _get_request_tags(metadata: dict, proxy_server_request: dict) -> List[str]:
request_tags = (
metadata.get("tags", [])
if isinstance(metadata.get("tags", []), list)
else []
)
def _get_request_tags(
litellm_params: dict, proxy_server_request: dict
) -> List[str]:
# check for 'tags' in both 'metadata' and 'litellm_metadata'
metadata = litellm_params.get("metadata") or {}
litellm_metadata = litellm_params.get("litellm_metadata") or {}
if metadata.get("tags", []):
request_tags = metadata.get("tags", [])
elif litellm_metadata.get("tags", []):
request_tags = litellm_metadata.get("tags", [])
else:
request_tags = []
user_agent_tags = StandardLoggingPayloadSetup._get_user_agent_tags(
proxy_server_request
)
@@ -4466,11 +4477,10 @@ class StandardLoggingPayloadSetup:
return request_tags
def _get_status_fields(
status: StandardLoggingPayloadStatus,
guardrail_information: Optional[dict],
error_str: Optional[str]
error_str: Optional[str],
) -> "StandardLoggingPayloadStatusFields":
"""
Determine status fields based on request status and guardrail information.
@@ -4490,7 +4500,7 @@ def _get_status_fields(
"guardrail_intervened": "guardrail_intervened", # direct
"failure": "guardrail_failed_to_respond", # legacy
"guardrail_failed_to_respond": "guardrail_failed_to_respond", # direct
"not_run": "not_run"
"not_run": "not_run",
}
# Set LLM API status
@@ -4506,8 +4516,7 @@ def _get_status_fields(
guardrail_status = GUARDRAIL_STATUS_MAP.get(raw_status, "not_run")
return StandardLoggingPayloadStatusFields(
llm_api_status=llm_api_status,
guardrail_status=guardrail_status
llm_api_status=llm_api_status, guardrail_status=guardrail_status
)
@@ -4556,7 +4565,7 @@ def get_standard_logging_object_payload(
)
# standardize this function to be used across, s3, dynamoDB, langfuse logging
litellm_params = kwargs.get("litellm_params", {})
litellm_params = kwargs.get("litellm_params", {}) or {}
proxy_server_request = litellm_params.get("proxy_server_request") or {}
metadata: dict = (
@@ -4581,7 +4590,7 @@ def get_standard_logging_object_payload(
_model_group = metadata.get("model_group", "")
request_tags = StandardLoggingPayloadSetup._get_request_tags(
metadata=metadata, proxy_server_request=proxy_server_request
litellm_params=litellm_params, proxy_server_request=proxy_server_request
)
# cleanup timestamps
@@ -4677,8 +4686,10 @@ def get_standard_logging_object_payload(
status=status,
status_fields=_get_status_fields(
status=status,
guardrail_information=metadata.get("standard_logging_guardrail_information", None),
error_str=error_str
guardrail_information=metadata.get(
"standard_logging_guardrail_information", None
),
error_str=error_str,
),
custom_llm_provider=cast(Optional[str], kwargs.get("custom_llm_provider")),
saved_cache_cost=saved_cache_cost,
litellm/proxy/_new_secret_config.yaml
+2 -2
View File
@@ -1,5 +1,5 @@
model_list:
- model_name: bedrock-anthropic-claude-sonnet-4-5-20250929-v1
- model_name: gpt-5-mini
litellm_params:
model: bedrock/global.anthropic.claude-sonnet-4-5-20250929-v1:0
- model_name: embedding-model
@@ -18,4 +18,4 @@ vector_store_registry:
litellm_embedding_config:
api_base: https://krris-mh44uf7y-eastus2.cognitiveservices.azure.com/
api_key: os.environ/AZURE_API_KEY
api_version: "2025-09-01"
api_version: "2025-09-01"
litellm/responses/main.py
+29 -31
View File
@@ -40,7 +40,7 @@ from litellm.types.llms.openai import (
# Handle ResponseText import with fallback
if TYPE_CHECKING:
from litellm.types.llms.openai import ResponseText
from litellm.types.llms.openai import ResponseText # type: ignore
else:
ResponseText = str # Fallback for ResponseText import
from litellm.types.responses.main import *
@@ -52,9 +52,7 @@ if TYPE_CHECKING:
else:
MCPTool = Any
from .streaming_iterator import (
BaseResponsesAPIStreamingIterator,
)
from .streaming_iterator import BaseResponsesAPIStreamingIterator
####### ENVIRONMENT VARIABLES ###################
# Initialize any necessary instances or variables here
@@ -212,7 +210,6 @@ async def aresponses_api_with_mcp(
if stream and mcp_tools_with_litellm_proxy:
# Generate MCP discovery events using the already processed tools
from litellm._uuid import uuid
from litellm.responses.mcp.mcp_streaming_iterator import (
create_mcp_list_tools_events,
)
@@ -583,11 +580,11 @@ def responses(
)
# get provider config
responses_api_provider_config: Optional[
BaseResponsesAPIConfig
] = ProviderConfigManager.get_provider_responses_api_config(
model=model,
provider=litellm.LlmProviders(custom_llm_provider),
responses_api_provider_config: Optional[BaseResponsesAPIConfig] = (
ProviderConfigManager.get_provider_responses_api_config(
model=model,
provider=litellm.LlmProviders(custom_llm_provider),
)
)
local_vars.update(kwargs)
@@ -626,8 +623,9 @@ def responses(
user=user,
optional_params=dict(responses_api_request_params),
litellm_params={
"litellm_call_id": litellm_call_id,
**responses_api_request_params,
"litellm_call_id": litellm_call_id,
"metadata": metadata,
},
custom_llm_provider=custom_llm_provider,
)
@@ -779,11 +777,11 @@ def delete_responses(
raise ValueError("custom_llm_provider is required but passed as None")
# get provider config
responses_api_provider_config: Optional[
BaseResponsesAPIConfig
] = ProviderConfigManager.get_provider_responses_api_config(
model=None,
provider=litellm.LlmProviders(custom_llm_provider),
responses_api_provider_config: Optional[BaseResponsesAPIConfig] = (
ProviderConfigManager.get_provider_responses_api_config(
model=None,
provider=litellm.LlmProviders(custom_llm_provider),
)
)
if responses_api_provider_config is None:
@@ -958,11 +956,11 @@ def get_responses(
raise ValueError("custom_llm_provider is required but passed as None")
# get provider config
responses_api_provider_config: Optional[
BaseResponsesAPIConfig
] = ProviderConfigManager.get_provider_responses_api_config(
model=None,
provider=litellm.LlmProviders(custom_llm_provider),
responses_api_provider_config: Optional[BaseResponsesAPIConfig] = (
ProviderConfigManager.get_provider_responses_api_config(
model=None,
provider=litellm.LlmProviders(custom_llm_provider),
)
)
if responses_api_provider_config is None:
@@ -1114,11 +1112,11 @@ def list_input_items(
if custom_llm_provider is None:
raise ValueError("custom_llm_provider is required but passed as None")
responses_api_provider_config: Optional[
BaseResponsesAPIConfig
] = ProviderConfigManager.get_provider_responses_api_config(
model=None,
provider=litellm.LlmProviders(custom_llm_provider),
responses_api_provider_config: Optional[BaseResponsesAPIConfig] = (
ProviderConfigManager.get_provider_responses_api_config(
model=None,
provider=litellm.LlmProviders(custom_llm_provider),
)
)
if responses_api_provider_config is None:
@@ -1271,11 +1269,11 @@ def cancel_responses(
raise ValueError("custom_llm_provider is required but passed as None")
# get provider config
responses_api_provider_config: Optional[
BaseResponsesAPIConfig
] = ProviderConfigManager.get_provider_responses_api_config(
model=None,
provider=litellm.LlmProviders(custom_llm_provider),
responses_api_provider_config: Optional[BaseResponsesAPIConfig] = (
ProviderConfigManager.get_provider_responses_api_config(
model=None,
provider=litellm.LlmProviders(custom_llm_provider),
)
)
if responses_api_provider_config is None:
tests/test_litellm/litellm_core_utils/test_litellm_logging.py
+117 -33
View File
@@ -242,7 +242,7 @@ def test_get_request_tags():
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
tags = StandardLoggingPayloadSetup._get_request_tags(
metadata={"tags": ["test-tag"]},
litellm_params={"metadata": {"tags": ["test-tag"]}},
proxy_server_request={
"headers": {
"user-agent": "litellm/0.1.0",
@@ -255,6 +255,90 @@ def test_get_request_tags():
assert "User-Agent: litellm/0.1.0" in tags
def test_get_request_tags_from_metadata_and_litellm_metadata():
"""
Test that _get_request_tags correctly picks tags from both 'metadata' and 'litellm_metadata'.
Scenarios tested:
1. Tags in metadata only
2. Tags in litellm_metadata only
3. Tags in both (metadata should take priority)
4. No tags in either
5. None values for metadata/litellm_metadata
"""
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
# Test case 1: Tags in metadata only
tags = StandardLoggingPayloadSetup._get_request_tags(
litellm_params={"metadata": {"tags": ["metadata-tag-1", "metadata-tag-2"]}},
proxy_server_request={},
)
assert "metadata-tag-1" in tags
assert "metadata-tag-2" in tags
assert len([t for t in tags if not t.startswith("User-Agent:")]) == 2
# Test case 2: Tags in litellm_metadata only
tags = StandardLoggingPayloadSetup._get_request_tags(
litellm_params={
"litellm_metadata": {
"tags": ["litellm-metadata-tag-1", "litellm-metadata-tag-2"]
}
},
proxy_server_request={},
)
assert "litellm-metadata-tag-1" in tags
assert "litellm-metadata-tag-2" in tags
assert len([t for t in tags if not t.startswith("User-Agent:")]) == 2
# Test case 3: Tags in both - metadata should take priority
tags = StandardLoggingPayloadSetup._get_request_tags(
litellm_params={
"metadata": {"tags": ["metadata-tag"]},
"litellm_metadata": {"tags": ["litellm-metadata-tag"]},
},
proxy_server_request={},
)
assert "metadata-tag" in tags
assert "litellm-metadata-tag" not in tags
assert len([t for t in tags if not t.startswith("User-Agent:")]) == 1
# Test case 4: No tags in either
tags = StandardLoggingPayloadSetup._get_request_tags(
litellm_params={"metadata": {}, "litellm_metadata": {}},
proxy_server_request={},
)
assert len([t for t in tags if not t.startswith("User-Agent:")]) == 0
# Test case 5: None values for metadata/litellm_metadata
tags = StandardLoggingPayloadSetup._get_request_tags(
litellm_params={"metadata": None, "litellm_metadata": None},
proxy_server_request={},
)
assert isinstance(tags, list)
assert len([t for t in tags if not t.startswith("User-Agent:")]) == 0
# Test case 6: Empty litellm_params
tags = StandardLoggingPayloadSetup._get_request_tags(
litellm_params={},
proxy_server_request={},
)
assert isinstance(tags, list)
assert len([t for t in tags if not t.startswith("User-Agent:")]) == 0
# Test case 7: Metadata tags combined with user-agent tags
tags = StandardLoggingPayloadSetup._get_request_tags(
litellm_params={"metadata": {"tags": ["custom-tag"]}},
proxy_server_request={
"headers": {
"user-agent": "litellm/1.0.0",
}
},
)
assert "custom-tag" in tags
assert "User-Agent: litellm" in tags
assert "User-Agent: litellm/1.0.0" in tags
def test_get_extra_header_tags():
"""Test the _get_extra_header_tags method with various scenarios."""
import litellm
@@ -489,23 +573,23 @@ async def test_e2e_generate_cold_storage_object_key_successful():
patch("litellm.integrations.s3.get_s3_object_key") as mock_get_s3_key:
# Mock the S3 object key generation to return a predictable result
mock_get_s3_key.return_value = "2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
mock_get_s3_key.return_value = (
"2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
)
# Call the function
result = StandardLoggingPayloadSetup._generate_cold_storage_object_key(
start_time=start_time,
response_id=response_id,
team_alias=team_alias
start_time=start_time, response_id=response_id, team_alias=team_alias
)
# Verify the S3 function was called with correct parameters
mock_get_s3_key.assert_called_once_with(
s3_path="", # Empty path as default
team_alias_prefix="", # No team alias prefix for cold storage
start_time=start_time,
s3_file_name="time-10-30-45-123456_chatcmpl-test-12345"
s3_file_name="time-10-30-45-123456_chatcmpl-test-12345",
)
# Verify the result
assert result == "2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
assert result is not None
@@ -525,7 +609,7 @@ async def test_e2e_generate_cold_storage_object_key_with_custom_logger_s3_path()
# Create test data
start_time = datetime(2025, 1, 15, 10, 30, 45, 123456, timezone.utc)
response_id = "chatcmpl-test-12345"
# Create mock custom logger with s3_path
mock_custom_logger = MagicMock()
mock_custom_logger.s3_path = "storage"
@@ -536,27 +620,30 @@ async def test_e2e_generate_cold_storage_object_key_with_custom_logger_s3_path()
# Setup mocks
mock_get_logger.return_value = mock_custom_logger
mock_get_s3_key.return_value = "storage/2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
mock_get_s3_key.return_value = (
"storage/2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
)
# Call the function
result = StandardLoggingPayloadSetup._generate_cold_storage_object_key(
start_time=start_time,
response_id=response_id
start_time=start_time, response_id=response_id
)
# Verify logger was queried correctly
mock_get_logger.assert_called_once_with("s3_v2")
# Verify the S3 function was called with the custom logger's s3_path
mock_get_s3_key.assert_called_once_with(
s3_path="storage", # Should use custom logger's s3_path
team_alias_prefix="",
start_time=start_time,
s3_file_name="time-10-30-45-123456_chatcmpl-test-12345"
s3_file_name="time-10-30-45-123456_chatcmpl-test-12345",
)
# Verify the result
assert result == "storage/2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
assert (
result == "storage/2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
)
@pytest.mark.asyncio
@@ -572,7 +659,7 @@ async def test_e2e_generate_cold_storage_object_key_with_logger_no_s3_path():
# Create test data
start_time = datetime(2025, 1, 15, 10, 30, 45, 123456, timezone.utc)
response_id = "chatcmpl-test-12345"
# Create mock custom logger without s3_path
mock_custom_logger = MagicMock()
mock_custom_logger.s3_path = None # or could be missing attribute
@@ -583,22 +670,23 @@ async def test_e2e_generate_cold_storage_object_key_with_logger_no_s3_path():
# Setup mocks
mock_get_logger.return_value = mock_custom_logger
mock_get_s3_key.return_value = "2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
mock_get_s3_key.return_value = (
"2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
)
# Call the function
result = StandardLoggingPayloadSetup._generate_cold_storage_object_key(
start_time=start_time,
response_id=response_id
start_time=start_time, response_id=response_id
)
# Verify the S3 function was called with empty s3_path (fallback)
mock_get_s3_key.assert_called_once_with(
s3_path="", # Should fall back to empty string
team_alias_prefix="",
start_time=start_time,
s3_file_name="time-10-30-45-123456_chatcmpl-test-12345"
s3_file_name="time-10-30-45-123456_chatcmpl-test-12345",
)
# Verify the result
assert result == "2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
@@ -623,12 +711,8 @@ async def test_e2e_generate_cold_storage_object_key_not_configured():
with patch.object(litellm, 'cold_storage_custom_logger', None):
# Call the function
result = StandardLoggingPayloadSetup._generate_cold_storage_object_key(
start_time=start_time,
response_id=response_id,
team_alias=team_alias
start_time=start_time, response_id=response_id, team_alias=team_alias
)
# Verify the result is None when cold storage is not configured
assert result is None