From 6dc9b390956e0fd9d0c3d3d980bb2534b16f609a Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Fri, 9 Aug 2024 13:26:21 -0700
Subject: [PATCH] test invalid tool namehandling

---
 litellm/tests/test_bedrock_completion.py | 125 ++++++++++++++++++++++-
 1 file changed, 124 insertions(+), 1 deletion(-)

diff --git a/litellm/tests/test_bedrock_completion.py b/litellm/tests/test_bedrock_completion.py
index 4b1e544c7f1..fa1894e74ce 100644
--- a/litellm/tests/test_bedrock_completion.py
+++ b/litellm/tests/test_bedrock_completion.py
@@ -25,8 +25,9 @@ from litellm import (
     completion_cost,
     embedding,
 )
-from litellm.llms.bedrock_httpx import BedrockLLM
+from litellm.llms.bedrock_httpx import BedrockLLM, ToolBlock
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
+from litellm.llms.prompt_templates.factory import _bedrock_tools_pt
 
 # litellm.num_retries = 3
 litellm.cache = None
@@ -983,3 +984,125 @@ def test_completion_bedrock_external_client_region():
         pass
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")
+
+
+def test_bedrock_tool_calling():
+    """
+    # related issue: https://github.com/BerriAI/litellm/issues/5007
+    # Bedrock tool names must satisfy regular expression pattern: [a-zA-Z][a-zA-Z0-9_]* ensure this is true
+    """
+    litellm.set_verbose = True
+    response = litellm.completion(
+        model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0",
+        fallbacks=["bedrock/meta.llama3-1-8b-instruct-v1:0"],
+        messages=[
+            {
+                "role": "user",
+                "content": "What's the weather like in Boston today in Fahrenheit?",
+            }
+        ],
+        tools=[
+            {
+                "type": "function",
+                "function": {
+                    "name": "-DoSomethingVeryCool-forLitellm_Testin999229291-0293993",
+                    "description": "do something very cool",
+                    "parameters": {"type": "object", "properties": {}},
+                },
+            }
+        ],
+    )
+
+    print("bedrock response")
+    print(response)
+
+
+def test_bedrock_tools_pt_valid_names():
+    """
+    # related issue: https://github.com/BerriAI/litellm/issues/5007
+    # Bedrock tool names must satisfy regular expression pattern: [a-zA-Z][a-zA-Z0-9_]* ensure this is true
+
+    """
+    tools = [
+        {
+            "type": "function",
+            "function": {
+                "name": "get_current_weather",
+                "description": "Get the current weather",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "location": {"type": "string"},
+                    },
+                    "required": ["location"],
+                },
+            },
+        },
+        {
+            "type": "function",
+            "function": {
+                "name": "search_restaurants",
+                "description": "Search for restaurants",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "cuisine": {"type": "string"},
+                    },
+                    "required": ["cuisine"],
+                },
+            },
+        },
+    ]
+
+    result = _bedrock_tools_pt(tools)
+
+    assert len(result) == 2
+    assert result[0]["toolSpec"]["name"] == "get_current_weather"
+    assert result[1]["toolSpec"]["name"] == "search_restaurants"
+
+
+def test_bedrock_tools_pt_invalid_names():
+    """
+    # related issue: https://github.com/BerriAI/litellm/issues/5007
+    # Bedrock tool names must satisfy regular expression pattern: [a-zA-Z][a-zA-Z0-9_]* ensure this is true
+
+    """
+
+    tools = [
+        {
+            "type": "function",
+            "function": {
+                "name": "123-invalid@name",
+                "description": "Invalid name test",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "test": {"type": "string"},
+                    },
+                    "required": ["test"],
+                },
+            },
+        },
+        {
+            "type": "function",
+            "function": {
+                "name": "another@invalid#name",
+                "description": "Another invalid name test",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "test": {"type": "string"},
+                    },
+                    "required": ["test"],
+                },
+            },
+        },
+    ]
+
+    result = _bedrock_tools_pt(tools)
+
+    print("bedrock tools after prompt formatting=", result)
+
+    assert len(result) == 2
+    assert result[0]["toolSpec"]["name"] == "a123_invalid_name"
+    assert result[1]["toolSpec"]["name"] == "another_invalid_name"