diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 8332463c5c..46fee677b9 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -34,7 +34,14 @@ from litellm.types.llms.openai import ( OpenAIChatCompletionToolParam, OpenAIMessageContentListBlock, ) -from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage +from litellm.types.utils import ( + ChatCompletionMessageToolCall, + Function, + Message, + ModelResponse, + PromptTokensDetailsWrapper, + Usage, +) from litellm.utils import add_dummy_tool, has_tool_call_blocks from ..common_utils import BedrockError, BedrockModelInfo, get_bedrock_tool_name @@ -690,6 +697,65 @@ class AmazonConverseConfig(BaseConfig): ) return openai_usage + def get_tool_call_names( + self, + tools: Optional[ + Union[List[ToolBlock], List[OpenAIChatCompletionToolParam]] + ] = None, + ) -> List[str]: + if tools is None: + return [] + tool_set: set[str] = set() + for tool in tools: + tool_spec = tool.get("toolSpec") + function = tool.get("function") + if tool_spec is not None: + _name = cast(dict, tool_spec).get("name") + if _name is not None and isinstance(_name, str): + tool_set.add(_name) + if function is not None: + _name = cast(dict, function).get("name") + if _name is not None and isinstance(_name, str): + tool_set.add(_name) + return list(tool_set) + + def apply_tool_call_transformation_if_needed( + self, + message: Message, + tools: Optional[List[ToolBlock]] = None, + initial_finish_reason: Optional[str] = None, + ) -> Tuple[Message, Optional[str]]: + """ + Apply tool call transformation to a message. + + LLM providers (e.g. Bedrock, Vertex AI) sometimes return tool call in the response content. + + If the response content is a JSON object, we can parse it and return the tool call in the tool_calls field. + """ + returned_finish_reason = initial_finish_reason + if tools is None: + return message, returned_finish_reason + + if message.content is not None: + try: + tool_call_names = self.get_tool_call_names(tools) + json_content = json.loads(message.content) + if ( + json_content.get("type") == "function" + and json_content.get("name") in tool_call_names + ): + tool_calls = [ + ChatCompletionMessageToolCall(function=Function(**json_content)) + ] + + message.tool_calls = tool_calls + message.content = None + returned_finish_reason = "tool_calls" + except Exception: + pass + + return message, returned_finish_reason + def _transform_response( self, model: str, @@ -819,11 +885,23 @@ class AmazonConverseConfig(BaseConfig): ## CALCULATING USAGE - bedrock returns usage in the headers usage = self._transform_usage(completion_response["usage"]) + ## HANDLE TOOL CALLS + _message = Message(**chat_completion_message) + initial_finish_reason = map_finish_reason(completion_response["stopReason"]) + + ( + returned_message, + returned_finish_reason, + ) = self.apply_tool_call_transformation_if_needed( + message=_message, + tools=optional_params.get("tools"), + initial_finish_reason=initial_finish_reason, + ) model_response.choices = [ litellm.Choices( - finish_reason=map_finish_reason(completion_response["stopReason"]), + finish_reason=returned_finish_reason, index=0, - message=litellm.Message(**chat_completion_message), + message=returned_message, ) ] model_response.created = int(time.time()) diff --git a/litellm/utils.py b/litellm/utils.py index 2403ded78c..7d9da60b5d 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -524,9 +524,9 @@ def function_setup( # noqa: PLR0915 function_id: Optional[str] = kwargs["id"] if "id" in kwargs else None ## DYNAMIC CALLBACKS ## - dynamic_callbacks: Optional[List[Union[str, Callable, CustomLogger]]] = ( - kwargs.pop("callbacks", None) - ) + dynamic_callbacks: Optional[ + List[Union[str, Callable, CustomLogger]] + ] = kwargs.pop("callbacks", None) all_callbacks = get_dynamic_callbacks(dynamic_callbacks=dynamic_callbacks) if len(all_callbacks) > 0: @@ -1210,9 +1210,9 @@ def client(original_function): # noqa: PLR0915 exception=e, retry_policy=kwargs.get("retry_policy"), ) - kwargs["retry_policy"] = ( - reset_retry_policy() - ) # prevent infinite loops + kwargs[ + "retry_policy" + ] = reset_retry_policy() # prevent infinite loops litellm.num_retries = ( None # set retries to None to prevent infinite loops ) @@ -2755,16 +2755,16 @@ def get_optional_params( # noqa: PLR0915 True # so that main.py adds the function call to the prompt ) if "tools" in non_default_params: - optional_params["functions_unsupported_model"] = ( - non_default_params.pop("tools") - ) + optional_params[ + "functions_unsupported_model" + ] = non_default_params.pop("tools") non_default_params.pop( "tool_choice", None ) # causes ollama requests to hang elif "functions" in non_default_params: - optional_params["functions_unsupported_model"] = ( - non_default_params.pop("functions") - ) + optional_params[ + "functions_unsupported_model" + ] = non_default_params.pop("functions") elif ( litellm.add_function_to_prompt ): # if user opts to add it to prompt instead @@ -2787,10 +2787,10 @@ def get_optional_params( # noqa: PLR0915 if "response_format" in non_default_params: if provider_config is not None: - non_default_params["response_format"] = ( - provider_config.get_json_schema_from_pydantic_object( - response_format=non_default_params["response_format"] - ) + non_default_params[ + "response_format" + ] = provider_config.get_json_schema_from_pydantic_object( + response_format=non_default_params["response_format"] ) else: non_default_params["response_format"] = type_to_response_format_param( @@ -3806,9 +3806,9 @@ def _count_characters(text: str) -> int: def get_response_string(response_obj: Union[ModelResponse, ModelResponseStream]) -> str: - _choices: Union[List[Union[Choices, StreamingChoices]], List[StreamingChoices]] = ( - response_obj.choices - ) + _choices: Union[ + List[Union[Choices, StreamingChoices]], List[StreamingChoices] + ] = response_obj.choices response_str = "" for choice in _choices: diff --git a/tests/litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/litellm/llms/bedrock/chat/test_converse_transformation.py index 063bde82b3..ef0b296fe0 100644 --- a/tests/litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/litellm/llms/bedrock/chat/test_converse_transformation.py @@ -41,6 +41,7 @@ def test_transform_usage(): assert openai_usage._cache_creation_input_tokens == usage["cacheWriteInputTokens"] assert openai_usage._cache_read_input_tokens == usage["cacheReadInputTokens"] + def test_transform_system_message(): config = AmazonConverseConfig() @@ -96,7 +97,11 @@ def test_transform_system_message(): { "role": "system", "content": [ - {"type": "text", "text": "Cache this!", "cache_control": {"type": "ephemeral"}}, + { + "type": "text", + "text": "Cache this!", + "cache_control": {"type": "ephemeral"}, + }, {"type": "text", "text": "Don't cache this!"}, ], }, @@ -121,6 +126,7 @@ def test_transform_system_message(): assert out_messages[1]["role"] == "assistant" assert system_blocks == [] + def test_transform_thinking_blocks_with_redacted_content(): thinking_blocks = [ { @@ -139,3 +145,34 @@ def test_transform_thinking_blocks_with_redacted_content(): assert transformed_thinking_blocks[0]["type"] == "thinking" assert transformed_thinking_blocks[1]["type"] == "redacted_thinking" + +def test_apply_tool_call_transformation_if_needed(): + from litellm.types.utils import Message + + config = AmazonConverseConfig() + tool_calls = [ + { + "type": "function", + "function": { + "name": "test_function", + "arguments": "test_arguments", + }, + }, + ] + tool_response = { + "type": "function", + "name": "test_function", + "parameters": {"test": "test"}, + } + message = Message( + role="user", + content=json.dumps(tool_response), + ) + transformed_message, _ = config.apply_tool_call_transformation_if_needed( + message, tool_calls + ) + assert len(transformed_message.tool_calls) == 1 + assert transformed_message.tool_calls[0].function.name == "test_function" + assert transformed_message.tool_calls[0].function.arguments == json.dumps( + tool_response["parameters"] + )