diff --git a/litellm/integrations/rubrik.py b/litellm/integrations/rubrik.py index e901ec5fbb7..299e95eb813 100644 --- a/litellm/integrations/rubrik.py +++ b/litellm/integrations/rubrik.py @@ -320,12 +320,8 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): "content": system_prompt_msg_list, } if isinstance(standard_logging_payload["messages"], list): - standard_logging_payload["messages"].insert( - 0, system_scaffold - ) - elif isinstance( - standard_logging_payload["messages"], (dict, str) - ): + standard_logging_payload["messages"].insert(0, system_scaffold) + elif isinstance(standard_logging_payload["messages"], (dict, str)): standard_logging_payload["messages"] = [ system_scaffold, standard_logging_payload["messages"], @@ -355,14 +351,10 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): exc_info=True, ) - async def async_log_success_event( - self, kwargs, response_obj, start_time, end_time - ): + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): await self._enqueue_log_event(kwargs, "success") - async def async_log_failure_event( - self, kwargs, response_obj, start_time, end_time - ): + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): await self._enqueue_log_event(kwargs, "failure") # -- Batch logging --------------------------------------------------------- diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index d6c5aacc35d..9e8467fe394 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -170,7 +170,9 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): anthropic_request.pop("model", None) anthropic_request.pop("stream", None) anthropic_request.pop("output_format", None) - if not _supports_factory(model=model, custom_llm_provider=None, key="supports_output_config"): + if not _supports_factory( + model=model, custom_llm_provider=None, key="supports_output_config" + ): anthropic_request.pop("output_config", None) if "anthropic_version" not in anthropic_request: anthropic_request["anthropic_version"] = self.anthropic_version diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index abb1920cc08..f5b0659068c 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -504,7 +504,9 @@ class AmazonAnthropicClaudeMessagesConfig( # 5b. Bedrock Invoke supports output_config (effort) for Claude 4.6+ models, # but older models do not — strip it to avoid request rejection. # Ref: https://github.com/BerriAI/litellm/issues/22797 - if not _supports_factory(model=model, custom_llm_provider=None, key="supports_output_config"): + if not _supports_factory( + model=model, custom_llm_provider=None, key="supports_output_config" + ): anthropic_messages_request.pop("output_config", None) # 5a. Remove `custom` field from tools (Bedrock doesn't support it) diff --git a/tests/test_litellm/integrations/rubrik_test_helpers.py b/tests/test_litellm/integrations/rubrik_test_helpers.py index 1bdb8cb247b..f33d4413ddd 100644 --- a/tests/test_litellm/integrations/rubrik_test_helpers.py +++ b/tests/test_litellm/integrations/rubrik_test_helpers.py @@ -5,9 +5,7 @@ from typing import Any, Dict from litellm.types.utils import GenericGuardrailAPIInputs -def make_tool_call_dict( - tc_id: str, name: str, arguments: str = "{}" -) -> Dict[str, Any]: +def make_tool_call_dict(tc_id: str, name: str, arguments: str = "{}") -> Dict[str, Any]: """Create a tool call dict matching the ChatCompletionMessageToolCall schema.""" return { "id": tc_id, diff --git a/tests/test_litellm/integrations/test_rubrik.py b/tests/test_litellm/integrations/test_rubrik.py index d543516f183..162a9e44335 100644 --- a/tests/test_litellm/integrations/test_rubrik.py +++ b/tests/test_litellm/integrations/test_rubrik.py @@ -58,9 +58,7 @@ class TestInitialization: def test_init_with_constructor_params(self): with patch("asyncio.create_task", Mock()): - handler = RubrikLogger( - api_key="ctor-key", api_base="http://ctor-host:9090" - ) + handler = RubrikLogger(api_key="ctor-key", api_base="http://ctor-host:9090") assert handler.key == "ctor-key" assert ( handler.tool_blocking_endpoint @@ -152,9 +150,7 @@ class TestInitialization: assert handler._headers["Content-Type"] == "application/json" def test_headers_without_api_key(self): - with patch.dict( - os.environ, {"RUBRIK_WEBHOOK_URL": "http://host"}, clear=True - ): + with patch.dict(os.environ, {"RUBRIK_WEBHOOK_URL": "http://host"}, clear=True): with patch("asyncio.create_task", Mock()): h = RubrikLogger() assert "Authorization" not in h._headers @@ -351,9 +347,7 @@ def _echo_service(): @pytest.mark.asyncio class TestApplyGuardrail: async def test_skips_requests(self, handler): - inputs = make_inputs_with_tools( - [make_tool_call_dict("call_1", "test_tool")] - ) + inputs = make_inputs_with_tools([make_tool_call_dict("call_1", "test_tool")]) result = await handler.apply_guardrail( inputs=inputs, request_data={}, input_type="request" ) @@ -464,9 +458,7 @@ class TestApplyGuardrail: async def test_blocking_service_payload_format(self, handler): tc1 = make_tool_call_dict("call_1", "get_weather", '{"location": "SF"}') - tc2 = make_tool_call_dict( - "call_2", "send_email", '{"to": "user@example.com"}' - ) + tc2 = make_tool_call_dict("call_2", "send_email", '{"to": "user@example.com"}') inputs = make_inputs_with_tools([tc1, tc2]) captured_payload: Dict[str, Any] = {} @@ -605,9 +597,7 @@ class TestApplyGuardrailAnthropicFormat: inputs = make_inputs_with_tools([tc]) mock_client = AsyncMock() - mock_client.post = AsyncMock( - side_effect=httpx.TimeoutException("Timeout") - ) + mock_client.post = AsyncMock(side_effect=httpx.TimeoutException("Timeout")) handler.tool_blocking_client = mock_client result = await handler.apply_guardrail( @@ -728,5 +718,7 @@ class TestResolveModel: response = Mock() response.model = "" - result = RubrikLogger._resolve_model({"response": response}, {"model": "fallback"}) + result = RubrikLogger._resolve_model( + {"response": response}, {"model": "fallback"} + ) assert result == "unknown" diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index a8042c22178..f1e7de23315 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -702,7 +702,10 @@ async def test_promote_message_start_cache_when_message_stop_omits_cache_fields( "delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"input_tokens": 10, "output_tokens": 181}, } - yield {"type": "message_stop", "usage": {"input_tokens": 10, "output_tokens": 181}} + yield { + "type": "message_stop", + "usage": {"input_tokens": 10, "output_tokens": 181}, + } merged: list[dict] = [] async for chunk in cfg._promote_message_stop_usage(_stream()): @@ -756,7 +759,11 @@ async def test_unified_bedrock_messages_cache_on_start_only_never_negative_cost( }, }, } - yield {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}} + yield { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + } yield { "type": "content_block_delta", "index": 0, @@ -768,7 +775,10 @@ async def test_unified_bedrock_messages_cache_on_start_only_never_negative_cost( "delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 181, "input_tokens": 10}, } - yield {"type": "message_stop", "usage": {"input_tokens": 10, "output_tokens": 181}} + yield { + "type": "message_stop", + "usage": {"input_tokens": 10, "output_tokens": 181}, + } logging_obj = LiteLLMLoggingObj( model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",