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:
Sunil Dutta
2025-12-19 13:35:53 -05:00
committed by GitHub
Unverified
parent defe0f1a89
commit 3b77192ad0
12 changed files with 958 additions and 0 deletions
+21
View File
@@ -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
+19
View File
@@ -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
+90
View File
@@ -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}
+1
View File
@@ -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 }
+1
View File
@@ -0,0 +1 @@
"""This sample has moved to python/packages/bedrock/samples/bedrock_sample.py."""
+18
View File
@@ -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"