mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
fix(conversere_transformation.py): fix test
This commit is contained in:
Krrish Dholakia
parent
447f1ea6bc
commit
c27f57fdab
1 changed files with 64 additions and 29 deletions
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue