Merge pull request #5091 from BerriAI/litellm_anthropic_streaming_tool_call_fix

fix(anthropic.py): handle anthropic returning empty argument string (invalid json str) for tool call while streaming

litellm/llms/anthropic.py

@@ -2,6 +2,7 @@ import copy
 import json
 import os
 import time
+import traceback
 import types
 from enum import Enum
 from functools import partial
@@ -36,6 +37,7 @@ from litellm.types.llms.anthropic import (
     AnthropicResponseUsageBlock,
     ContentBlockDelta,
     ContentBlockStart,
+    ContentBlockStop,
     ContentJsonBlockDelta,
     ContentTextBlockDelta,
     MessageBlockDelta,
@@ -920,7 +922,12 @@ class AnthropicChatCompletion(BaseLLM):
                     model=model, messages=messages, custom_llm_provider="anthropic"
                 )
             except Exception as e:
-                raise AnthropicError(status_code=400, message=str(e))
+                raise AnthropicError(
+                    status_code=400,
+                    message="{}\n{}\nReceived Messages={}".format(
+                        str(e), traceback.format_exc(), messages
+                    ),
+                )
 
         ## Load Config
         config = litellm.AnthropicConfig.get_config()
@@ -1079,10 +1086,30 @@ class ModelResponseIterator:
     def __init__(self, streaming_response, sync_stream: bool):
         self.streaming_response = streaming_response
         self.response_iterator = self.streaming_response
+        self.content_blocks: List[ContentBlockDelta] = []
+
+    def check_empty_tool_call_args(self) -> bool:
+        """
+        Check if the tool call block so far has been an empty string
+        """
+        args = ""
+        # if text content block -> skip
+        if len(self.content_blocks) == 0:
+            return False
+
+        if self.content_blocks[0]["delta"]["type"] == "text_delta":
+            return False
+
+        for block in self.content_blocks:
+            if block["delta"]["type"] == "input_json_delta":
+                args += block["delta"].get("partial_json", "")  # type: ignore
+
+        if len(args) == 0:
+            return True
+        return False
 
     def chunk_parser(self, chunk: dict) -> GenericStreamingChunk:
         try:
-            verbose_logger.debug(f"\n\nRaw chunk:\n{chunk}\n")
             type_chunk = chunk.get("type", "") or ""
 
             text = ""
@@ -1098,6 +1125,7 @@ class ModelResponseIterator:
                 chunk = {'type': 'content_block_delta', 'index': 0, 'delta': {'type': 'text_delta', 'text': 'Hello'}}
                 """
                 content_block = ContentBlockDelta(**chunk)  # type: ignore
+                self.content_blocks.append(content_block)
                 if "text" in content_block["delta"]:
                     text = content_block["delta"]["text"]
                 elif "partial_json" in content_block["delta"]:
@@ -1116,6 +1144,7 @@ class ModelResponseIterator:
                 data: {"type":"content_block_start","index":1,"content_block":{"type":"tool_use","id":"toolu_01T1x1fJ34qAmk2tNTrN7Up6","name":"get_weather","input":{}}}
                 """
                 content_block_start = ContentBlockStart(**chunk)  # type: ignore
+                self.content_blocks = []  # reset content blocks when new block starts
                 if content_block_start["content_block"]["type"] == "text":
                     text = content_block_start["content_block"]["text"]
                 elif content_block_start["content_block"]["type"] == "tool_use":
@@ -1128,6 +1157,20 @@ class ModelResponseIterator:
                         },
                         "index": content_block_start["index"],
                     }
+            elif type_chunk == "content_block_stop":
+                content_block_stop = ContentBlockStop(**chunk)  # type: ignore
+                # check if tool call content block
+                is_empty = self.check_empty_tool_call_args()
+                if is_empty:
+                    tool_use = {
+                        "id": None,
+                        "type": "function",
+                        "function": {
+                            "name": None,
+                            "arguments": "{}",
+                        },
+                        "index": content_block_stop["index"],
+                    }
             elif type_chunk == "message_delta":
                 """
                 Anthropic

litellm/llms/bedrock_httpx.py

@@ -27,6 +27,7 @@ import httpx  # type: ignore
 import requests  # type: ignore
 
 import litellm
+from litellm import verbose_logger
 from litellm.caching import DualCache
 from litellm.litellm_core_utils.core_helpers import map_finish_reason
 from litellm.litellm_core_utils.litellm_logging import Logging
@@ -1969,6 +1970,7 @@ class BedrockConverseLLM(BaseLLM):
         # Tool Config
         if bedrock_tool_config is not None:
             _data["toolConfig"] = bedrock_tool_config
+
         data = json.dumps(_data)
         ## COMPLETION CALL
 
@@ -2109,9 +2111,31 @@ class AWSEventStreamDecoder:
 
         self.model = model
         self.parser = EventStreamJSONParser()
+        self.content_blocks: List[ContentBlockDeltaEvent] = []
+
+    def check_empty_tool_call_args(self) -> bool:
+        """
+        Check if the tool call block so far has been an empty string
+        """
+        args = ""
+        # if text content block -> skip
+        if len(self.content_blocks) == 0:
+            return False
+
+        if "text" in self.content_blocks[0]:
+            return False
+
+        for block in self.content_blocks:
+            if "toolUse" in block:
+                args += block["toolUse"]["input"]
+
+        if len(args) == 0:
+            return True
+        return False
 
     def converse_chunk_parser(self, chunk_data: dict) -> GChunk:
         try:
+            verbose_logger.debug("\n\nRaw Chunk: {}\n\n".format(chunk_data))
             text = ""
             tool_use: Optional[ChatCompletionToolCallChunk] = None
             is_finished = False
@@ -2121,6 +2145,7 @@ class AWSEventStreamDecoder:
             index = int(chunk_data.get("contentBlockIndex", 0))
             if "start" in chunk_data:
                 start_obj = ContentBlockStartEvent(**chunk_data["start"])
+                self.content_blocks = []  # reset
                 if (
                     start_obj is not None
                     and "toolUse" in start_obj
@@ -2137,6 +2162,7 @@ class AWSEventStreamDecoder:
                     }
             elif "delta" in chunk_data:
                 delta_obj = ContentBlockDeltaEvent(**chunk_data["delta"])
+                self.content_blocks.append(delta_obj)
                 if "text" in delta_obj:
                     text = delta_obj["text"]
                 elif "toolUse" in delta_obj:
@@ -2149,6 +2175,20 @@ class AWSEventStreamDecoder:
                         },
                         "index": index,
                     }
+            elif (
+                "contentBlockIndex" in chunk_data
+            ):  # stop block, no 'start' or 'delta' object
+                is_empty = self.check_empty_tool_call_args()
+                if is_empty:
+                    tool_use = {
+                        "id": None,
+                        "type": "function",
+                        "function": {
+                            "name": None,
+                            "arguments": "{}",
+                        },
+                        "index": chunk_data["contentBlockIndex"],
+                    }
             elif "stopReason" in chunk_data:
                 finish_reason = map_finish_reason(chunk_data.get("stopReason", "stop"))
                 is_finished = True
@@ -2255,6 +2295,7 @@ class AWSEventStreamDecoder:
     def _parse_message_from_event(self, event) -> Optional[str]:
         response_dict = event.to_response_dict()
         parsed_response = self.parser.parse(response_dict, get_response_stream_shape())
+
         if response_dict["status_code"] != 200:
             raise ValueError(f"Bad response code, expected 200: {response_dict}")
         if "chunk" in parsed_response:

litellm/llms/prompt_templates/factory.py

@@ -2345,7 +2345,9 @@ def _bedrock_tools_pt(tools: List) -> List[BedrockToolBlock]:
     for tool in tools:
         parameters = tool.get("function", {}).get("parameters", None)
         name = tool.get("function", {}).get("name", "")
-        description = tool.get("function", {}).get("description", "")
+        description = tool.get("function", {}).get(
+            "description", name
+        )  # converse api requires a description
         tool_input_schema = BedrockToolInputSchemaBlock(json=parameters)
         tool_spec = BedrockToolSpecBlock(
             inputSchema=tool_input_schema, name=name, description=description

litellm/main.py

@@ -5114,7 +5114,9 @@ def stream_chunk_builder(
                 prev_index = curr_index
                 prev_id = curr_id
 
-        combined_arguments = "".join(argument_list)
+        combined_arguments = (
+            "".join(argument_list) or "{}"
+        )  # base case, return empty dict
         tool_calls_list.append(
             {
                 "id": id,

litellm/tests/test_completion.py

@@ -938,6 +938,7 @@ def test_completion_function_plus_image(model):
             }
         ]
 
+    try:
         response = completion(
             model=model,
             messages=[image_message],
@@ -949,8 +950,6 @@ def test_completion_function_plus_image(model):
         print(response)
     except litellm.InternalServerError:
         pass
-    except Exception as e:
-        pytest.fail(f"error occurred: {str(e)}")
 
 
 @pytest.mark.parametrize(

litellm/tests/test_streaming.py

@@ -2,6 +2,7 @@
 #    This tests streaming for the completion endpoint
 
 import asyncio
+import json
 import os
 import sys
 import time
@@ -2596,8 +2597,8 @@ def streaming_and_function_calling_format_tests(idx, chunk):
 @pytest.mark.parametrize(
     "model",
     [
-        "gpt-3.5-turbo",
-        "anthropic.claude-3-sonnet-20240229-v1:0",
+        # "gpt-3.5-turbo",
+        # "anthropic.claude-3-sonnet-20240229-v1:0",
         "claude-3-haiku-20240307",
     ],
 )
@@ -2627,7 +2628,7 @@ def test_streaming_and_function_calling(model):
 
     messages = [{"role": "user", "content": "What is the weather like in Boston?"}]
     try:
-        litellm.set_verbose = True
+        # litellm.set_verbose = True
         response: litellm.CustomStreamWrapper = completion(
             model=model,
             tools=tools,
@@ -2639,7 +2640,7 @@ def test_streaming_and_function_calling(model):
         json_str = ""
         for idx, chunk in enumerate(response):
             # continue
-            print("\n{}\n".format(chunk))
+            # print("\n{}\n".format(chunk))
             if idx == 0:
                 assert (
                     chunk.choices[0].delta.tool_calls[0].function.arguments is not None
@@ -3688,3 +3689,71 @@ def test_unit_test_custom_stream_wrapper_function_call():
     print("\n\n{}\n\n".format(new_model))
 
     assert len(new_model.choices[0].delta.tool_calls) > 0
+
+
+@pytest.mark.parametrize(
+    "model",
+    [
+        "gpt-3.5-turbo",
+        "claude-3-5-sonnet-20240620",
+        "anthropic.claude-3-sonnet-20240229-v1:0",
+        "vertex_ai/claude-3-5-sonnet@20240620",
+    ],
+)
+def test_streaming_tool_calls_valid_json_str(model):
+    if "vertex_ai" in model:
+        from litellm.tests.test_amazing_vertex_completion import (
+            load_vertex_ai_credentials,
+        )
+
+        load_vertex_ai_credentials()
+        vertex_location = "us-east5"
+    else:
+        vertex_location = None
+    litellm.set_verbose = False
+    messages = [
+        {"role": "user", "content": "Hit the snooze button."},
+    ]
+
+    tools = [
+        {
+            "type": "function",
+            "function": {
+                "name": "snooze",
+                "parameters": {
+                    "type": "object",
+                    "properties": {},
+                    "required": [],
+                },
+            },
+        }
+    ]
+
+    stream = litellm.completion(
+        model, messages, tools=tools, stream=True, vertex_location=vertex_location
+    )
+    chunks = [*stream]
+    print(f"chunks: {chunks}")
+    tool_call_id_arg_map = {}
+    curr_tool_call_id = None
+    curr_tool_call_str = ""
+    for chunk in chunks:
+        if chunk.choices[0].delta.tool_calls is not None:
+            if chunk.choices[0].delta.tool_calls[0].id is not None:
+                # flush prev tool call
+                if curr_tool_call_id is not None:
+                    tool_call_id_arg_map[curr_tool_call_id] = curr_tool_call_str
+                    curr_tool_call_str = ""
+                curr_tool_call_id = chunk.choices[0].delta.tool_calls[0].id
+                tool_call_id_arg_map[curr_tool_call_id] = ""
+            if chunk.choices[0].delta.tool_calls[0].function.arguments is not None:
+                curr_tool_call_str += (
+                    chunk.choices[0].delta.tool_calls[0].function.arguments
+                )
+    # flush prev tool call
+    if curr_tool_call_id is not None:
+        tool_call_id_arg_map[curr_tool_call_id] = curr_tool_call_str
+
+    for k, v in tool_call_id_arg_map.items():
+        print("k={}, v={}".format(k, v))
+        json.loads(v)  # valid json str

litellm/types/llms/anthropic.py

@@ -141,6 +141,11 @@ class ContentBlockDelta(TypedDict):
     delta: Union[ContentTextBlockDelta, ContentJsonBlockDelta]
 
 
+class ContentBlockStop(TypedDict):
+    type: Literal["content_block_stop"]
+    index: int
+
+
 class ToolUseBlock(TypedDict):
     """
     "content_block":{"type":"tool_use","id":"toolu_01T1x1fJ34qAmk2tNTrN7Up6","name":"get_weather","input":{}}