mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: Introducing support for Bedrock-hosted models (Anthropic, Cohere, etc.) (#2610)
* Pushing the bedrock related changes to the new branch after addressing the review comments * 2524 Addressed the second round review comments * 2524 Addressed few more minor comments on the PR * resolving the merge conflict * 2524 resolved the uv.lock conflicts * 2524 addressed more comments * 2524 removed the print statement to fix the checks failure * 2524 resolved the CI failure issues * 2524 fixing the CI breaks * 2524 Addressed the review comment * 2524 resolved conflict --------- Co-authored-by: Sunil Dutta <sunil.dutta@penske.com> Co-authored-by: budgetboardingai <apurva.sharma31@gmail.com>
This commit is contained in:
committed by
GitHub
Unverified
parent
defe0f1a89
commit
3b77192ad0
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) Microsoft Corporation.
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE
|
||||
@@ -0,0 +1,19 @@
|
||||
# Get Started with Microsoft Agent Framework Bedrock
|
||||
|
||||
Install the provider package:
|
||||
|
||||
```bash
|
||||
pip install agent-framework-bedrock --pre
|
||||
```
|
||||
|
||||
## Bedrock Integration
|
||||
|
||||
The Bedrock integration enables Microsoft Agent Framework applications to call Amazon Bedrock models with familiar chat abstractions, including tool/function calling when you attach tools through `ChatOptions`.
|
||||
|
||||
### Basic Usage Example
|
||||
|
||||
See the [Bedrock sample script](samples/bedrock_sample.py) for a runnable end-to-end script that:
|
||||
|
||||
- Loads credentials from the `BEDROCK_*` environment variables
|
||||
- Instantiates `BedrockChatClient`
|
||||
- Sends a simple conversation turn and prints the response
|
||||
@@ -0,0 +1,15 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import importlib.metadata
|
||||
|
||||
from ._chat_client import BedrockChatClient
|
||||
|
||||
try:
|
||||
__version__ = importlib.metadata.version(__name__)
|
||||
except importlib.metadata.PackageNotFoundError:
|
||||
__version__ = "0.0.0"
|
||||
|
||||
__all__ = [
|
||||
"BedrockChatClient",
|
||||
"__version__",
|
||||
]
|
||||
@@ -0,0 +1,527 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from collections import deque
|
||||
from collections.abc import AsyncIterable, MutableMapping, MutableSequence, Sequence
|
||||
from typing import Any, ClassVar
|
||||
from uuid import uuid4
|
||||
|
||||
from agent_framework import (
|
||||
AGENT_FRAMEWORK_USER_AGENT,
|
||||
AIFunction,
|
||||
BaseChatClient,
|
||||
ChatMessage,
|
||||
ChatOptions,
|
||||
ChatResponse,
|
||||
ChatResponseUpdate,
|
||||
Contents,
|
||||
FinishReason,
|
||||
FunctionCallContent,
|
||||
FunctionResultContent,
|
||||
Role,
|
||||
TextContent,
|
||||
ToolProtocol,
|
||||
UsageContent,
|
||||
UsageDetails,
|
||||
get_logger,
|
||||
prepare_function_call_results,
|
||||
use_chat_middleware,
|
||||
use_function_invocation,
|
||||
)
|
||||
from agent_framework._pydantic import AFBaseSettings
|
||||
from agent_framework.exceptions import ServiceInitializationError, ServiceInvalidResponseError
|
||||
from agent_framework.observability import use_instrumentation
|
||||
from boto3.session import Session as Boto3Session
|
||||
from botocore.client import BaseClient
|
||||
from botocore.config import Config as BotoConfig
|
||||
from pydantic import SecretStr, ValidationError
|
||||
|
||||
logger = get_logger("agent_framework.bedrock")
|
||||
|
||||
DEFAULT_REGION = "us-east-1"
|
||||
DEFAULT_MAX_TOKENS = 1024
|
||||
|
||||
ROLE_MAP: dict[Role, str] = {
|
||||
Role.USER: "user",
|
||||
Role.ASSISTANT: "assistant",
|
||||
Role.SYSTEM: "user",
|
||||
Role.TOOL: "user",
|
||||
}
|
||||
|
||||
FINISH_REASON_MAP: dict[str, FinishReason] = {
|
||||
"end_turn": FinishReason.STOP,
|
||||
"stop_sequence": FinishReason.STOP,
|
||||
"max_tokens": FinishReason.LENGTH,
|
||||
"length": FinishReason.LENGTH,
|
||||
"content_filtered": FinishReason.CONTENT_FILTER,
|
||||
"tool_use": FinishReason.TOOL_CALLS,
|
||||
}
|
||||
|
||||
|
||||
class BedrockSettings(AFBaseSettings):
|
||||
"""Bedrock configuration settings pulled from environment variables or .env files."""
|
||||
|
||||
env_prefix: ClassVar[str] = "BEDROCK_"
|
||||
|
||||
region: str = DEFAULT_REGION
|
||||
chat_model_id: str | None = None
|
||||
access_key: SecretStr | None = None
|
||||
secret_key: SecretStr | None = None
|
||||
session_token: SecretStr | None = None
|
||||
|
||||
|
||||
@use_function_invocation
|
||||
@use_instrumentation
|
||||
@use_chat_middleware
|
||||
class BedrockChatClient(BaseChatClient):
|
||||
"""Async chat client for Amazon Bedrock's Converse API."""
|
||||
|
||||
OTEL_PROVIDER_NAME: ClassVar[str] = "aws.bedrock" # type: ignore[reportIncompatibleVariableOverride, misc]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
region: str | None = None,
|
||||
model_id: str | None = None,
|
||||
access_key: str | None = None,
|
||||
secret_key: str | None = None,
|
||||
session_token: str | None = None,
|
||||
client: BaseClient | None = None,
|
||||
boto3_session: Boto3Session | None = None,
|
||||
env_file_path: str | None = None,
|
||||
env_file_encoding: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Create a Bedrock chat client and load AWS credentials.
|
||||
|
||||
Args:
|
||||
region: Region to send Bedrock requests to; falls back to BEDROCK_REGION.
|
||||
model_id: Default model identifier; falls back to BEDROCK_CHAT_MODEL_ID.
|
||||
access_key: Optional AWS access key for manual credential injection.
|
||||
secret_key: Optional AWS secret key paired with ``access_key``.
|
||||
session_token: Optional AWS session token for temporary credentials.
|
||||
client: Preconfigured Bedrock runtime client; when omitted a boto3 session is created.
|
||||
boto3_session: Custom boto3 session used to build the runtime client if provided.
|
||||
env_file_path: Optional .env file path used by ``BedrockSettings`` to load defaults.
|
||||
env_file_encoding: Encoding for the optional .env file.
|
||||
kwargs: Additional arguments forwarded to ``BaseChatClient``.
|
||||
"""
|
||||
try:
|
||||
settings = BedrockSettings(
|
||||
region=region,
|
||||
chat_model_id=model_id,
|
||||
access_key=access_key, # type: ignore[arg-type]
|
||||
secret_key=secret_key, # type: ignore[arg-type]
|
||||
session_token=session_token, # type: ignore[arg-type]
|
||||
env_file_path=env_file_path,
|
||||
env_file_encoding=env_file_encoding,
|
||||
)
|
||||
except ValidationError as ex:
|
||||
raise ServiceInitializationError("Failed to initialize Bedrock settings.", ex) from ex
|
||||
|
||||
if client is None:
|
||||
session = boto3_session or self._create_session(settings)
|
||||
client = session.client(
|
||||
"bedrock-runtime",
|
||||
region_name=settings.region,
|
||||
config=BotoConfig(user_agent_extra=AGENT_FRAMEWORK_USER_AGENT),
|
||||
)
|
||||
|
||||
super().__init__(**kwargs)
|
||||
self._bedrock_client = client
|
||||
self.model_id = settings.chat_model_id
|
||||
self.region = settings.region
|
||||
|
||||
@staticmethod
|
||||
def _create_session(settings: BedrockSettings) -> Boto3Session:
|
||||
session_kwargs: dict[str, Any] = {"region_name": settings.region or DEFAULT_REGION}
|
||||
if settings.access_key and settings.secret_key:
|
||||
session_kwargs["aws_access_key_id"] = settings.access_key.get_secret_value()
|
||||
session_kwargs["aws_secret_access_key"] = settings.secret_key.get_secret_value()
|
||||
if settings.session_token:
|
||||
session_kwargs["aws_session_token"] = settings.session_token.get_secret_value()
|
||||
return Boto3Session(**session_kwargs)
|
||||
|
||||
async def _inner_get_response(
|
||||
self,
|
||||
*,
|
||||
messages: MutableSequence[ChatMessage],
|
||||
chat_options: ChatOptions,
|
||||
**kwargs: Any,
|
||||
) -> ChatResponse:
|
||||
request = self._build_converse_request(messages, chat_options, **kwargs)
|
||||
raw_response = await asyncio.to_thread(self._bedrock_client.converse, **request)
|
||||
return self._process_converse_response(raw_response)
|
||||
|
||||
async def _inner_get_streaming_response(
|
||||
self,
|
||||
*,
|
||||
messages: MutableSequence[ChatMessage],
|
||||
chat_options: ChatOptions,
|
||||
**kwargs: Any,
|
||||
) -> AsyncIterable[ChatResponseUpdate]:
|
||||
response = await self._inner_get_response(messages=messages, chat_options=chat_options, **kwargs)
|
||||
contents = list(response.messages[0].contents if response.messages else [])
|
||||
if response.usage_details:
|
||||
contents.append(UsageContent(details=response.usage_details))
|
||||
yield ChatResponseUpdate(
|
||||
response_id=response.response_id,
|
||||
contents=contents,
|
||||
model_id=response.model_id,
|
||||
finish_reason=response.finish_reason,
|
||||
raw_representation=response.raw_representation,
|
||||
)
|
||||
|
||||
def _build_converse_request(
|
||||
self,
|
||||
messages: MutableSequence[ChatMessage],
|
||||
chat_options: ChatOptions,
|
||||
**kwargs: Any,
|
||||
) -> dict[str, Any]:
|
||||
model_id = chat_options.model_id or self.model_id
|
||||
if not model_id:
|
||||
raise ServiceInitializationError(
|
||||
"Bedrock model_id is required. Set via chat options or BEDROCK_CHAT_MODEL_ID environment variable."
|
||||
)
|
||||
|
||||
system_prompts, conversation = self._prepare_bedrock_messages(messages)
|
||||
if not conversation:
|
||||
raise ServiceInitializationError("At least one non-system message is required for Bedrock requests.")
|
||||
|
||||
payload: dict[str, Any] = {
|
||||
"modelId": model_id,
|
||||
"messages": conversation,
|
||||
}
|
||||
if system_prompts:
|
||||
payload["system"] = system_prompts
|
||||
|
||||
inference_config: dict[str, Any] = {}
|
||||
inference_config["maxTokens"] = (
|
||||
chat_options.max_tokens if chat_options.max_tokens is not None else DEFAULT_MAX_TOKENS
|
||||
)
|
||||
if chat_options.temperature is not None:
|
||||
inference_config["temperature"] = chat_options.temperature
|
||||
if chat_options.top_p is not None:
|
||||
inference_config["topP"] = chat_options.top_p
|
||||
if chat_options.stop is not None:
|
||||
inference_config["stopSequences"] = chat_options.stop
|
||||
if inference_config:
|
||||
payload["inferenceConfig"] = inference_config
|
||||
|
||||
tool_config = self._convert_tools_to_bedrock_config(chat_options.tools)
|
||||
if tool_choice := self._convert_tool_choice(chat_options.tool_choice):
|
||||
if tool_config is None:
|
||||
tool_config = {}
|
||||
tool_config["toolChoice"] = tool_choice
|
||||
if tool_config:
|
||||
payload["toolConfig"] = tool_config
|
||||
|
||||
if chat_options.additional_properties:
|
||||
payload.update(chat_options.additional_properties)
|
||||
if kwargs:
|
||||
payload.update(kwargs)
|
||||
return payload
|
||||
|
||||
def _prepare_bedrock_messages(
|
||||
self, messages: Sequence[ChatMessage]
|
||||
) -> tuple[list[dict[str, str]], list[dict[str, Any]]]:
|
||||
prompts: list[dict[str, str]] = []
|
||||
conversation: list[dict[str, Any]] = []
|
||||
pending_tool_use_ids: deque[str] = deque()
|
||||
for message in messages:
|
||||
if message.role == Role.SYSTEM:
|
||||
text_value = message.text
|
||||
if text_value:
|
||||
prompts.append({"text": text_value})
|
||||
continue
|
||||
|
||||
content_blocks = self._convert_message_to_content_blocks(message)
|
||||
if not content_blocks:
|
||||
continue
|
||||
|
||||
role = ROLE_MAP.get(message.role, "user")
|
||||
if role == "assistant":
|
||||
pending_tool_use_ids = deque(
|
||||
block["toolUse"]["toolUseId"]
|
||||
for block in content_blocks
|
||||
if isinstance(block, MutableMapping) and "toolUse" in block
|
||||
)
|
||||
elif message.role == Role.TOOL:
|
||||
content_blocks = self._align_tool_results_with_pending(content_blocks, pending_tool_use_ids)
|
||||
pending_tool_use_ids.clear()
|
||||
if not content_blocks:
|
||||
continue
|
||||
else:
|
||||
pending_tool_use_ids.clear()
|
||||
|
||||
conversation.append({"role": role, "content": content_blocks})
|
||||
|
||||
return prompts, conversation
|
||||
|
||||
def _align_tool_results_with_pending(
|
||||
self, content_blocks: list[dict[str, Any]], pending_tool_use_ids: deque[str]
|
||||
) -> list[dict[str, Any]]:
|
||||
if not content_blocks:
|
||||
return content_blocks
|
||||
if not pending_tool_use_ids:
|
||||
# No pending tool calls; drop toolResult blocks to avoid Bedrock validation errors
|
||||
return [
|
||||
block for block in content_blocks if not (isinstance(block, MutableMapping) and "toolResult" in block)
|
||||
]
|
||||
|
||||
aligned_blocks: list[dict[str, Any]] = []
|
||||
pending = deque(pending_tool_use_ids)
|
||||
for block in content_blocks:
|
||||
if not isinstance(block, MutableMapping):
|
||||
aligned_blocks.append(block)
|
||||
continue
|
||||
tool_result = block.get("toolResult")
|
||||
if not tool_result:
|
||||
aligned_blocks.append(block)
|
||||
continue
|
||||
if not pending:
|
||||
logger.debug("Dropping extra tool result block due to missing pending tool uses: %s", block)
|
||||
continue
|
||||
tool_use_id = tool_result.get("toolUseId")
|
||||
if tool_use_id:
|
||||
try:
|
||||
pending.remove(tool_use_id)
|
||||
except ValueError:
|
||||
logger.debug("Tool result references unknown toolUseId '%s'. Dropping block.", tool_use_id)
|
||||
continue
|
||||
else:
|
||||
tool_result["toolUseId"] = pending.popleft()
|
||||
aligned_blocks.append(block)
|
||||
|
||||
return aligned_blocks
|
||||
|
||||
def _convert_message_to_content_blocks(self, message: ChatMessage) -> list[dict[str, Any]]:
|
||||
blocks: list[dict[str, Any]] = []
|
||||
for content in message.contents:
|
||||
block = self._convert_content_to_bedrock_block(content)
|
||||
if block is None:
|
||||
logger.debug("Skipping unsupported content type for Bedrock: %s", type(content))
|
||||
continue
|
||||
blocks.append(block)
|
||||
return blocks
|
||||
|
||||
def _convert_content_to_bedrock_block(self, content: Contents) -> dict[str, Any] | None:
|
||||
if isinstance(content, TextContent):
|
||||
return {"text": content.text}
|
||||
if isinstance(content, FunctionCallContent):
|
||||
arguments = content.parse_arguments() or {}
|
||||
return {
|
||||
"toolUse": {
|
||||
"toolUseId": content.call_id or self._generate_tool_call_id(),
|
||||
"name": content.name,
|
||||
"input": arguments,
|
||||
}
|
||||
}
|
||||
if isinstance(content, FunctionResultContent):
|
||||
tool_result_block = {
|
||||
"toolResult": {
|
||||
"toolUseId": content.call_id,
|
||||
"content": self._convert_tool_result_to_blocks(content.result),
|
||||
"status": "error" if content.exception else "success",
|
||||
}
|
||||
}
|
||||
if content.exception:
|
||||
tool_result = tool_result_block["toolResult"]
|
||||
existing_content = tool_result.get("content")
|
||||
content_list: list[dict[str, Any]]
|
||||
if isinstance(existing_content, list):
|
||||
content_list = existing_content
|
||||
else:
|
||||
content_list = []
|
||||
tool_result["content"] = content_list
|
||||
content_list.append({"text": str(content.exception)})
|
||||
return tool_result_block
|
||||
return None
|
||||
|
||||
def _convert_tool_result_to_blocks(self, result: Any) -> list[dict[str, Any]]:
|
||||
prepared_result = prepare_function_call_results(result)
|
||||
try:
|
||||
parsed_result = json.loads(prepared_result)
|
||||
except json.JSONDecodeError:
|
||||
return [{"text": prepared_result}]
|
||||
|
||||
return self._convert_prepared_tool_result_to_blocks(parsed_result)
|
||||
|
||||
def _convert_prepared_tool_result_to_blocks(self, value: Any) -> list[dict[str, Any]]:
|
||||
if isinstance(value, list):
|
||||
blocks: list[dict[str, Any]] = []
|
||||
for item in value:
|
||||
blocks.extend(self._convert_prepared_tool_result_to_blocks(item))
|
||||
return blocks or [{"text": ""}]
|
||||
return [self._normalize_tool_result_value(value)]
|
||||
|
||||
def _normalize_tool_result_value(self, value: Any) -> dict[str, Any]:
|
||||
if isinstance(value, dict):
|
||||
return {"json": value}
|
||||
if isinstance(value, (list, tuple)):
|
||||
return {"json": list(value)}
|
||||
if isinstance(value, str):
|
||||
return {"text": value}
|
||||
if isinstance(value, (int, float, bool)) or value is None:
|
||||
return {"json": value}
|
||||
if isinstance(value, TextContent) and getattr(value, "text", None):
|
||||
return {"text": value.text}
|
||||
if hasattr(value, "to_dict"):
|
||||
try:
|
||||
return {"json": value.to_dict()} # type: ignore[call-arg]
|
||||
except Exception: # pragma: no cover - defensive
|
||||
return {"text": str(value)}
|
||||
return {"text": str(value)}
|
||||
|
||||
def _convert_tools_to_bedrock_config(
|
||||
self, tools: list[ToolProtocol | MutableMapping[str, Any]] | None
|
||||
) -> dict[str, Any] | None:
|
||||
if not tools:
|
||||
return None
|
||||
converted: list[dict[str, Any]] = []
|
||||
for tool in tools:
|
||||
if isinstance(tool, MutableMapping):
|
||||
converted.append(dict(tool))
|
||||
continue
|
||||
if isinstance(tool, AIFunction):
|
||||
converted.append({
|
||||
"toolSpec": {
|
||||
"name": tool.name,
|
||||
"description": tool.description or "",
|
||||
"inputSchema": {"json": tool.parameters()},
|
||||
}
|
||||
})
|
||||
continue
|
||||
logger.debug("Ignoring unsupported tool type for Bedrock: %s", type(tool))
|
||||
return {"tools": converted} if converted else None
|
||||
|
||||
def _convert_tool_choice(self, tool_choice: Any) -> dict[str, Any] | None:
|
||||
if not tool_choice:
|
||||
return None
|
||||
mode = tool_choice.mode if hasattr(tool_choice, "mode") else str(tool_choice)
|
||||
required_name = getattr(tool_choice, "required_function_name", None)
|
||||
match mode:
|
||||
case "auto":
|
||||
return {"auto": {}}
|
||||
case "none":
|
||||
return {"none": {}}
|
||||
case "required":
|
||||
if required_name:
|
||||
return {"tool": {"name": required_name}}
|
||||
return {"any": {}}
|
||||
case _:
|
||||
logger.debug("Unsupported tool choice mode for Bedrock: %s", mode)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _generate_tool_call_id() -> str:
|
||||
return f"tool-call-{uuid4().hex}"
|
||||
|
||||
def _process_converse_response(self, response: dict[str, Any]) -> ChatResponse:
|
||||
output = response.get("output", {})
|
||||
message = output.get("message", {})
|
||||
content_blocks = message.get("content", []) or []
|
||||
contents = self._parse_message_contents(content_blocks)
|
||||
chat_message = ChatMessage(role=Role.ASSISTANT, contents=contents, raw_representation=message)
|
||||
usage_details = self._parse_usage(response.get("usage") or output.get("usage"))
|
||||
finish_reason = self._map_finish_reason(output.get("completionReason") or response.get("stopReason"))
|
||||
response_id = response.get("responseId") or message.get("id")
|
||||
model_id = response.get("modelId") or output.get("modelId") or self.model_id
|
||||
return ChatResponse(
|
||||
response_id=response_id,
|
||||
messages=[chat_message],
|
||||
usage_details=usage_details,
|
||||
model_id=model_id,
|
||||
finish_reason=finish_reason,
|
||||
raw_representation=response,
|
||||
)
|
||||
|
||||
def _parse_usage(self, usage: dict[str, Any] | None) -> UsageDetails | None:
|
||||
if not usage:
|
||||
return None
|
||||
details = UsageDetails()
|
||||
if (input_tokens := usage.get("inputTokens")) is not None:
|
||||
details.input_token_count = input_tokens
|
||||
if (output_tokens := usage.get("outputTokens")) is not None:
|
||||
details.output_token_count = output_tokens
|
||||
if (total_tokens := usage.get("totalTokens")) is not None:
|
||||
details.additional_counts["bedrock.total_tokens"] = total_tokens
|
||||
return details
|
||||
|
||||
def _parse_message_contents(self, content_blocks: Sequence[MutableMapping[str, Any]]) -> list[Any]:
|
||||
contents: list[Any] = []
|
||||
for block in content_blocks:
|
||||
if text_value := block.get("text"):
|
||||
contents.append(TextContent(text=text_value, raw_representation=block))
|
||||
continue
|
||||
if (json_value := block.get("json")) is not None:
|
||||
contents.append(TextContent(text=json.dumps(json_value), raw_representation=block))
|
||||
continue
|
||||
tool_use = block.get("toolUse")
|
||||
if isinstance(tool_use, MutableMapping):
|
||||
tool_name = tool_use.get("name")
|
||||
if not tool_name:
|
||||
raise ServiceInvalidResponseError("Bedrock response missing required tool name in toolUse block.")
|
||||
contents.append(
|
||||
FunctionCallContent(
|
||||
call_id=tool_use.get("toolUseId") or self._generate_tool_call_id(),
|
||||
name=tool_name,
|
||||
arguments=tool_use.get("input"),
|
||||
raw_representation=block,
|
||||
)
|
||||
)
|
||||
continue
|
||||
tool_result = block.get("toolResult")
|
||||
if isinstance(tool_result, MutableMapping):
|
||||
status = (tool_result.get("status") or "success").lower()
|
||||
exception = None
|
||||
if status not in {"success", "ok"}:
|
||||
exception = RuntimeError(f"Bedrock tool result status: {status}")
|
||||
result_value = self._convert_bedrock_tool_result_to_value(tool_result.get("content"))
|
||||
contents.append(
|
||||
FunctionResultContent(
|
||||
call_id=tool_result.get("toolUseId") or self._generate_tool_call_id(),
|
||||
result=result_value,
|
||||
exception=exception,
|
||||
raw_representation=block,
|
||||
)
|
||||
)
|
||||
continue
|
||||
logger.debug("Ignoring unsupported Bedrock content block: %s", block)
|
||||
return contents
|
||||
|
||||
def _map_finish_reason(self, reason: str | None) -> FinishReason | None:
|
||||
if not reason:
|
||||
return None
|
||||
return FINISH_REASON_MAP.get(reason.lower())
|
||||
|
||||
def service_url(self) -> str:
|
||||
"""Returns the service URL for the Bedrock runtime in the configured AWS region.
|
||||
|
||||
Returns:
|
||||
str: The Bedrock runtime service URL.
|
||||
"""
|
||||
return f"https://bedrock-runtime.{self.region}.amazonaws.com"
|
||||
|
||||
def _convert_bedrock_tool_result_to_value(self, content: Any) -> Any:
|
||||
if not content:
|
||||
return None
|
||||
if isinstance(content, Sequence) and not isinstance(content, (str, bytes, bytearray)):
|
||||
values: list[Any] = []
|
||||
for item in content:
|
||||
if isinstance(item, MutableMapping):
|
||||
if (text_value := item.get("text")) is not None:
|
||||
values.append(text_value)
|
||||
continue
|
||||
if "json" in item:
|
||||
values.append(item["json"])
|
||||
continue
|
||||
values.append(item)
|
||||
return values[0] if len(values) == 1 else values
|
||||
if isinstance(content, MutableMapping):
|
||||
if (text_value := content.get("text")) is not None:
|
||||
return text_value
|
||||
if "json" in content:
|
||||
return content["json"]
|
||||
return content
|
||||
@@ -0,0 +1,90 @@
|
||||
[project]
|
||||
name = "agent-framework-bedrock"
|
||||
description = "Amazon Bedrock integration for Microsoft Agent Framework."
|
||||
authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
version = "1.0.0b251120"
|
||||
license-files = ["LICENSE"]
|
||||
urls.homepage = "https://aka.ms/agent-framework"
|
||||
urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
|
||||
urls.release_notes = "https://github.com/microsoft/agent-framework/releases?q=tag%3Apython-1&expanded=true"
|
||||
urls.issues = "https://github.com/microsoft/agent-framework/issues"
|
||||
classifiers = [
|
||||
"License :: OSI Approved :: MIT License",
|
||||
"Development Status :: 4 - Beta",
|
||||
"Intended Audience :: Developers",
|
||||
"Programming Language :: Python :: 3",
|
||||
"Programming Language :: Python :: 3.10",
|
||||
"Programming Language :: Python :: 3.11",
|
||||
"Programming Language :: Python :: 3.12",
|
||||
"Programming Language :: Python :: 3.13",
|
||||
"Programming Language :: Python :: 3.14",
|
||||
"Typing :: Typed",
|
||||
]
|
||||
dependencies = [
|
||||
"agent-framework-core",
|
||||
"boto3>=1.35.0,<2.0.0",
|
||||
"botocore>=1.35.0,<2.0.0",
|
||||
]
|
||||
|
||||
|
||||
[tool.uv]
|
||||
prerelease = "if-necessary-or-explicit"
|
||||
environments = [
|
||||
"sys_platform == 'darwin'",
|
||||
"sys_platform == 'linux'",
|
||||
"sys_platform == 'win32'"
|
||||
]
|
||||
|
||||
[tool.uv-dynamic-versioning]
|
||||
fallback-version = "0.0.0"
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
testpaths = 'tests'
|
||||
addopts = "-ra -q -r fEX"
|
||||
asyncio_mode = "auto"
|
||||
asyncio_default_fixture_loop_scope = "function"
|
||||
filterwarnings = []
|
||||
timeout = 120
|
||||
|
||||
[tool.ruff]
|
||||
extend = "../../pyproject.toml"
|
||||
|
||||
[tool.coverage.run]
|
||||
omit = [
|
||||
"**/__init__.py"
|
||||
]
|
||||
|
||||
[tool.pyright]
|
||||
extends = "../../pyproject.toml"
|
||||
|
||||
[tool.mypy]
|
||||
plugins = ['pydantic.mypy']
|
||||
strict = true
|
||||
python_version = "3.10"
|
||||
ignore_missing_imports = true
|
||||
disallow_untyped_defs = true
|
||||
no_implicit_optional = true
|
||||
check_untyped_defs = true
|
||||
warn_return_any = true
|
||||
show_error_codes = true
|
||||
warn_unused_ignores = false
|
||||
disallow_incomplete_defs = true
|
||||
disallow_untyped_decorators = true
|
||||
|
||||
[tool.bandit]
|
||||
targets = ["agent_framework_bedrock"]
|
||||
exclude_dirs = ["tests"]
|
||||
|
||||
[tool.poe]
|
||||
executor.type = "uv"
|
||||
include = "../../shared_tasks.toml"
|
||||
|
||||
[tool.poe.tasks]
|
||||
mypy = "mypy --config-file $POE_ROOT/pyproject.toml agent_framework_bedrock"
|
||||
test = "pytest --cov=agent_framework_bedrock --cov-report=term-missing:skip-covered tests"
|
||||
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
build-backend = "hatchling.build"
|
||||
@@ -0,0 +1,64 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from collections.abc import Sequence
|
||||
|
||||
from agent_framework import (
|
||||
AgentRunResponse,
|
||||
ChatAgent,
|
||||
FunctionCallContent,
|
||||
FunctionResultContent,
|
||||
Role,
|
||||
TextContent,
|
||||
ToolMode,
|
||||
ai_function,
|
||||
)
|
||||
|
||||
from agent_framework_bedrock import BedrockChatClient
|
||||
|
||||
|
||||
@ai_function
|
||||
def get_weather(city: str) -> dict[str, str]:
|
||||
"""Return a mock forecast for the requested city."""
|
||||
normalized = city.strip() or "New York"
|
||||
return {"city": normalized, "forecast": "72F and sunny"}
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
"""Run the Bedrock sample agent, invoke the weather tool, and log the response."""
|
||||
agent = ChatAgent(
|
||||
chat_client=BedrockChatClient(),
|
||||
instructions="You are a concise travel assistant.",
|
||||
name="BedrockWeatherAgent",
|
||||
tool_choice=ToolMode.AUTO,
|
||||
tools=[get_weather],
|
||||
)
|
||||
|
||||
response = await agent.run("Use the weather tool to check the forecast for new york.")
|
||||
logging.info("\nAssistant reply:", response.text or "<no text returned>")
|
||||
_log_response(response)
|
||||
|
||||
|
||||
def _log_response(response: AgentRunResponse) -> None:
|
||||
logging.info("\nConversation transcript:")
|
||||
for idx, message in enumerate(response.messages, start=1):
|
||||
tag = f"{idx}. {message.role.value if isinstance(message.role, Role) else message.role}"
|
||||
_log_contents(tag, message.contents)
|
||||
|
||||
|
||||
def _log_contents(tag: str, contents: Sequence[object]) -> None:
|
||||
logging.info(f"[{tag}] {len(contents)} content blocks")
|
||||
for idx, content in enumerate(contents, start=1):
|
||||
if isinstance(content, TextContent):
|
||||
logging.info(f" {idx}. text -> {content.text}")
|
||||
elif isinstance(content, FunctionCallContent):
|
||||
logging.info(f" {idx}. tool_call ({content.name}) -> {content.arguments}")
|
||||
elif isinstance(content, FunctionResultContent):
|
||||
logging.info(f" {idx}. tool_result ({content.call_id}) -> {content.result}")
|
||||
else: # pragma: no cover - defensive
|
||||
logging.info(f" {idx}. {content.type}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,69 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from agent_framework import ChatMessage, ChatOptions, Role, TextContent
|
||||
from agent_framework.exceptions import ServiceInitializationError
|
||||
|
||||
from agent_framework_bedrock import BedrockChatClient
|
||||
|
||||
|
||||
class _StubBedrockRuntime:
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[dict[str, Any]] = []
|
||||
|
||||
def converse(self, **kwargs: Any) -> dict[str, Any]:
|
||||
self.calls.append(kwargs)
|
||||
return {
|
||||
"modelId": kwargs["modelId"],
|
||||
"responseId": "resp-123",
|
||||
"usage": {"inputTokens": 10, "outputTokens": 5, "totalTokens": 15},
|
||||
"output": {
|
||||
"completionReason": "end_turn",
|
||||
"message": {
|
||||
"id": "msg-1",
|
||||
"role": "assistant",
|
||||
"content": [{"text": "Bedrock says hi"}],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def test_get_response_invokes_bedrock_runtime() -> None:
|
||||
stub = _StubBedrockRuntime()
|
||||
client = BedrockChatClient(
|
||||
model_id="amazon.titan-text",
|
||||
region="us-west-2",
|
||||
client=stub,
|
||||
)
|
||||
|
||||
messages = [
|
||||
ChatMessage(role=Role.SYSTEM, contents=[TextContent(text="You are concise.")]),
|
||||
ChatMessage(role=Role.USER, contents=[TextContent(text="hello")]),
|
||||
]
|
||||
|
||||
response = asyncio.run(client.get_response(messages=messages, chat_options=ChatOptions(max_tokens=32)))
|
||||
|
||||
assert stub.calls, "Expected the runtime client to be called"
|
||||
payload = stub.calls[0]
|
||||
assert payload["modelId"] == "amazon.titan-text"
|
||||
assert payload["messages"][0]["content"][0]["text"] == "hello"
|
||||
assert response.messages[0].contents[0].text == "Bedrock says hi"
|
||||
assert response.usage_details and response.usage_details.input_token_count == 10
|
||||
|
||||
|
||||
def test_build_request_requires_non_system_messages() -> None:
|
||||
client = BedrockChatClient(
|
||||
model_id="amazon.titan-text",
|
||||
region="us-west-2",
|
||||
client=_StubBedrockRuntime(),
|
||||
)
|
||||
|
||||
messages = [ChatMessage(role=Role.SYSTEM, contents=[TextContent(text="Only system text")])]
|
||||
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
client._build_converse_request(messages, ChatOptions())
|
||||
@@ -0,0 +1,133 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from agent_framework import (
|
||||
AIFunction,
|
||||
ChatMessage,
|
||||
ChatOptions,
|
||||
FunctionCallContent,
|
||||
FunctionResultContent,
|
||||
Role,
|
||||
TextContent,
|
||||
ToolMode,
|
||||
)
|
||||
from pydantic import BaseModel
|
||||
|
||||
from agent_framework_bedrock._chat_client import BedrockChatClient, BedrockSettings
|
||||
|
||||
|
||||
class _WeatherArgs(BaseModel):
|
||||
location: str
|
||||
|
||||
|
||||
def _build_client() -> BedrockChatClient:
|
||||
fake_runtime = MagicMock()
|
||||
fake_runtime.converse.return_value = {}
|
||||
return BedrockChatClient(model_id="test-model", client=fake_runtime)
|
||||
|
||||
|
||||
def _dummy_weather(location: str) -> str: # pragma: no cover - helper
|
||||
return f"Weather in {location}"
|
||||
|
||||
|
||||
def test_settings_load_from_environment(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("BEDROCK_REGION", "us-west-2")
|
||||
monkeypatch.setenv("BEDROCK_CHAT_MODEL_ID", "anthropic.claude-v2")
|
||||
settings = BedrockSettings()
|
||||
assert settings.region == "us-west-2"
|
||||
assert settings.chat_model_id == "anthropic.claude-v2"
|
||||
|
||||
|
||||
def test_build_request_includes_tool_config() -> None:
|
||||
client = _build_client()
|
||||
|
||||
tool = AIFunction(name="get_weather", description="desc", func=_dummy_weather, input_model=_WeatherArgs)
|
||||
options = ChatOptions(tools=[tool], tool_choice=ToolMode.REQUIRED("get_weather"))
|
||||
messages = [ChatMessage(role=Role.USER, contents=[TextContent(text="hi")])]
|
||||
|
||||
request = client._build_converse_request(messages, options)
|
||||
|
||||
assert request["toolConfig"]["tools"][0]["toolSpec"]["name"] == "get_weather"
|
||||
assert request["toolConfig"]["toolChoice"] == {"tool": {"name": "get_weather"}}
|
||||
|
||||
|
||||
def test_build_request_serializes_tool_history() -> None:
|
||||
client = _build_client()
|
||||
options = ChatOptions()
|
||||
messages = [
|
||||
ChatMessage(role=Role.USER, contents=[TextContent(text="how's weather?")]),
|
||||
ChatMessage(
|
||||
role=Role.ASSISTANT,
|
||||
contents=[FunctionCallContent(call_id="call-1", name="get_weather", arguments='{"location": "SEA"}')],
|
||||
),
|
||||
ChatMessage(
|
||||
role=Role.TOOL,
|
||||
contents=[FunctionResultContent(call_id="call-1", result={"answer": "72F"})],
|
||||
),
|
||||
]
|
||||
|
||||
request = client._build_converse_request(messages, options)
|
||||
assistant_block = request["messages"][1]["content"][0]["toolUse"]
|
||||
result_block = request["messages"][2]["content"][0]["toolResult"]
|
||||
|
||||
assert assistant_block["name"] == "get_weather"
|
||||
assert assistant_block["input"] == {"location": "SEA"}
|
||||
assert result_block["toolUseId"] == "call-1"
|
||||
assert result_block["content"][0]["json"] == {"answer": "72F"}
|
||||
|
||||
|
||||
def test_process_response_parses_tool_use_and_result() -> None:
|
||||
client = _build_client()
|
||||
response = {
|
||||
"modelId": "model",
|
||||
"output": {
|
||||
"message": {
|
||||
"id": "msg-1",
|
||||
"content": [
|
||||
{"toolUse": {"toolUseId": "call-1", "name": "get_weather", "input": {"location": "NYC"}}},
|
||||
{"text": "Calling tool"},
|
||||
],
|
||||
},
|
||||
"completionReason": "tool_use",
|
||||
},
|
||||
}
|
||||
|
||||
chat_response = client._process_converse_response(response)
|
||||
contents = chat_response.messages[0].contents
|
||||
|
||||
assert isinstance(contents[0], FunctionCallContent)
|
||||
assert contents[0].name == "get_weather"
|
||||
assert isinstance(contents[1], TextContent)
|
||||
assert chat_response.finish_reason == client._map_finish_reason("tool_use")
|
||||
|
||||
|
||||
def test_process_response_parses_tool_result() -> None:
|
||||
client = _build_client()
|
||||
response = {
|
||||
"modelId": "model",
|
||||
"output": {
|
||||
"message": {
|
||||
"id": "msg-2",
|
||||
"content": [
|
||||
{
|
||||
"toolResult": {
|
||||
"toolUseId": "call-1",
|
||||
"status": "success",
|
||||
"content": [{"json": {"answer": 42}}],
|
||||
}
|
||||
}
|
||||
],
|
||||
},
|
||||
"completionReason": "end_turn",
|
||||
},
|
||||
}
|
||||
|
||||
chat_response = client._process_converse_response(response)
|
||||
contents = chat_response.messages[0].contents
|
||||
|
||||
assert isinstance(contents[0], FunctionResultContent)
|
||||
assert contents[0].result == {"answer": 42}
|
||||
@@ -90,6 +90,7 @@ agent-framework-azure-ai-search = { workspace = true }
|
||||
agent-framework-anthropic = { workspace = true }
|
||||
agent-framework-azure-ai = { workspace = true }
|
||||
agent-framework-azurefunctions = { workspace = true }
|
||||
agent-framework-bedrock = { workspace = true }
|
||||
agent-framework-chatkit = { workspace = true }
|
||||
agent-framework-copilotstudio = { workspace = true }
|
||||
agent-framework-declarative = { workspace = true }
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""This sample has moved to python/packages/bedrock/samples/bedrock_sample.py."""
|
||||
Generated
+18
@@ -33,6 +33,7 @@ members = [
|
||||
"agent-framework-azure-ai",
|
||||
"agent-framework-azure-ai-search",
|
||||
"agent-framework-azurefunctions",
|
||||
"agent-framework-bedrock",
|
||||
"agent-framework-chatkit",
|
||||
"agent-framework-copilotstudio",
|
||||
"agent-framework-core",
|
||||
@@ -275,6 +276,23 @@ requires-dist = [
|
||||
[package.metadata.requires-dev]
|
||||
dev = [{ name = "types-python-dateutil", specifier = ">=2.9.0" }]
|
||||
|
||||
[[package]]
|
||||
name = "agent-framework-bedrock"
|
||||
version = "1.0.0b251120"
|
||||
source = { editable = "packages/bedrock" }
|
||||
dependencies = [
|
||||
{ name = "agent-framework-core", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
|
||||
{ name = "boto3", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
|
||||
{ name = "botocore", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
|
||||
]
|
||||
|
||||
[package.metadata]
|
||||
requires-dist = [
|
||||
{ name = "agent-framework-core", editable = "packages/core" },
|
||||
{ name = "boto3", specifier = ">=1.35.0,<2.0.0" },
|
||||
{ name = "botocore", specifier = ">=1.35.0,<2.0.0" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "agent-framework-chatkit"
|
||||
version = "1.0.0b251218"
|
||||
|
||||
Reference in New Issue
Block a user