tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-02 18:21:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(converse_transformation.py): handle meta llama tool call response (#10541)

Browse Source 
* fix(converse_transformation.py): handle meta llama tool call response

Fixes issue where bedrock meta llama would return tool call response as content str

* test(test_converse_transformation.py): add unit testing for new function

* fix: fix linting error

* fix: fix linting error
This commit is contained in:
Krish Dholakia
2025-05-03 21:31:14 -07:00
committed by GitHub
parent 7c4fae477a
commit 522e334fbd
3 changed files with 138 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/llms/bedrock/chat/converse_transformation.py
+81 -3
View File
@@ -34,7 +34,14 @@ from litellm.types.llms.openai import (
OpenAIChatCompletionToolParam,
OpenAIMessageContentListBlock,
)
from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage
from litellm.types.utils import (
ChatCompletionMessageToolCall,
Function,
Message,
ModelResponse,
PromptTokensDetailsWrapper,
Usage,
)
from litellm.utils import add_dummy_tool, has_tool_call_blocks
from ..common_utils import BedrockError, BedrockModelInfo, get_bedrock_tool_name
@@ -690,6 +697,65 @@ class AmazonConverseConfig(BaseConfig):
)
return openai_usage
def get_tool_call_names(
self,
tools: Optional[
Union[List[ToolBlock], List[OpenAIChatCompletionToolParam]]
] = None,
) -> List[str]:
if tools is None:
return []
tool_set: set[str] = set()
for tool in tools:
tool_spec = tool.get("toolSpec")
function = tool.get("function")
if tool_spec is not None:
_name = cast(dict, tool_spec).get("name")
if _name is not None and isinstance(_name, str):
tool_set.add(_name)
if function is not None:
_name = cast(dict, function).get("name")
if _name is not None and isinstance(_name, str):
tool_set.add(_name)
return list(tool_set)
def apply_tool_call_transformation_if_needed(
self,
message: Message,
tools: Optional[List[ToolBlock]] = None,
initial_finish_reason: Optional[str] = None,
) -> Tuple[Message, Optional[str]]:
"""
Apply tool call transformation to a message.
LLM providers (e.g. Bedrock, Vertex AI) sometimes return tool call in the response content.
If the response content is a JSON object, we can parse it and return the tool call in the tool_calls field.
"""
returned_finish_reason = initial_finish_reason
if tools is None:
return message, returned_finish_reason
if message.content is not None:
try:
tool_call_names = self.get_tool_call_names(tools)
json_content = json.loads(message.content)
if (
json_content.get("type") == "function"
and json_content.get("name") in tool_call_names
):
tool_calls = [
ChatCompletionMessageToolCall(function=Function(**json_content))
]
message.tool_calls = tool_calls
message.content = None
returned_finish_reason = "tool_calls"
except Exception:
pass
return message, returned_finish_reason
def _transform_response(
self,
model: str,
@@ -819,11 +885,23 @@ class AmazonConverseConfig(BaseConfig):
## CALCULATING USAGE - bedrock returns usage in the headers
usage = self._transform_usage(completion_response["usage"])
## HANDLE TOOL CALLS
_message = Message(**chat_completion_message)
initial_finish_reason = map_finish_reason(completion_response["stopReason"])
(
returned_message,
returned_finish_reason,
) = self.apply_tool_call_transformation_if_needed(
message=_message,
tools=optional_params.get("tools"),
initial_finish_reason=initial_finish_reason,
)
model_response.choices = [
litellm.Choices(
finish_reason=map_finish_reason(completion_response["stopReason"]),
finish_reason=returned_finish_reason,
index=0,
message=litellm.Message(**chat_completion_message),
message=returned_message,
)
]
model_response.created = int(time.time())
litellm/utils.py
+19 -19
View File
@@ -524,9 +524,9 @@ def function_setup( # noqa: PLR0915
function_id: Optional[str] = kwargs["id"] if "id" in kwargs else None
## DYNAMIC CALLBACKS ##
dynamic_callbacks: Optional[List[Union[str, Callable, CustomLogger]]] = (
kwargs.pop("callbacks", None)
)
dynamic_callbacks: Optional[
List[Union[str, Callable, CustomLogger]]
] = kwargs.pop("callbacks", None)
all_callbacks = get_dynamic_callbacks(dynamic_callbacks=dynamic_callbacks)
if len(all_callbacks) > 0:
@@ -1210,9 +1210,9 @@ def client(original_function): # noqa: PLR0915
exception=e,
retry_policy=kwargs.get("retry_policy"),
)
kwargs["retry_policy"] = (
reset_retry_policy()
) # prevent infinite loops
kwargs[
"retry_policy"
] = reset_retry_policy() # prevent infinite loops
litellm.num_retries = (
None # set retries to None to prevent infinite loops
)
@@ -2755,16 +2755,16 @@ def get_optional_params( # noqa: PLR0915
True # so that main.py adds the function call to the prompt
)
if "tools" in non_default_params:
optional_params["functions_unsupported_model"] = (
non_default_params.pop("tools")
)
optional_params[
"functions_unsupported_model"
] = non_default_params.pop("tools")
non_default_params.pop(
"tool_choice", None
) # causes ollama requests to hang
elif "functions" in non_default_params:
optional_params["functions_unsupported_model"] = (
non_default_params.pop("functions")
)
optional_params[
"functions_unsupported_model"
] = non_default_params.pop("functions")
elif (
litellm.add_function_to_prompt
): # if user opts to add it to prompt instead
@@ -2787,10 +2787,10 @@ def get_optional_params( # noqa: PLR0915
if "response_format" in non_default_params:
if provider_config is not None:
non_default_params["response_format"] = (
provider_config.get_json_schema_from_pydantic_object(
response_format=non_default_params["response_format"]
)
non_default_params[
"response_format"
] = provider_config.get_json_schema_from_pydantic_object(
response_format=non_default_params["response_format"]
)
else:
non_default_params["response_format"] = type_to_response_format_param(
@@ -3806,9 +3806,9 @@ def _count_characters(text: str) -> int:
def get_response_string(response_obj: Union[ModelResponse, ModelResponseStream]) -> str:
_choices: Union[List[Union[Choices, StreamingChoices]], List[StreamingChoices]] = (
response_obj.choices
)
_choices: Union[
List[Union[Choices, StreamingChoices]], List[StreamingChoices]
] = response_obj.choices
response_str = ""
for choice in _choices:
tests/litellm/llms/bedrock/chat/test_converse_transformation.py
+38 -1
View File
@@ -41,6 +41,7 @@ def test_transform_usage():
assert openai_usage._cache_creation_input_tokens == usage["cacheWriteInputTokens"]
assert openai_usage._cache_read_input_tokens == usage["cacheReadInputTokens"]
def test_transform_system_message():
config = AmazonConverseConfig()
@@ -96,7 +97,11 @@ def test_transform_system_message():
{
"role": "system",
"content": [
{"type": "text", "text": "Cache this!", "cache_control": {"type": "ephemeral"}},
{
"type": "text",
"text": "Cache this!",
"cache_control": {"type": "ephemeral"},
},
{"type": "text", "text": "Don't cache this!"},
],
},
@@ -121,6 +126,7 @@ def test_transform_system_message():
assert out_messages[1]["role"] == "assistant"
assert system_blocks == []
def test_transform_thinking_blocks_with_redacted_content():
thinking_blocks = [
{
@@ -139,3 +145,34 @@ def test_transform_thinking_blocks_with_redacted_content():
assert transformed_thinking_blocks[0]["type"] == "thinking"
assert transformed_thinking_blocks[1]["type"] == "redacted_thinking"
def test_apply_tool_call_transformation_if_needed():
from litellm.types.utils import Message
config = AmazonConverseConfig()
tool_calls = [
{
"type": "function",
"function": {
"name": "test_function",
"arguments": "test_arguments",
},
},
]
tool_response = {
"type": "function",
"name": "test_function",
"parameters": {"test": "test"},
}
message = Message(
role="user",
content=json.dumps(tool_response),
)
transformed_message, _ = config.apply_tool_call_transformation_if_needed(
message, tool_calls
)
assert len(transformed_message.tool_calls) == 1
assert transformed_message.tool_calls[0].function.name == "test_function"
assert transformed_message.tool_calls[0].function.arguments == json.dumps(
tool_response["parameters"]
)