482 lines
18 KiB
Python
482 lines
18 KiB
Python
"""AWS Bedrock-specific utilities.
|
|
|
|
This module contains utilities specific to the AWS Bedrock provider,
|
|
including reask functions, response handlers, and message formatting.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import json
|
|
import mimetypes
|
|
from textwrap import dedent
|
|
from typing import Any
|
|
|
|
from ...mode import Mode
|
|
|
|
|
|
def generate_bedrock_schema(response_model: type[Any]) -> dict[str, Any]:
|
|
"""
|
|
Generate Bedrock tool schema from a Pydantic model.
|
|
|
|
Bedrock Converse API expects tools in this format:
|
|
{
|
|
"toolSpec": {
|
|
"name": "tool_name",
|
|
"description": "tool description",
|
|
"inputSchema": {
|
|
"json": { JSON Schema }
|
|
}
|
|
}
|
|
}
|
|
"""
|
|
schema = response_model.model_json_schema()
|
|
|
|
return {
|
|
"toolSpec": {
|
|
"name": response_model.__name__,
|
|
"description": response_model.__doc__
|
|
or f"Correctly extracted `{response_model.__name__}` with all the required parameters with correct types",
|
|
"inputSchema": {"json": schema},
|
|
}
|
|
}
|
|
|
|
|
|
def reask_bedrock_json(
|
|
kwargs: dict[str, Any],
|
|
response: Any,
|
|
exception: Exception,
|
|
):
|
|
"""
|
|
Handle reask for Bedrock JSON mode when validation fails.
|
|
|
|
Kwargs modifications:
|
|
- Adds: "messages" (user message requesting JSON correction)
|
|
"""
|
|
kwargs = kwargs.copy()
|
|
reask_msgs = [response["output"]["message"]]
|
|
reask_msgs.append(
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"text": f"Correct your JSON ONLY RESPONSE, based on the following errors:\n{exception}"
|
|
},
|
|
],
|
|
}
|
|
)
|
|
kwargs["messages"].extend(reask_msgs)
|
|
return kwargs
|
|
|
|
|
|
def reask_bedrock_tools(
|
|
kwargs: dict[str, Any],
|
|
response: Any,
|
|
exception: Exception,
|
|
):
|
|
"""
|
|
Handle reask for Bedrock tools mode when validation fails.
|
|
|
|
Kwargs modifications:
|
|
- Adds: "messages" (assistant message with tool use, then user message with tool result error)
|
|
"""
|
|
kwargs = kwargs.copy()
|
|
|
|
# Add the assistant's response message
|
|
assistant_message = response["output"]["message"]
|
|
reask_msgs = [assistant_message]
|
|
|
|
# Find the tool use ID from the assistant's response to reference in the error
|
|
tool_use_id = None
|
|
if "content" in assistant_message:
|
|
for content_block in assistant_message["content"]:
|
|
if "toolUse" in content_block:
|
|
tool_use_id = content_block["toolUse"]["toolUseId"]
|
|
break
|
|
|
|
# Add a user message with tool result indicating validation error
|
|
if tool_use_id:
|
|
reask_msgs.append(
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"toolResult": {
|
|
"toolUseId": tool_use_id,
|
|
"content": [
|
|
{
|
|
"text": f"Validation Error found:\n{exception}\nRecall the function correctly, fix the errors"
|
|
}
|
|
],
|
|
"status": "error",
|
|
}
|
|
}
|
|
],
|
|
}
|
|
)
|
|
else:
|
|
# Fallback if no tool use ID found
|
|
reask_msgs.append(
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"text": f"Validation Error due to no tool invocation:\n{exception}\nRecall the function correctly, fix the errors"
|
|
}
|
|
],
|
|
}
|
|
)
|
|
|
|
kwargs["messages"].extend(reask_msgs)
|
|
return kwargs
|
|
|
|
|
|
def _normalize_bedrock_image_format(mime_or_ext: str) -> str:
|
|
"""
|
|
Map common/variant image types to Bedrock's required image.format enum:
|
|
one of {'gif','jpeg','png','webp'}.
|
|
"""
|
|
if not mime_or_ext:
|
|
return "jpeg"
|
|
val = mime_or_ext.strip().lower()
|
|
if "/" in val:
|
|
val = val.split("/", 1)[1] # take subtype, e.g., 'image/jpeg' -> 'jpeg'
|
|
if val in ("jpg", "pjpeg", "x-jpeg", "x-jpg"):
|
|
return "jpeg"
|
|
if val in ("png", "x-png"):
|
|
return "png"
|
|
if val in ("gif", "x-gif"):
|
|
return "gif"
|
|
if val in ("webp", "image/webp"):
|
|
return "webp"
|
|
return "jpeg"
|
|
|
|
|
|
def _openai_image_part_to_bedrock(part: dict[str, Any]) -> dict[str, Any]:
|
|
"""
|
|
Convert OpenAI-style image part:
|
|
{"type":"image_url","image_url":{"url": "<data:...>"}}
|
|
into Bedrock Converse image content:
|
|
{"image":{"format": "<fmt>","source":{"bytes": <raw-bytes>}}}
|
|
"""
|
|
image_url = (part.get("image_url") or {}).get("url")
|
|
if not image_url:
|
|
raise ValueError("image_url.url is required for OpenAI-style image parts")
|
|
|
|
# data URL to bytes
|
|
if image_url.startswith("data:"):
|
|
try:
|
|
header, b64 = image_url.split(",", 1)
|
|
except ValueError as e:
|
|
raise ValueError("Invalid data URL in image_url.url") from e
|
|
if ";base64" not in header:
|
|
raise ValueError("Only base64 data URLs are supported for Bedrock")
|
|
meta = header[5:] # strip 'data:'
|
|
mime = meta.split(";", 1)[0]
|
|
if not mime or "/" not in mime:
|
|
guessed = None
|
|
for token in meta.split(";")[1:]:
|
|
if token.startswith("name="):
|
|
name = token[len("name=") :].strip().strip('"')
|
|
guessed = mimetypes.guess_type(name)[0]
|
|
if guessed:
|
|
break
|
|
mime = guessed or "image/jpeg"
|
|
fmt = _normalize_bedrock_image_format(mime)
|
|
return {"image": {"format": fmt, "source": {"bytes": base64.b64decode(b64)}}}
|
|
|
|
else:
|
|
raise ValueError(
|
|
"Unsupported image_url scheme for Bedrock. "
|
|
"Use data:image/...;base64,... or pass Bedrock-native image bytes."
|
|
)
|
|
|
|
|
|
def _to_bedrock_content_items(content: Any) -> list[dict[str, Any]]:
|
|
"""
|
|
Normalize content into Bedrock Converse content list.
|
|
|
|
Allowed inputs:
|
|
- string -> [{"text": "..."}]
|
|
- list of parts:
|
|
OpenAI-style:
|
|
{"type":"text","text":"..."}
|
|
{"type":"input_text","text":"..."}
|
|
{"type":"image_url","image_url":{"url":"<data:... or https:...>"}}
|
|
Bedrock-native (passed through as-is):
|
|
{"text":"..."}
|
|
{"image":{"format":"jpeg|png|gif|webp","source":{"bytes": <raw bytes>}}}
|
|
{"document":{"format":"pdf|csv|doc|docx|xls|xlsx|html|txt|md","name":"...","source":{"bytes": <raw bytes>}}}
|
|
|
|
Note:
|
|
- We do not validate or normalize Bedrock-native image/document blocks here.
|
|
Caller is responsible for providing valid 'format' and raw 'bytes'.
|
|
"""
|
|
# Plain string
|
|
if isinstance(content, str):
|
|
return [{"text": content}]
|
|
|
|
# List of parts
|
|
if isinstance(content, list):
|
|
items: list[dict[str, Any]] = []
|
|
for p in content:
|
|
# OpenAI-style parts (have "type")
|
|
if isinstance(p, dict) and "type" in p:
|
|
t = p.get("type")
|
|
if t in ("text", "input_text"):
|
|
txt = p.get("text") or p.get("input_text") or ""
|
|
items.append({"text": txt})
|
|
continue
|
|
if t == "image_url":
|
|
items.append(_openai_image_part_to_bedrock(p))
|
|
continue
|
|
raise ValueError(f"Unsupported OpenAI-style part type for Bedrock: {t}")
|
|
|
|
# Bedrock-native pass-throughs (no "type")
|
|
if isinstance(p, dict):
|
|
# Pass-through pure text
|
|
if (
|
|
"text" in p
|
|
and isinstance(p["text"], str)
|
|
and set(p.keys()) == {"text"}
|
|
):
|
|
items.append(p)
|
|
continue
|
|
# Pass-through Bedrock-native image as-is (assumes correct format and raw bytes)
|
|
if "image" in p and isinstance(p["image"], dict):
|
|
items.append(p)
|
|
continue
|
|
# Pass-through Bedrock-native document as-is (assumes correct format and raw bytes)
|
|
if "document" in p and isinstance(p["document"], dict):
|
|
items.append(p)
|
|
continue
|
|
# Pass-through Bedrock cache point as-is
|
|
if "cachePoint" in p:
|
|
items.append(p)
|
|
continue
|
|
|
|
raise ValueError(f"Unsupported dict content for Bedrock: {p}")
|
|
|
|
# Plain string elements inside list
|
|
if isinstance(p, str):
|
|
items.append({"text": p})
|
|
continue
|
|
|
|
raise ValueError(f"Unsupported content part for Bedrock: {type(p)}")
|
|
return items
|
|
|
|
raise ValueError(f"Unsupported message content type for Bedrock: {type(content)}")
|
|
|
|
|
|
def _prepare_bedrock_converse_kwargs_internal(
|
|
call_kwargs: dict[str, Any],
|
|
) -> dict[str, Any]:
|
|
"""
|
|
Prepare kwargs for the Bedrock Converse API.
|
|
|
|
Kwargs modifications:
|
|
- Moves: system list to messages as a system role
|
|
- Renames: "model" -> "modelId"
|
|
- Collects: temperature, max_tokens, top_p, stop into inferenceConfig
|
|
- Converts: messages content to Bedrock format
|
|
"""
|
|
# Handle Bedrock-native system parameter format: system=[{'text': '...'}]
|
|
# Convert to OpenAI format by adding to messages as system role
|
|
if "system" in call_kwargs and isinstance(call_kwargs["system"], list):
|
|
system_content = call_kwargs.pop("system")
|
|
if (
|
|
system_content
|
|
and isinstance(system_content[0], dict)
|
|
and "text" in system_content[0]
|
|
):
|
|
# Convert system=[{'text': '...'}] to OpenAI format
|
|
system_text = system_content[0]["text"]
|
|
if "messages" not in call_kwargs:
|
|
call_kwargs["messages"] = []
|
|
# Insert system message at beginning
|
|
call_kwargs["messages"].insert(
|
|
0, {"role": "system", "content": system_text}
|
|
)
|
|
|
|
# Bedrock expects 'modelId' over 'model'
|
|
if "model" in call_kwargs and "modelId" not in call_kwargs:
|
|
call_kwargs["modelId"] = call_kwargs.pop("model")
|
|
|
|
# Prepare inferenceConfig for parameters like temperature, maxTokens, etc.
|
|
inference_config_params = {}
|
|
|
|
# Temperature
|
|
if "temperature" in call_kwargs:
|
|
inference_config_params["temperature"] = call_kwargs.pop("temperature")
|
|
|
|
# Max Tokens (OpenAI uses max_tokens)
|
|
if "max_tokens" in call_kwargs:
|
|
inference_config_params["maxTokens"] = call_kwargs.pop("max_tokens")
|
|
elif "maxTokens" in call_kwargs: # If Bedrock-style maxTokens is already top-level
|
|
inference_config_params["maxTokens"] = call_kwargs.pop("maxTokens")
|
|
|
|
# Top P (OpenAI uses top_p)
|
|
if "top_p" in call_kwargs:
|
|
inference_config_params["topP"] = call_kwargs.pop("top_p")
|
|
elif "topP" in call_kwargs: # If Bedrock-style topP is already top-level
|
|
inference_config_params["topP"] = call_kwargs.pop("topP")
|
|
|
|
# Stop Sequences (OpenAI uses 'stop')
|
|
# Bedrock 'Converse' API expects 'stopSequences'
|
|
if "stop" in call_kwargs:
|
|
stop_val = call_kwargs.pop("stop")
|
|
if isinstance(stop_val, str):
|
|
inference_config_params["stopSequences"] = [stop_val]
|
|
elif isinstance(stop_val, list):
|
|
inference_config_params["stopSequences"] = stop_val
|
|
elif "stop_sequences" in call_kwargs:
|
|
inference_config_params["stopSequences"] = call_kwargs.pop("stop_sequences")
|
|
elif (
|
|
"stopSequences" in call_kwargs
|
|
): # If Bedrock-style stopSequences is already top-level
|
|
inference_config_params["stopSequences"] = call_kwargs.pop("stopSequences")
|
|
|
|
# If any inference parameters were collected, add them to inferenceConfig
|
|
# Merge with existing inferenceConfig if user provided one.
|
|
# User-provided inferenceConfig keys take precedence over top-level params if conflicts.
|
|
if inference_config_params:
|
|
if "inferenceConfig" in call_kwargs:
|
|
# Merge, giving precedence to what's already in call_kwargs["inferenceConfig"]
|
|
# This could be more sophisticated, but for now, if inferenceConfig is set, assume it's intentional.
|
|
existing_inference_config = call_kwargs["inferenceConfig"]
|
|
for key, value in inference_config_params.items():
|
|
if key not in existing_inference_config:
|
|
existing_inference_config[key] = value
|
|
else:
|
|
call_kwargs["inferenceConfig"] = inference_config_params
|
|
|
|
# Process messages for Bedrock: separate system prompts and format text content.
|
|
if "messages" in call_kwargs and isinstance(call_kwargs["messages"], list):
|
|
original_input_messages = call_kwargs.pop("messages")
|
|
|
|
bedrock_system_list: list[dict[str, Any]] = []
|
|
bedrock_user_assistant_messages_list: list[dict[str, Any]] = []
|
|
|
|
for msg_dict in original_input_messages:
|
|
if not isinstance(msg_dict, dict):
|
|
# If an item in the messages list is not a dictionary,
|
|
# pass it through to the user/assistant messages list as is.
|
|
# This allows non-standard message items to be handled by subsequent Boto3 validation
|
|
# or if they represent something other than standard role/content messages.
|
|
bedrock_user_assistant_messages_list.append(msg_dict)
|
|
continue
|
|
|
|
# Make a copy to avoid modifying the original dict if it's part of a larger structure
|
|
# or if the original list/dicts are expected to remain unchanged by the caller.
|
|
current_message_for_api = msg_dict.copy()
|
|
role = current_message_for_api.get("role")
|
|
content = current_message_for_api.get(
|
|
"content"
|
|
) # content can be None or other types
|
|
|
|
if role == "system":
|
|
if isinstance(content, str):
|
|
bedrock_system_list.append({"text": content})
|
|
else: # System message content is not a string (could be None, list, int, etc.)
|
|
raise ValueError(
|
|
"System message content must be a string for Bedrock processing by this handler. "
|
|
f"Found type: {type(content)}."
|
|
)
|
|
else: # For user, assistant, or other roles that go into Bedrock's 'messages' list
|
|
if "content" in current_message_for_api:
|
|
# Sort out the content from the messages
|
|
current_message_for_api["content"] = _to_bedrock_content_items(
|
|
content
|
|
)
|
|
bedrock_user_assistant_messages_list.append(current_message_for_api)
|
|
|
|
if bedrock_system_list:
|
|
call_kwargs["system"] = bedrock_system_list
|
|
|
|
# Always re-assign the 'messages' key with the processed list.
|
|
# If original_input_messages was empty or only contained system messages that were extracted,
|
|
# bedrock_user_assistant_messages_list will be empty, correctly resulting in `messages: []`.
|
|
call_kwargs["messages"] = bedrock_user_assistant_messages_list
|
|
return call_kwargs
|
|
|
|
|
|
def handle_bedrock_json(
|
|
response_model: type[Any], new_kwargs: dict[str, Any]
|
|
) -> tuple[type[Any], dict[str, Any]]:
|
|
"""
|
|
Handle Bedrock JSON mode.
|
|
|
|
Kwargs modifications:
|
|
- Adds: "response_format" with json_schema
|
|
- Adds/Modifies: "system" (prepends JSON instructions)
|
|
- Applies: _prepare_bedrock_converse_kwargs_internal transformations
|
|
"""
|
|
new_kwargs = _prepare_bedrock_converse_kwargs_internal(new_kwargs)
|
|
json_message = dedent(
|
|
f"""
|
|
As a genius expert, your task is to understand the content and provide
|
|
the parsed objects in json that match the following json_schema:\n
|
|
|
|
{json.dumps(response_model.model_json_schema(), indent=2, ensure_ascii=False)}
|
|
|
|
Make sure to return an instance of the JSON, not the schema itself
|
|
and don't include any other text in the response apart from the json
|
|
"""
|
|
)
|
|
system_message = new_kwargs.pop("system", None)
|
|
if not system_message:
|
|
new_kwargs["system"] = [{"text": json_message}]
|
|
else:
|
|
if not isinstance(system_message, list):
|
|
raise ValueError(
|
|
"""system must be a list of SystemMessage, refer to:
|
|
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-runtime/client/converse.html
|
|
"""
|
|
)
|
|
system_message.append({"text": json_message})
|
|
new_kwargs["system"] = system_message
|
|
|
|
return response_model, new_kwargs
|
|
|
|
|
|
def handle_bedrock_tools(
|
|
response_model: type[Any] | None, new_kwargs: dict[str, Any]
|
|
) -> tuple[type[Any] | None, dict[str, Any]]:
|
|
"""
|
|
Handle Bedrock tools mode.
|
|
|
|
Kwargs modifications:
|
|
- When response_model is None: Only applies _prepare_bedrock_converse_kwargs_internal transformations
|
|
- When response_model is provided:
|
|
- Adds: "toolConfig" with tools list and toolChoice configuration
|
|
- Applies: _prepare_bedrock_converse_kwargs_internal transformations
|
|
"""
|
|
new_kwargs = _prepare_bedrock_converse_kwargs_internal(new_kwargs)
|
|
|
|
if response_model is None:
|
|
return None, new_kwargs
|
|
|
|
# Generate Bedrock tool schema
|
|
tool_schema = generate_bedrock_schema(response_model)
|
|
|
|
# Set up tools configuration for Bedrock Converse API
|
|
new_kwargs["toolConfig"] = {
|
|
"tools": [tool_schema],
|
|
"toolChoice": {"tool": {"name": response_model.__name__}},
|
|
}
|
|
|
|
return response_model, new_kwargs
|
|
|
|
|
|
# Handler registry for Bedrock
|
|
BEDROCK_HANDLERS = {
|
|
Mode.BEDROCK_JSON: {
|
|
"reask": reask_bedrock_json,
|
|
"response": handle_bedrock_json,
|
|
},
|
|
Mode.BEDROCK_TOOLS: {
|
|
"reask": reask_bedrock_tools,
|
|
"response": handle_bedrock_tools,
|
|
},
|
|
}
|