diff --git a/enterprise/litellm_enterprise/integrations/custom_guardrail.py b/enterprise/litellm_enterprise/integrations/custom_guardrail.py index db7e557ac5..b165d788f3 100644 --- a/enterprise/litellm_enterprise/integrations/custom_guardrail.py +++ b/enterprise/litellm_enterprise/integrations/custom_guardrail.py @@ -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, ) diff --git a/enterprise/litellm_enterprise/integrations/prometheus.py b/enterprise/litellm_enterprise/integrations/prometheus.py index 4e81206122..49037678df 100644 --- a/enterprise/litellm_enterprise/integrations/prometheus.py +++ b/enterprise/litellm_enterprise/integrations/prometheus.py @@ -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) diff --git a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py index b411a21757..236935778d 100644 --- a/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py +++ b/litellm/integrations/vector_store_integrations/vector_store_pre_call_hook.py @@ -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 diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 7d72755478..f495852725 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -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, diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 0ee11fdb3f..2880e31e29 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -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" \ No newline at end of file + api_version: "2025-09-01" diff --git a/litellm/responses/main.py b/litellm/responses/main.py index c190d8e2d2..c297a87077 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -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: diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index aae79532ed..5c2f84cd87 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -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 - -