From c27f57fdabc8ac80a4e321042a2c31de3c9f2e3a Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 6 Sep 2025 11:06:15 -0700 Subject: [PATCH] fix(conversere_transformation.py): fix test --- .../bedrock/chat/converse_transformation.py | 93 +++++++++++++------ 1 file changed, 64 insertions(+), 29 deletions(-) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 88b65132138..e06f92a8380 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -10,6 +10,7 @@ from typing import List, Literal, Optional, Tuple, Union, cast, overload import httpx import litellm +from litellm._logging import verbose_logger from litellm.constants import RESPONSE_FORMAT_TOOL_NAME from litellm.litellm_core_utils.core_helpers import map_finish_reason from litellm.litellm_core_utils.litellm_logging import Logging @@ -49,14 +50,19 @@ from litellm.types.utils import ( ) from litellm.utils import add_dummy_tool, has_tool_call_blocks, supports_reasoning -from ..common_utils import BedrockError, BedrockModelInfo, get_bedrock_tool_name, get_anthropic_beta_from_headers +from ..common_utils import ( + BedrockError, + BedrockModelInfo, + get_anthropic_beta_from_headers, + get_bedrock_tool_name, +) # Computer use tool prefixes supported by Bedrock BEDROCK_COMPUTER_USE_TOOLS = [ "computer_use_preview", "computer_", "bash_", - "text_editor_" + "text_editor_", ] @@ -236,7 +242,7 @@ class AmazonConverseConfig(BaseConfig): """Check if computer use tools are being used in the request.""" if tools is None: return False - + for tool in tools: if "type" in tool: tool_type = tool["type"] @@ -250,17 +256,17 @@ class AmazonConverseConfig(BaseConfig): ) -> List[dict]: """Transform computer use tools to Bedrock format.""" transformed_tools: List[dict] = [] - + for tool in computer_use_tools: tool_type = tool.get("type", "") - + # Check if this is a computer use tool with the startswith method is_computer_use_tool = False for computer_use_prefix in BEDROCK_COMPUTER_USE_TOOLS: if tool_type.startswith(computer_use_prefix): is_computer_use_tool = True break - + transformed_tool: dict = {} if is_computer_use_tool: if tool_type.startswith("computer_") and "function" in tool: @@ -269,7 +275,7 @@ class AmazonConverseConfig(BaseConfig): transformed_tool = { "type": tool_type, "name": func.get("name", "computer"), - **func.get("parameters", {}) + **func.get("parameters", {}), } else: # Direct tools - just need to ensure name is present @@ -282,27 +288,29 @@ class AmazonConverseConfig(BaseConfig): else: # Pass through other tools as-is transformed_tool = dict(tool) - + transformed_tools.append(transformed_tool) - + return transformed_tools def _separate_computer_use_tools( self, tools: List[OpenAIChatCompletionToolParam], model: str - ) -> Tuple[List[OpenAIChatCompletionToolParam], List[OpenAIChatCompletionToolParam]]: + ) -> Tuple[ + List[OpenAIChatCompletionToolParam], List[OpenAIChatCompletionToolParam] + ]: """ Separate computer use tools from regular function tools. - + Args: tools: List of tools to separate model: The model name to check if it supports computer use - + Returns: Tuple of (computer_use_tools, regular_tools) """ computer_use_tools = [] regular_tools = [] - + for tool in tools: if "type" in tool: tool_type = tool["type"] @@ -317,9 +325,8 @@ class AmazonConverseConfig(BaseConfig): regular_tools.append(tool) else: regular_tools.append(tool) - - return computer_use_tools, regular_tools + return computer_use_tools, regular_tools def _create_json_tool_call_for_response_format( self, @@ -345,6 +352,8 @@ class AmazonConverseConfig(BaseConfig): "properties": {}, } else: + # Use the schema as-is for Bedrock + # Bedrock requires the tool schema to be of type "object" and doesn't need unwrapping _input_schema = json_schema tool_param_function_chunk = ChatCompletionToolParamFunctionChunk( @@ -426,9 +435,7 @@ class AmazonConverseConfig(BaseConfig): ): optional_params["tool_choice"] = ToolChoiceValuesBlock( - tool=SpecificToolChoiceBlock( - name=RESPONSE_FORMAT_TOOL_NAME - ) + tool=SpecificToolChoiceBlock(name=RESPONSE_FORMAT_TOOL_NAME) ) optional_params["json_mode"] = True if non_default_params.get("stream", False) is True: @@ -602,7 +609,6 @@ class AmazonConverseConfig(BaseConfig): return {} - def _transform_request_helper( self, model: str, @@ -658,36 +664,38 @@ class AmazonConverseConfig(BaseConfig): ) original_tools = inference_params.pop("tools", []) - + # Initialize bedrock_tools bedrock_tools: List[ToolBlock] = [] - + # Collect anthropic_beta values from user headers anthropic_beta_list = [] if headers: user_betas = get_anthropic_beta_from_headers(headers) anthropic_beta_list.extend(user_betas) - + # Only separate tools if computer use tools are actually present if original_tools and self.is_computer_use_tool_used(original_tools, model): # Separate computer use tools from regular function tools computer_use_tools, regular_tools = self._separate_computer_use_tools( original_tools, model ) - + # Process regular function tools using existing logic bedrock_tools = _bedrock_tools_pt(regular_tools) - + # Add computer use tools and anthropic_beta if needed (only when computer use tools are present) if computer_use_tools: anthropic_beta_list.append("computer-use-2024-10-22") # Transform computer use tools to proper Bedrock format - transformed_computer_tools = self._transform_computer_use_tools(computer_use_tools) + transformed_computer_tools = self._transform_computer_use_tools( + computer_use_tools + ) additional_request_params["tools"] = transformed_computer_tools else: # No computer use tools, process all tools as regular tools bedrock_tools = _bedrock_tools_pt(original_tools) - + # Set anthropic_beta in additional_request_params if we have any beta features if anthropic_beta_list: # Remove duplicates while preserving order @@ -698,7 +706,7 @@ class AmazonConverseConfig(BaseConfig): unique_betas.append(beta) seen.add(beta) additional_request_params["anthropic_beta"] = unique_betas - + bedrock_tool_config: Optional[ToolConfigBlock] = None if len(bedrock_tools) > 0: tool_choice_values: ToolChoiceValuesBlock = inference_params.pop( @@ -1124,9 +1132,37 @@ class AmazonConverseConfig(BaseConfig): self._transform_thinking_blocks(reasoningContentBlocks) ) chat_completion_message["content"] = content_str - if json_mode is True and tools is not None and len(tools) == 1 and tools[0]["function"]["name"] == RESPONSE_FORMAT_TOOL_NAME: + if ( + json_mode is True + and tools is not None + and len(tools) == 1 + and tools[0]["function"].get("name") == RESPONSE_FORMAT_TOOL_NAME + ): + verbose_logger.debug( + "Processing JSON tool call response for response_format" + ) json_mode_content_str: Optional[str] = tools[0]["function"].get("arguments") if json_mode_content_str is not None: + import json + + # Bedrock returns the response wrapped in a "properties" object + # We need to extract the actual content from this wrapper + try: + + response_data = json.loads(json_mode_content_str) + + # If Bedrock wrapped the response in "properties", extract the content + if ( + isinstance(response_data, dict) + and "properties" in response_data + and len(response_data) == 1 + ): + response_data = response_data["properties"] + json_mode_content_str = json.dumps(response_data) + except json.JSONDecodeError: + # If parsing fails, use the original response + pass + chat_completion_message["content"] = json_mode_content_str else: chat_completion_message["tool_calls"] = tools @@ -1186,7 +1222,6 @@ class AmazonConverseConfig(BaseConfig): if api_key: headers["Authorization"] = f"Bearer {api_key}" return headers - def should_fake_stream( self,