Merge pull request #5144 from BerriAI/litellm_bedrock_put_modified_tool_Call_name_in_output

[Fix] Place bedrock modified tool call name in output

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,

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,

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
 

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():
     """