Update AWSBedrock docstrings
This commit is contained in:
@@ -4,6 +4,13 @@
|
||||
# SPDX-License-Identifier: BSD 2-Clause License
|
||||
#
|
||||
|
||||
"""AWS Bedrock integration for Large Language Model services.
|
||||
|
||||
This module provides AWS Bedrock LLM service implementation with support for
|
||||
Amazon Nova and Anthropic Claude models, including vision capabilities and
|
||||
function calling.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import copy
|
||||
@@ -61,17 +68,50 @@ except ModuleNotFoundError as e:
|
||||
|
||||
@dataclass
|
||||
class AWSBedrockContextAggregatorPair:
|
||||
"""Container for AWS Bedrock context aggregators.
|
||||
|
||||
Provides convenient access to both user and assistant context aggregators
|
||||
for AWS Bedrock LLM operations.
|
||||
|
||||
Parameters:
|
||||
_user: The user context aggregator instance.
|
||||
_assistant: The assistant context aggregator instance.
|
||||
"""
|
||||
|
||||
_user: "AWSBedrockUserContextAggregator"
|
||||
_assistant: "AWSBedrockAssistantContextAggregator"
|
||||
|
||||
def user(self) -> "AWSBedrockUserContextAggregator":
|
||||
"""Get the user context aggregator.
|
||||
|
||||
Returns:
|
||||
The user context aggregator instance.
|
||||
"""
|
||||
return self._user
|
||||
|
||||
def assistant(self) -> "AWSBedrockAssistantContextAggregator":
|
||||
"""Get the assistant context aggregator.
|
||||
|
||||
Returns:
|
||||
The assistant context aggregator instance.
|
||||
"""
|
||||
return self._assistant
|
||||
|
||||
|
||||
class AWSBedrockLLMContext(OpenAILLMContext):
|
||||
"""AWS Bedrock-specific LLM context implementation.
|
||||
|
||||
Extends OpenAI LLM context to handle AWS Bedrock's specific message format
|
||||
and system message handling. Manages conversion between OpenAI and Bedrock
|
||||
message formats.
|
||||
|
||||
Args:
|
||||
messages: List of conversation messages in OpenAI format.
|
||||
tools: List of available function calling tools.
|
||||
tool_choice: Tool selection strategy or specific tool choice.
|
||||
system: System message content for AWS Bedrock.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
messages: Optional[List[dict]] = None,
|
||||
@@ -85,6 +125,14 @@ class AWSBedrockLLMContext(OpenAILLMContext):
|
||||
|
||||
@staticmethod
|
||||
def upgrade_to_bedrock(obj: OpenAILLMContext) -> "AWSBedrockLLMContext":
|
||||
"""Upgrade an OpenAI LLM context to AWS Bedrock format.
|
||||
|
||||
Args:
|
||||
obj: The OpenAI LLM context to upgrade.
|
||||
|
||||
Returns:
|
||||
The upgraded AWS Bedrock LLM context.
|
||||
"""
|
||||
logger.debug(f"Upgrading to AWS Bedrock: {obj}")
|
||||
if isinstance(obj, OpenAILLMContext) and not isinstance(obj, AWSBedrockLLMContext):
|
||||
obj.__class__ = AWSBedrockLLMContext
|
||||
@@ -95,6 +143,14 @@ class AWSBedrockLLMContext(OpenAILLMContext):
|
||||
|
||||
@classmethod
|
||||
def from_openai_context(cls, openai_context: OpenAILLMContext):
|
||||
"""Create AWS Bedrock context from OpenAI context.
|
||||
|
||||
Args:
|
||||
openai_context: The OpenAI LLM context to convert.
|
||||
|
||||
Returns:
|
||||
New AWS Bedrock LLM context instance.
|
||||
"""
|
||||
self = cls(
|
||||
messages=openai_context.messages,
|
||||
tools=openai_context.tools,
|
||||
@@ -106,12 +162,28 @@ class AWSBedrockLLMContext(OpenAILLMContext):
|
||||
|
||||
@classmethod
|
||||
def from_messages(cls, messages: List[dict]) -> "AWSBedrockLLMContext":
|
||||
"""Create AWS Bedrock context from message list.
|
||||
|
||||
Args:
|
||||
messages: List of messages in OpenAI format.
|
||||
|
||||
Returns:
|
||||
New AWS Bedrock LLM context instance.
|
||||
"""
|
||||
self = cls(messages=messages)
|
||||
self._restructure_from_openai_messages()
|
||||
return self
|
||||
|
||||
@classmethod
|
||||
def from_image_frame(cls, frame: VisionImageRawFrame) -> "AWSBedrockLLMContext":
|
||||
"""Create AWS Bedrock context from vision image frame.
|
||||
|
||||
Args:
|
||||
frame: The vision image frame to convert.
|
||||
|
||||
Returns:
|
||||
New AWS Bedrock LLM context instance.
|
||||
"""
|
||||
context = cls()
|
||||
context.add_image_frame_message(
|
||||
format=frame.format, size=frame.size, image=frame.image, text=frame.text
|
||||
@@ -119,10 +191,14 @@ class AWSBedrockLLMContext(OpenAILLMContext):
|
||||
return context
|
||||
|
||||
def set_messages(self, messages: List):
|
||||
"""Set the messages list and restructure for Bedrock format.
|
||||
|
||||
Args:
|
||||
messages: List of messages to set.
|
||||
"""
|
||||
self._messages[:] = messages
|
||||
self._restructure_from_openai_messages()
|
||||
|
||||
# convert a message in AWS Bedrock format into one or more messages in OpenAI format
|
||||
def to_standard_messages(self, obj):
|
||||
"""Convert AWS Bedrock message format to standard structured format.
|
||||
|
||||
@@ -295,6 +371,14 @@ class AWSBedrockLLMContext(OpenAILLMContext):
|
||||
def add_image_frame_message(
|
||||
self, *, format: str, size: tuple[int, int], image: bytes, text: str = None
|
||||
):
|
||||
"""Add an image message to the context.
|
||||
|
||||
Args:
|
||||
format: The image format (e.g., 'RGB', 'RGBA').
|
||||
size: The image dimensions as (width, height).
|
||||
image: The raw image data as bytes.
|
||||
text: Optional text to accompany the image.
|
||||
"""
|
||||
buffer = io.BytesIO()
|
||||
Image.frombytes(format, size, image).save(buffer, format="JPEG")
|
||||
encoded_image = base64.b64encode(buffer.getvalue()).decode("utf-8")
|
||||
@@ -306,6 +390,14 @@ class AWSBedrockLLMContext(OpenAILLMContext):
|
||||
self.add_message({"role": "user", "content": content})
|
||||
|
||||
def add_message(self, message):
|
||||
"""Add a message to the context, merging with previous message if same role.
|
||||
|
||||
AWS Bedrock requires alternating roles, so consecutive messages from the
|
||||
same role are merged together.
|
||||
|
||||
Args:
|
||||
message: The message to add to the context.
|
||||
"""
|
||||
try:
|
||||
if self.messages:
|
||||
# AWS Bedrock requires that roles alternate. If this message's
|
||||
@@ -330,10 +422,10 @@ class AWSBedrockLLMContext(OpenAILLMContext):
|
||||
logger.error(f"Error adding message: {e}")
|
||||
|
||||
def _restructure_from_bedrock_messages(self):
|
||||
"""Restructure messages in AWS Bedrock format by handling system
|
||||
messages, merging consecutive messages with the same role, and ensuring
|
||||
proper content formatting.
|
||||
"""Restructure messages in AWS Bedrock format.
|
||||
|
||||
Handles system messages, merging consecutive messages with the same role,
|
||||
and ensuring proper content formatting.
|
||||
"""
|
||||
# Handle system message if present at the beginning
|
||||
if self.messages and self.messages[0]["role"] == "system":
|
||||
@@ -416,12 +508,22 @@ class AWSBedrockLLMContext(OpenAILLMContext):
|
||||
message["content"] = [{"type": "text", "text": "(empty)"}]
|
||||
|
||||
def get_messages_for_persistent_storage(self):
|
||||
"""Get messages formatted for persistent storage.
|
||||
|
||||
Returns:
|
||||
List of messages including system message if present.
|
||||
"""
|
||||
messages = super().get_messages_for_persistent_storage()
|
||||
if self.system:
|
||||
messages.insert(0, {"role": "system", "content": self.system})
|
||||
return messages
|
||||
|
||||
def get_messages_for_logging(self) -> str:
|
||||
"""Get messages formatted for logging with sensitive data redacted.
|
||||
|
||||
Returns:
|
||||
JSON string representation of messages with image data redacted.
|
||||
"""
|
||||
msgs = []
|
||||
for message in self.messages:
|
||||
msg = copy.deepcopy(message)
|
||||
@@ -435,11 +537,36 @@ class AWSBedrockLLMContext(OpenAILLMContext):
|
||||
|
||||
|
||||
class AWSBedrockUserContextAggregator(LLMUserContextAggregator):
|
||||
"""User context aggregator for AWS Bedrock LLM service.
|
||||
|
||||
Handles aggregation of user messages and frames for AWS Bedrock format.
|
||||
Inherits all functionality from the base LLM user context aggregator.
|
||||
|
||||
Args:
|
||||
context: The LLM context to aggregate messages into.
|
||||
params: Configuration parameters for the aggregator.
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class AWSBedrockAssistantContextAggregator(LLMAssistantContextAggregator):
|
||||
"""Assistant context aggregator for AWS Bedrock LLM service.
|
||||
|
||||
Handles aggregation of assistant responses and function calls for AWS Bedrock
|
||||
format, including tool use and tool result handling.
|
||||
|
||||
Args:
|
||||
context: The LLM context to aggregate messages into.
|
||||
params: Configuration parameters for the aggregator.
|
||||
"""
|
||||
|
||||
async def handle_function_call_in_progress(self, frame: FunctionCallInProgressFrame):
|
||||
"""Handle function call in progress frame.
|
||||
|
||||
Args:
|
||||
frame: The function call in progress frame to handle.
|
||||
"""
|
||||
# Format tool use according to AWS Bedrock API
|
||||
self._context.add_message(
|
||||
{
|
||||
@@ -470,6 +597,11 @@ class AWSBedrockAssistantContextAggregator(LLMAssistantContextAggregator):
|
||||
)
|
||||
|
||||
async def handle_function_call_result(self, frame: FunctionCallResultFrame):
|
||||
"""Handle function call result frame.
|
||||
|
||||
Args:
|
||||
frame: The function call result frame to handle.
|
||||
"""
|
||||
if frame.result:
|
||||
result = json.dumps(frame.result)
|
||||
await self._update_function_call_result(frame.function_name, frame.tool_call_id, result)
|
||||
@@ -479,6 +611,11 @@ class AWSBedrockAssistantContextAggregator(LLMAssistantContextAggregator):
|
||||
)
|
||||
|
||||
async def handle_function_call_cancel(self, frame: FunctionCallCancelFrame):
|
||||
"""Handle function call cancel frame.
|
||||
|
||||
Args:
|
||||
frame: The function call cancel frame to handle.
|
||||
"""
|
||||
await self._update_function_call_result(
|
||||
frame.function_name, frame.tool_call_id, "CANCELLED"
|
||||
)
|
||||
@@ -497,6 +634,11 @@ class AWSBedrockAssistantContextAggregator(LLMAssistantContextAggregator):
|
||||
content["toolResult"]["content"] = [{"text": result}]
|
||||
|
||||
async def handle_user_image_frame(self, frame: UserImageRawFrame):
|
||||
"""Handle user image frame.
|
||||
|
||||
Args:
|
||||
frame: The user image frame to handle.
|
||||
"""
|
||||
await self._update_function_call_result(
|
||||
frame.request.function_name, frame.request.tool_call_id, "COMPLETED"
|
||||
)
|
||||
@@ -509,18 +651,38 @@ class AWSBedrockAssistantContextAggregator(LLMAssistantContextAggregator):
|
||||
|
||||
|
||||
class AWSBedrockLLMService(LLMService):
|
||||
"""This class implements inference with AWS Bedrock models including Amazon
|
||||
Nova and Anthropic Claude.
|
||||
"""AWS Bedrock Large Language Model service implementation.
|
||||
|
||||
Requires AWS credentials to be configured in the environment or through
|
||||
boto3 configuration.
|
||||
Provides inference capabilities for AWS Bedrock models including Amazon Nova
|
||||
and Anthropic Claude. Supports streaming responses, function calling, and
|
||||
vision capabilities.
|
||||
|
||||
Args:
|
||||
model: The AWS Bedrock model identifier to use.
|
||||
aws_access_key: AWS access key ID. If None, uses default credentials.
|
||||
aws_secret_key: AWS secret access key. If None, uses default credentials.
|
||||
aws_session_token: AWS session token for temporary credentials.
|
||||
aws_region: AWS region for the Bedrock service.
|
||||
params: Model parameters and configuration.
|
||||
client_config: Custom boto3 client configuration.
|
||||
**kwargs: Additional arguments passed to parent LLMService.
|
||||
"""
|
||||
|
||||
# Overriding the default adapter to use the Anthropic one.
|
||||
adapter_class = AWSBedrockLLMAdapter
|
||||
|
||||
class InputParams(BaseModel):
|
||||
"""Input parameters for AWS Bedrock LLM service.
|
||||
|
||||
Parameters:
|
||||
max_tokens: Maximum number of tokens to generate.
|
||||
temperature: Sampling temperature between 0.0 and 1.0.
|
||||
top_p: Nucleus sampling parameter between 0.0 and 1.0.
|
||||
stop_sequences: List of strings that stop generation.
|
||||
latency: Performance mode - "standard" or "optimized".
|
||||
additional_model_request_fields: Additional model-specific parameters.
|
||||
"""
|
||||
|
||||
max_tokens: Optional[int] = Field(default_factory=lambda: 4096, ge=1)
|
||||
temperature: Optional[float] = Field(default_factory=lambda: 0.7, ge=0.0, le=1.0)
|
||||
top_p: Optional[float] = Field(default_factory=lambda: 0.999, ge=0.0, le=1.0)
|
||||
@@ -573,6 +735,11 @@ class AWSBedrockLLMService(LLMService):
|
||||
logger.info(f"Using AWS Bedrock model: {model}")
|
||||
|
||||
def can_generate_metrics(self) -> bool:
|
||||
"""Check if the service can generate usage metrics.
|
||||
|
||||
Returns:
|
||||
True if metrics generation is supported.
|
||||
"""
|
||||
return True
|
||||
|
||||
def create_context_aggregator(
|
||||
@@ -582,21 +749,21 @@ class AWSBedrockLLMService(LLMService):
|
||||
user_params: LLMUserAggregatorParams = LLMUserAggregatorParams(),
|
||||
assistant_params: LLMAssistantAggregatorParams = LLMAssistantAggregatorParams(),
|
||||
) -> AWSBedrockContextAggregatorPair:
|
||||
"""Create an instance of AWSBedrockContextAggregatorPair from an
|
||||
OpenAILLMContext. Constructor keyword arguments for both the user and
|
||||
assistant aggregators can be provided.
|
||||
"""Create AWS Bedrock-specific context aggregators.
|
||||
|
||||
Creates a pair of context aggregators optimized for AWS Bedrocks's message
|
||||
format, including support for function calls, tool usage, and image handling.
|
||||
|
||||
Args:
|
||||
context (OpenAILLMContext): The LLM context.
|
||||
user_params (LLMUserAggregatorParams, optional): User aggregator
|
||||
parameters.
|
||||
assistant_params (LLMAssistantAggregatorParams, optional): User
|
||||
aggregator parameters.
|
||||
context: The LLM context to create aggregators for.
|
||||
user_params: Parameters for user message aggregation.
|
||||
assistant_params: Parameters for assistant message aggregation.
|
||||
|
||||
Returns:
|
||||
AWSBedrockContextAggregatorPair: A pair of context aggregators, one
|
||||
for the user and one for the assistant, encapsulated in an
|
||||
AWSBedrockContextAggregatorPair: A pair of context aggregators, one for
|
||||
the user and one for the assistant, encapsulated in an
|
||||
AWSBedrockContextAggregatorPair.
|
||||
|
||||
"""
|
||||
context.set_llm_adapter(self.get_llm_adapter())
|
||||
|
||||
@@ -792,6 +959,12 @@ class AWSBedrockLLMService(LLMService):
|
||||
)
|
||||
|
||||
async def process_frame(self, frame: Frame, direction: FrameDirection):
|
||||
"""Process incoming frames and handle LLM-specific frame types.
|
||||
|
||||
Args:
|
||||
frame: The frame to process.
|
||||
direction: The direction of frame processing.
|
||||
"""
|
||||
await super().process_frame(frame, direction)
|
||||
|
||||
context = None
|
||||
|
||||
Reference in New Issue
Block a user