mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-05 08:23:17 +00:00
Merge pull request #14005 from moshemorad/bedrock_fix_structure_output
Bedrock fix structure output
This commit is contained in:
Krish Dholakia
committed by
Mohse Morad
Mohse Morad
parent
0c9051abba
commit
fd0f47d48d
Vendored
BIN
Binary file not shown.
@@ -10,6 +10,7 @@ from typing import List, Literal, Optional, Tuple, Union, cast, overload
|
||||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm.constants import RESPONSE_FORMAT_TOOL_NAME
|
||||
from litellm.litellm_core_utils.core_helpers import map_finish_reason
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
|
||||
@@ -322,7 +323,6 @@ class AmazonConverseConfig(BaseConfig):
|
||||
def _create_json_tool_call_for_response_format(
|
||||
self,
|
||||
json_schema: Optional[dict] = None,
|
||||
schema_name: str = "json_tool_call",
|
||||
description: Optional[str] = None,
|
||||
) -> ChatCompletionToolParam:
|
||||
"""
|
||||
@@ -347,7 +347,7 @@ class AmazonConverseConfig(BaseConfig):
|
||||
_input_schema = json_schema
|
||||
|
||||
tool_param_function_chunk = ChatCompletionToolParamFunctionChunk(
|
||||
name=schema_name, parameters=_input_schema
|
||||
name=RESPONSE_FORMAT_TOOL_NAME, parameters=_input_schema
|
||||
)
|
||||
if description:
|
||||
tool_param_function_chunk["description"] = description
|
||||
@@ -391,14 +391,11 @@ class AmazonConverseConfig(BaseConfig):
|
||||
continue
|
||||
|
||||
json_schema: Optional[dict] = None
|
||||
schema_name: str = ""
|
||||
description: Optional[str] = None
|
||||
if "response_schema" in value:
|
||||
json_schema = value["response_schema"]
|
||||
schema_name = "json_tool_call"
|
||||
elif "json_schema" in value:
|
||||
json_schema = value["json_schema"]["schema"]
|
||||
schema_name = value["json_schema"]["name"]
|
||||
description = value["json_schema"].get("description")
|
||||
|
||||
if "type" in value and value["type"] == "text":
|
||||
@@ -414,7 +411,6 @@ class AmazonConverseConfig(BaseConfig):
|
||||
"""
|
||||
_tool = self._create_json_tool_call_for_response_format(
|
||||
json_schema=json_schema,
|
||||
schema_name=schema_name if schema_name != "" else "json_tool_call",
|
||||
description=description,
|
||||
)
|
||||
optional_params = self._add_tools_to_optional_params(
|
||||
@@ -430,7 +426,7 @@ class AmazonConverseConfig(BaseConfig):
|
||||
|
||||
optional_params["tool_choice"] = ToolChoiceValuesBlock(
|
||||
tool=SpecificToolChoiceBlock(
|
||||
name=schema_name if schema_name != "" else "json_tool_call"
|
||||
name=RESPONSE_FORMAT_TOOL_NAME
|
||||
)
|
||||
)
|
||||
optional_params["json_mode"] = True
|
||||
@@ -1119,8 +1115,7 @@ class AmazonConverseConfig(BaseConfig):
|
||||
self._transform_thinking_blocks(reasoningContentBlocks)
|
||||
)
|
||||
chat_completion_message["content"] = content_str
|
||||
if json_mode is True and tools is not None and len(tools) == 1:
|
||||
# to support 'json_schema' logic on bedrock models
|
||||
if json_mode is True and tools is not None and len(tools) == 1 and tools[0]["function"]["name"] == RESPONSE_FORMAT_TOOL_NAME:
|
||||
json_mode_content_str: Optional[str] = tools[0]["function"].get("arguments")
|
||||
if json_mode_content_str is not None:
|
||||
chat_completion_message["content"] = json_mode_content_str
|
||||
|
||||
@@ -475,6 +475,239 @@ def test_transform_response_with_bash_tool():
|
||||
assert args["command"] == "ls -la *.py"
|
||||
|
||||
|
||||
def test_transform_response_with_structured_response_being_called():
|
||||
"""Test response transformation with structured response."""
|
||||
from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
# Simulate a Bedrock Converse response with a bash tool call
|
||||
response_json = {
|
||||
"additionalModelResponseFields": {},
|
||||
"metrics": {"latencyMs": 100.0},
|
||||
"output": {
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"toolUse": {
|
||||
"toolUseId": "tooluse_456",
|
||||
"name": "json_tool_call",
|
||||
"input": {
|
||||
"Current_Temperature": 62,
|
||||
"Weather_Explanation": "San Francisco typically has mild, cool weather year-round due to its coastal location and marine influence. The city is known for its fog, moderate temperatures, and relatively stable climate with little seasonal variation."},
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"stopReason": "tool_use",
|
||||
"usage": {
|
||||
"inputTokens": 8,
|
||||
"outputTokens": 3,
|
||||
"totalTokens": 11,
|
||||
"cacheReadInputTokenCount": 0,
|
||||
"cacheReadInputTokens": 0,
|
||||
"cacheWriteInputTokenCount": 0,
|
||||
"cacheWriteInputTokens": 0,
|
||||
},
|
||||
}
|
||||
# Mock httpx.Response
|
||||
class MockResponse:
|
||||
def json(self):
|
||||
return response_json
|
||||
@property
|
||||
def text(self):
|
||||
return json.dumps(response_json)
|
||||
|
||||
config = AmazonConverseConfig()
|
||||
model_response = ModelResponse()
|
||||
optional_params = {
|
||||
"json_mode": True,
|
||||
"tools": [
|
||||
{
|
||||
'type': 'function',
|
||||
'function': {
|
||||
'name': 'get_weather',
|
||||
'description': 'Get the current weather in a given location',
|
||||
'parameters': {
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'location': {
|
||||
'type': 'string',
|
||||
'description': 'The city and state, e.g. San Francisco, CA'
|
||||
},
|
||||
'unit': {
|
||||
'type': 'string',
|
||||
'enum': ['celsius', 'fahrenheit']
|
||||
}
|
||||
},
|
||||
'required': ['location']
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
'type': 'function',
|
||||
'function': {
|
||||
'name': 'json_tool_call',
|
||||
'parameters': {
|
||||
'$schema': 'http://json-schema.org/draft-07/schema#',
|
||||
'type': 'object',
|
||||
'required': ['Weather_Explanation', 'Current_Temperature'],
|
||||
'properties': {
|
||||
'Weather_Explanation': {
|
||||
'type': ['string', 'null'],
|
||||
'description': '1-2 sentences explaining the weather in the location'
|
||||
},
|
||||
'Current_Temperature': {
|
||||
'type': ['number', 'null'],
|
||||
'description': 'Current temperature in the location'
|
||||
}
|
||||
},
|
||||
'additionalProperties': False
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
# Call the transformation logic
|
||||
result = config._transform_response(
|
||||
model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
|
||||
response=MockResponse(),
|
||||
model_response=model_response,
|
||||
stream=False,
|
||||
logging_obj=None,
|
||||
optional_params=optional_params,
|
||||
api_key=None,
|
||||
data=None,
|
||||
messages=[],
|
||||
encoding=None,
|
||||
)
|
||||
# Check that the tool call is present in the returned message
|
||||
assert result.choices[0].message.tool_calls is None
|
||||
|
||||
assert result.choices[0].message.content is not None
|
||||
assert result.choices[0].message.content == '{"Current_Temperature": 62, "Weather_Explanation": "San Francisco typically has mild, cool weather year-round due to its coastal location and marine influence. The city is known for its fog, moderate temperatures, and relatively stable climate with little seasonal variation."}'
|
||||
|
||||
def test_transform_response_with_structured_response_calling_tool():
|
||||
"""Test response transformation with structured response."""
|
||||
from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
# Simulate a Bedrock Converse response with a bash tool call
|
||||
response_json = {
|
||||
"metrics": {
|
||||
"latencyMs": 1148
|
||||
},
|
||||
"output": {
|
||||
"message":
|
||||
{
|
||||
"content": [
|
||||
{
|
||||
"text": "I\'ll check the current weather in San Francisco for you."
|
||||
},
|
||||
{
|
||||
"toolUse": {
|
||||
"input": {
|
||||
"location": "San Francisco, CA",
|
||||
"unit": "celsius"
|
||||
},
|
||||
"name": "get_weather",
|
||||
"toolUseId": "tooluse_oKk__QrqSUmufMw3Q7vGaQ"
|
||||
}
|
||||
}
|
||||
],
|
||||
"role": "assistant"
|
||||
}
|
||||
},
|
||||
"stopReason": "tool_use",
|
||||
"usage": {
|
||||
"cacheReadInputTokenCount": 0,
|
||||
"cacheReadInputTokens": 0,
|
||||
"cacheWriteInputTokenCount": 0,
|
||||
"cacheWriteInputTokens": 0,
|
||||
"inputTokens": 534,
|
||||
"outputTokens": 69,
|
||||
"totalTokens": 603
|
||||
}
|
||||
}
|
||||
# Mock httpx.Response
|
||||
class MockResponse:
|
||||
def json(self):
|
||||
return response_json
|
||||
@property
|
||||
def text(self):
|
||||
return json.dumps(response_json)
|
||||
|
||||
config = AmazonConverseConfig()
|
||||
model_response = ModelResponse()
|
||||
optional_params = {
|
||||
"json_mode": True,
|
||||
"tools": [
|
||||
{
|
||||
'type': 'function',
|
||||
'function': {
|
||||
'name': 'get_weather',
|
||||
'description': 'Get the current weather in a given location',
|
||||
'parameters': {
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'location': {
|
||||
'type': 'string',
|
||||
'description': 'The city and state, e.g. San Francisco, CA'
|
||||
},
|
||||
'unit': {
|
||||
'type': 'string',
|
||||
'enum': ['celsius', 'fahrenheit']
|
||||
}
|
||||
},
|
||||
'required': ['location']
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
'type': 'function',
|
||||
'function': {
|
||||
'name': 'json_tool_call',
|
||||
'parameters': {
|
||||
'$schema': 'http://json-schema.org/draft-07/schema#',
|
||||
'type': 'object',
|
||||
'required': ['Weather_Explanation', 'Current_Temperature'],
|
||||
'properties': {
|
||||
'Weather_Explanation': {
|
||||
'type': ['string', 'null'],
|
||||
'description': '1-2 sentences explaining the weather in the location'
|
||||
},
|
||||
'Current_Temperature': {
|
||||
'type': ['number', 'null'],
|
||||
'description': 'Current temperature in the location'
|
||||
}
|
||||
},
|
||||
'additionalProperties': False
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
# Call the transformation logic
|
||||
result = config._transform_response(
|
||||
model="bedrock/eu.anthropic.claude-sonnet-4-20250514-v1:0",
|
||||
response=MockResponse(),
|
||||
model_response=model_response,
|
||||
stream=False,
|
||||
logging_obj=None,
|
||||
optional_params=optional_params,
|
||||
api_key=None,
|
||||
data=None,
|
||||
messages=[],
|
||||
encoding=None,
|
||||
)
|
||||
# Check that the tool call is present in the returned message
|
||||
assert result.choices[0].message.tool_calls is not None
|
||||
assert len(result.choices[0].message.tool_calls) == 1
|
||||
assert result.choices[0].message.tool_calls[0].function.name == "get_weather"
|
||||
assert result.choices[0].message.tool_calls[0].function.arguments == '{"location": "San Francisco, CA", "unit": "celsius"}'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bedrock_bash_tool_acompletion():
|
||||
"""Test Bedrock with bash tool for ls command using acompletion."""
|
||||
@@ -938,6 +1171,68 @@ def test_transform_request_with_function_tool():
|
||||
assert request_data["toolConfig"]["tools"][0]["toolSpec"]["name"] == "get_weather"
|
||||
|
||||
|
||||
def test_map_openai_params_with_response_format():
|
||||
"""Test map_openai_params with response_format."""
|
||||
config = AmazonConverseConfig()
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "Get the current weather in a given location",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"location": {
|
||||
"type": "string",
|
||||
"description": "The city and state, e.g. San Francisco, CA",
|
||||
},
|
||||
"unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
|
||||
},
|
||||
"required": ["location"],
|
||||
},
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
json_schema = {
|
||||
"type": "json_schema",
|
||||
"json_schema": {
|
||||
"name": "WeatherResult",
|
||||
"schema": {
|
||||
"$schema": "http://json-schema.org/draft-07/schema#",
|
||||
"type": "object",
|
||||
"required": ["Weather_Explanation", "Current_Temperature"],
|
||||
"properties": {
|
||||
"Weather_Explanation": {
|
||||
"type": ["string", "null"],
|
||||
"description": "1-2 sentences explaining the weather in the location",
|
||||
},
|
||||
"Current_Temperature": {
|
||||
"type": ["number", "null"],
|
||||
"description": "Current temperature in the location",
|
||||
},
|
||||
},
|
||||
"additionalProperties": False,
|
||||
},
|
||||
"strict": False,
|
||||
},
|
||||
}
|
||||
|
||||
optional_params = config.map_openai_params(
|
||||
non_default_params={"response_format": json_schema},
|
||||
optional_params={"tools": tools},
|
||||
model="eu.anthropic.claude-sonnet-4-20250514-v1:0",
|
||||
drop_params=False
|
||||
)
|
||||
|
||||
assert "tools" in optional_params
|
||||
assert len(optional_params["tools"]) == 2
|
||||
assert optional_params["tools"][1]["type"] == "function"
|
||||
assert optional_params["tools"][1]["function"]["name"] == "json_tool_call"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_assistant_message_cache_control():
|
||||
"""Test that assistant messages with cache_control generate cachePoint blocks."""
|
||||
|
||||
Reference in New Issue
Block a user