diff --git a/litellm/__init__.py b/litellm/__init__.py index 8d0e1b31495..7e23e92a4ac 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -866,6 +866,7 @@ from .llms.bedrock_httpx import ( AmazonCohereChatConfig, AmazonConverseConfig, BEDROCK_CONVERSE_MODELS, + bedrock_tool_name_mappings, ) from .llms.bedrock import ( AmazonTitanConfig, diff --git a/litellm/llms/bedrock_httpx.py b/litellm/llms/bedrock_httpx.py index 49f080bd06b..a92288b6ffa 100644 --- a/litellm/llms/bedrock_httpx.py +++ b/litellm/llms/bedrock_httpx.py @@ -28,7 +28,7 @@ import requests # type: ignore import litellm from litellm import verbose_logger -from litellm.caching import DualCache +from litellm.caching import DualCache, InMemoryCache from litellm.litellm_core_utils.core_helpers import map_finish_reason from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.http_handler import ( @@ -89,6 +89,9 @@ BEDROCK_CONVERSE_MODELS = [ iam_cache = DualCache() _response_stream_shape_cache = None +bedrock_tool_name_mappings: InMemoryCache = InMemoryCache( + max_size_in_memory=50, default_ttl=600 +) class AmazonCohereChatConfig: @@ -1495,8 +1498,14 @@ class BedrockConverseLLM(BaseLLM): if "text" in content: content_str += content["text"] if "toolUse" in content: + + ## check tool name was formatted by litellm + _response_tool_name = content["toolUse"]["name"] + response_tool_name = get_bedrock_tool_name( + response_tool_name=_response_tool_name + ) _function_chunk = ChatCompletionToolCallFunctionChunk( - name=content["toolUse"]["name"], + name=response_tool_name, arguments=json.dumps(content["toolUse"]["input"]), ) _tool_response_chunk = ChatCompletionToolCallChunk( @@ -2105,6 +2114,24 @@ def get_response_stream_shape(): return _response_stream_shape_cache +def get_bedrock_tool_name(response_tool_name: str) -> str: + """ + If litellm formatted the input tool name, we need to convert it back to the original name. + + Args: + response_tool_name (str): The name of the tool as received from the response. + + Returns: + str: The original name of the tool. + """ + + if response_tool_name in litellm.bedrock_tool_name_mappings.cache_dict: + response_tool_name = litellm.bedrock_tool_name_mappings.cache_dict[ + response_tool_name + ] + return response_tool_name + + class AWSEventStreamDecoder: def __init__(self, model: str) -> None: from botocore.parsers import EventStreamJSONParser @@ -2151,11 +2178,16 @@ class AWSEventStreamDecoder: and "toolUse" in start_obj and start_obj["toolUse"] is not None ): + ## check tool name was formatted by litellm + _response_tool_name = start_obj["toolUse"]["name"] + response_tool_name = get_bedrock_tool_name( + response_tool_name=_response_tool_name + ) tool_use = { "id": start_obj["toolUse"]["toolUseId"], "type": "function", "function": { - "name": start_obj["toolUse"]["name"], + "name": response_tool_name, "arguments": "", }, "index": index, diff --git a/litellm/llms/prompt_templates/factory.py b/litellm/llms/prompt_templates/factory.py index db655c01d0d..6c055300606 100644 --- a/litellm/llms/prompt_templates/factory.py +++ b/litellm/llms/prompt_templates/factory.py @@ -1,3 +1,4 @@ +import copy import json import re import traceback @@ -2310,13 +2311,20 @@ def make_valid_bedrock_tool_name(input_tool_name: str) -> str: # If the string is empty, return a default valid identifier if input_tool_name is None or len(input_tool_name) == 0: return input_tool_name - + bedrock_tool_name = copy.copy(input_tool_name) # If it doesn't start with a letter, prepend 'a' - if not input_tool_name[0].isalpha(): - input_tool_name = "a" + input_tool_name + if not bedrock_tool_name[0].isalpha(): + bedrock_tool_name = "a" + bedrock_tool_name # Replace any invalid characters with underscores - valid_string = "".join(replace_invalid(char) for char in input_tool_name) + valid_string = "".join(replace_invalid(char) for char in bedrock_tool_name) + + if input_tool_name != valid_string: + # passed tool name was formatted to become valid + # store it internally so we can use for the response + litellm.bedrock_tool_name_mappings.set_cache( + key=valid_string, value=input_tool_name + ) return valid_string diff --git a/litellm/tests/test_bedrock_completion.py b/litellm/tests/test_bedrock_completion.py index fa1894e74ce..f97fbf0303d 100644 --- a/litellm/tests/test_bedrock_completion.py +++ b/litellm/tests/test_bedrock_completion.py @@ -1006,7 +1006,7 @@ def test_bedrock_tool_calling(): "type": "function", "function": { "name": "-DoSomethingVeryCool-forLitellm_Testin999229291-0293993", - "description": "do something very cool", + "description": "use this to get the current weather", "parameters": {"type": "object", "properties": {}}, }, } @@ -1016,6 +1016,18 @@ def test_bedrock_tool_calling(): print("bedrock response") print(response) + # Assert that the tools in response have the same function name as the input + _choice_1 = response.choices[0] + if _choice_1.message.tool_calls is not None: + print(_choice_1.message.tool_calls) + for tool_call in _choice_1.message.tool_calls: + _tool_Call_name = tool_call.function.name + if _tool_Call_name is not None and "DoSomethingVeryCool" in _tool_Call_name: + assert ( + _tool_Call_name + == "-DoSomethingVeryCool-forLitellm_Testin999229291-0293993" + ) + def test_bedrock_tools_pt_valid_names(): """