scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-30 01:52:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 70

fix(bedrock): support tool search results and surface citations as annotations

Browse source 
Add an optional tool-message search_results path that maps directly to Bedrock toolResult.searchResult blocks, and convert Converse citationsContent into chat completion annotations for user-facing citation metadata.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Sameer Kankute 2026-05-28 13:06:23 +05:30
parent adb10dc938
commit 83c6a3d7eb
No known key found for this signature in database
6 changed files with 735 additions and 386 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

171
litellm/litellm_core_utils/prompt_templates/factory.py
View file

@ -4107,89 +4107,110 @@ def _convert_to_bedrock_tool_call_result(
-
"""
tool_result_content_blocks: List[BedrockToolResultContentBlock] = []
if isinstance(message["content"], str):
tool_result_content_blocks.append(
BedrockToolResultContentBlock(text=message["content"])
)
elif isinstance(message["content"], List):
content_list = message["content"]
for content in content_list:
if content["type"] == "text":
tool_result_content_blocks.append(
BedrockToolResultContentBlock(text=content["text"])
)
elif content["type"] == "image_url":
format: Optional[str] = None
if isinstance(content["image_url"], dict):
image_url = content["image_url"]["url"]
format = content["image_url"].get("format")
else:
image_url = content["image_url"]
_block: BedrockContentBlock = BedrockImageProcessor.process_image_sync(
image_url=image_url,
format=format,
)
if "image" in _block:
# Optional OpenAI tool-message extension:
# allow structured Bedrock search results on tool messages and map them
# directly to toolResult.content[].searchResult for Converse API.
#
# If `search_results` is present, we intentionally prefer it over `content`
# to avoid generating mixed text + searchResult blocks.
search_results = message.get("search_results")
if isinstance(search_results, list):
for result in search_results:
if not isinstance(result, dict):
continue
tool_result_content_blocks.append(
BedrockToolResultContentBlock(searchResult=result) # type: ignore[arg-type]
)
else:
if isinstance(message["content"], str):
tool_result_content_blocks.append(
BedrockToolResultContentBlock(text=message["content"])
)
elif isinstance(message["content"], List):
content_list = message["content"]
for content in content_list:
if content["type"] == "text":
tool_result_content_blocks.append(
BedrockToolResultContentBlock(image=_block["image"])
BedrockToolResultContentBlock(text=content["text"])
)
elif "document" in _block:
tool_result_content_blocks.append(
BedrockToolResultContentBlock(document=_block["document"])
elif content["type"] == "image_url":
format: Optional[str] = None
if isinstance(content["image_url"], dict):
image_url = content["image_url"]["url"]
format = content["image_url"].get("format")
else:
image_url = content["image_url"]
_block: BedrockContentBlock = (
BedrockImageProcessor.process_image_sync(
image_url=image_url,
format=format,
)
)
else:
verbose_logger.warning(
"Bedrock Converse: unrecognized BedrockContentBlock keys "
"%s for image_url tool-result block %s; dropping.",
list(_block.keys()),
content,
)
elif content["type"] == "file":
# Match the user-message path (_process_file_message): accept
# either file_data (base64 data URI) or file_id (server-side
# reference / URL) and hand off to BedrockImageProcessor. Raise
# BadRequestError on both-None rather than silently dropping.
file_obj = content.get("file") or {}
file_data = file_obj.get("file_data")
file_id = file_obj.get("file_id")
if file_data is None and file_id is None:
raise litellm.BadRequestError(
message="file_data and file_id cannot both be None. Got={}".format(
content
),
model="",
llm_provider="bedrock",
)
file_format = file_obj.get("format")
_file_block: BedrockContentBlock = (
BedrockImageProcessor.process_image_sync(
image_url=cast(str, file_id or file_data),
format=file_format,
)
)
if "document" in _file_block:
tool_result_content_blocks.append(
BedrockToolResultContentBlock(document=_file_block["document"])
)
elif "image" in _file_block:
tool_result_content_blocks.append(
BedrockToolResultContentBlock(image=_file_block["image"])
)
else:
verbose_logger.warning(
"Bedrock Converse: unrecognized BedrockContentBlock keys "
"%s for file tool-result block %s; dropping.",
list(_file_block.keys()),
content,
if "image" in _block:
tool_result_content_blocks.append(
BedrockToolResultContentBlock(image=_block["image"])
)
elif "document" in _block:
tool_result_content_blocks.append(
BedrockToolResultContentBlock(document=_block["document"])
)
else:
verbose_logger.warning(
"Bedrock Converse: unrecognized BedrockContentBlock keys "
"%s for image_url tool-result block %s; dropping.",
list(_block.keys()),
content,
)
elif content["type"] == "file":
# Match the user-message path (_process_file_message): accept
# either file_data (base64 data URI) or file_id (server-side
# reference / URL) and hand off to BedrockImageProcessor. Raise
# BadRequestError on both-None rather than silently dropping.
file_obj = content.get("file") or {}
file_data = file_obj.get("file_data")
file_id = file_obj.get("file_id")
if file_data is None and file_id is None:
raise litellm.BadRequestError(
message="file_data and file_id cannot both be None. Got={}".format(
content
),
model="",
llm_provider="bedrock",
)
file_format = file_obj.get("format")
_file_block: BedrockContentBlock = (
BedrockImageProcessor.process_image_sync(
image_url=cast(str, file_id or file_data),
format=file_format,
)
)
if "document" in _file_block:
tool_result_content_blocks.append(
BedrockToolResultContentBlock(
document=_file_block["document"]
)
)
elif "image" in _file_block:
tool_result_content_blocks.append(
BedrockToolResultContentBlock(image=_file_block["image"])
)
else:
verbose_logger.warning(
"Bedrock Converse: unrecognized BedrockContentBlock keys "
"%s for file tool-result block %s; dropping.",
list(_file_block.keys()),
content,
)
message.get("name", "")
id = str(message.get("tool_call_id", str(uuid.uuid4())))
tool_result = BedrockToolResultBlock(
content=tool_result_content_blocks,
toolUseId=id,
content=tool_result_content_blocks, toolUseId=id
)
if isinstance(search_results, list):
tool_result["status"] = cast(Literal["success"], "success")
content_block = BedrockContentBlock(toolResult=tool_result)
@ -5584,7 +5605,9 @@ def default_response_schema_prompt(response_schema: dict) -> str:
prompt_str = """Use this JSON schema:
```json
{}
```""".format(response_schema)
```""".format(
response_schema
)
return prompt_str

75
litellm/llms/bedrock/chat/converse_transformation.py
View file

@ -41,6 +41,7 @@ from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMExcepti
from litellm.types.llms.bedrock import *
from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionAnnotation,
ChatCompletionAssistantMessage,
ChatCompletionRedactedThinkingBlock,
ChatCompletionResponseMessage,
@ -1890,6 +1891,67 @@ class AmazonConverseConfig(BaseConfig):
return content_str, tools, reasoningContentBlocks, citationsContentBlocks
@staticmethod
def _transform_citations_to_annotations(
citations_content_blocks: Optional[List[CitationsContentBlock]],
) -> Tuple[Optional[str], Optional[List[ChatCompletionAnnotation]]]:
"""
Convert Bedrock citationsContent blocks into OpenAI-style annotations.
Returns:
citations_text: concatenated text from citationsContent.content
annotations: OpenAI URL citation annotations
"""
if not citations_content_blocks:
return None, None
annotations: List[ChatCompletionAnnotation] = []
citations_text_parts: List[str] = []
content_offset = 0
for citations_block in citations_content_blocks:
block_text = ""
for content_part in citations_block.get("content", []):
if isinstance(content_part, dict):
_text = content_part.get("text")
if isinstance(_text, str):
block_text += _text
if block_text:
citations_text_parts.append(block_text)
for citation in citations_block.get("citations", []):
if not isinstance(citation, dict):
continue
location = citation.get("location", {})
search_location = (
location.get("searchResultLocation", {})
if isinstance(location, dict)
else {}
)
start = search_location.get("start")
end = search_location.get("end")
if not isinstance(start, int) or not isinstance(end, int):
continue
annotations.append(
ChatCompletionAnnotation(
type="url_citation",
url_citation={
"start_index": content_offset + start,
"end_index": content_offset + end,
"title": str(citation.get("title", "")),
"url": str(citation.get("source", "")),
},
)
)
content_offset += len(block_text)
citations_text = "".join(citations_text_parts) if citations_text_parts else None
return citations_text, annotations or None
@staticmethod
def _unwrap_bedrock_properties(json_str: str) -> str:
"""
@ -2072,6 +2134,19 @@ class AmazonConverseConfig(BaseConfig):
provider_specific_fields
)
citations_text, annotations = self._transform_citations_to_annotations(
citationsContentBlocks
)
if citations_text:
if not content_str:
content_str = citations_text
elif content_str.strip() == ".":
# Bedrock may emit the cited sentence in citationsContent and only
# punctuation in text blocks; stitch them for user-facing content.
content_str = citations_text + content_str
if annotations:
chat_completion_message["annotations"] = annotations
if reasoningContentBlocks is not None:
chat_completion_message["reasoning_content"] = (
self._transform_reasoning_content(reasoningContentBlocks)

15
litellm/types/llms/bedrock.py
View file

@ -49,9 +49,24 @@ class DocumentBlock(TypedDict):
name: str
class SearchResultBlock(TypedDict, total=False):
"""
Search result block used in Bedrock toolResult content.
Reference:
https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_SearchResultBlock.html
"""
source: str
title: str
content: List[dict]
citations: dict
class ToolResultContentBlock(TypedDict, total=False):
image: ImageBlock
document: DocumentBlock
searchResult: SearchResultBlock
json: dict
text: str

1
litellm/types/llms/openai.py
View file

@ -965,6 +965,7 @@ ChatCompletionAssistantContentValue = (
class ChatCompletionResponseMessage(TypedDict, total=False):
content: Optional[ChatCompletionAssistantContentValue]
annotations: Optional[List[ChatCompletionAnnotation]]
tool_calls: Optional[List[ChatCompletionToolCallChunk]]
role: Literal["assistant"]
function_call: Optional[ChatCompletionToolCallFunctionChunk]

707
tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py
View file

File diff suppressed because it is too large Load diff

152
tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py
View file

@ -1646,6 +1646,74 @@ async def test_tool_message_string_content_cache_control():
assert tool_message_content[1]["cachePoint"]["type"] == "default"
@pytest.mark.asyncio
async def test_tool_message_search_results_maps_to_bedrock_search_result_block():
"""OpenAI tool message search_results should map to Bedrock searchResult blocks."""
from litellm.litellm_core_utils.prompt_templates.factory import (
BedrockConverseMessagesProcessor,
_bedrock_converse_messages_pt,
)
messages = [
{"role": "user", "content": "What is Apptio?"},
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "tooluse_a4rBqeZNRTKj2lTskvaO4H",
"type": "function",
"function": {
"name": "RAGRequest",
"arguments": '{"query":"What is Apptio?"}',
},
}
],
},
{
"role": "tool",
"tool_call_id": "tooluse_a4rBqeZNRTKj2lTskvaO4H",
"content": "Apptio is a company that makes calls to Bedrock using passthrough APIs via LiteLLM",
"search_results": [
{
"source": "Great Source of Information About Apptio",
"title": "12adbd74-46bd-4a88-88b2-0048755f6eb5",
"content": [
{
"text": "Apptio is a company that makes calls to Bedrock using passthrough APIs via LiteLLM"
}
],
"citations": {"enabled": True},
}
],
},
]
result = _bedrock_converse_messages_pt(
messages=messages,
model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
llm_provider="bedrock_converse",
)
async_result = (
await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async(
messages=messages,
model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
llm_provider="bedrock_converse",
)
)
assert result == async_result
tool_result = result[2]["content"][0]["toolResult"]
assert tool_result["toolUseId"] == "tooluse_a4rBqeZNRTKj2lTskvaO4H"
assert tool_result["status"] == "success"
assert len(tool_result["content"]) == 1
assert "searchResult" in tool_result["content"][0]
assert (
tool_result["content"][0]["searchResult"]["title"]
== "12adbd74-46bd-4a88-88b2-0048755f6eb5"
)
@pytest.mark.asyncio
async def test_assistant_tool_calls_cache_control():
"""Test that assistant tool_calls with cache_control generate cachePoint blocks."""
@ -4326,6 +4394,90 @@ def test_transform_response_finish_reason_stop_when_json_mode_filters_all_tools(
assert result.choices[0].finish_reason == "stop"
def test_transform_response_citations_content_maps_to_annotations():
from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig
from litellm.types.utils import ModelResponse
response_json = {
"metrics": {"latencyMs": 100},
"output": {
"message": {
"role": "assistant",
"content": [
{
"citationsContent": {
"content": [
{
"text": "Apptio is a company that makes calls to Bedrock using passthrough APIs via LiteLLM"
}
],
"citations": [
{
"location": {
"searchResultLocation": {
"start": 0,
"end": 42,
"searchResultIndex": 0,
}
},
"source": "https://www.apptio.com/about",
"title": "About Apptio",
}
],
}
},
{"text": "."},
],
}
},
"stopReason": "end_turn",
"usage": {
"inputTokens": 10,
"outputTokens": 5,
"totalTokens": 15,
"cacheReadInputTokenCount": 0,
"cacheReadInputTokens": 0,
"cacheWriteInputTokenCount": 0,
"cacheWriteInputTokens": 0,
},
}
class MockResponse:
def json(self):
return response_json
@property
def text(self):
return json.dumps(response_json)
config = AmazonConverseConfig()
model_response = ModelResponse()
result = config._transform_response(
model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
response=MockResponse(),
model_response=model_response,
stream=False,
logging_obj=None,
optional_params={},
api_key=None,
data=None,
messages=[],
encoding=None,
)
message = result.choices[0].message
assert message.content.startswith("Apptio is a company")
assert message.annotations is not None
assert len(message.annotations) == 1
annotation = message.annotations[0]
assert annotation["type"] == "url_citation"
assert annotation["url_citation"]["start_index"] == 0
assert annotation["url_citation"]["end_index"] == 42
assert annotation["url_citation"]["title"] == "About Apptio"
assert annotation["url_citation"]["url"] == "https://www.apptio.com/about"
def test_bedrock_tool_message_openai_file_pdf_becomes_document():
"""
OpenAI Chat Completions `{type: "file", file: {file_data: "data:application/pdf;...", filename}}`