ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(bedrock_httpx.py): handle empty arguments returned during tool calling streaming

Browse Source
This commit is contained in:
Krrish Dholakia 2024-08-07 09:54:50 -07:00
parent 4919cc4d25
commit 2ccb5a48b7
4 changed files with 99 additions and 49 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

41
litellm/llms/bedrock_httpx.py
View File

@ -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:

4
litellm/llms/prompt_templates/factory.py
View File

@ -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

48
litellm/tests/test_completion.py
View File

@ -4346,51 +4346,3 @@ def test_moderation():
# test_moderation()
@pytest.mark.parametrize("model", ["gpt-3.5-turbo", "claude-3-5-sonnet-20240620"])
def test_streaming_tool_calls_valid_json_str(model):
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)
chunks = [*stream]
print(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

55
litellm/tests/test_streaming.py
View File

@ -2,6 +2,7 @@
# This tests streaming for the completion endpoint
import asyncio
import json
import os
import sys
import time
@ -3688,3 +3689,57 @@ 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",
],
)
def test_streaming_tool_calls_valid_json_str(model):
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)
chunks = [*stream]
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