diff --git a/python/packages/bedrock/LICENSE b/python/packages/bedrock/LICENSE new file mode 100644 index 0000000000..79656060de --- /dev/null +++ b/python/packages/bedrock/LICENSE @@ -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 \ No newline at end of file diff --git a/python/packages/bedrock/README.md b/python/packages/bedrock/README.md new file mode 100644 index 0000000000..6bcd9ff53a --- /dev/null +++ b/python/packages/bedrock/README.md @@ -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 diff --git a/python/packages/bedrock/agent_framework_bedrock/__init__.py b/python/packages/bedrock/agent_framework_bedrock/__init__.py new file mode 100644 index 0000000000..84f3e5946c --- /dev/null +++ b/python/packages/bedrock/agent_framework_bedrock/__init__.py @@ -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__", +] diff --git a/python/packages/bedrock/agent_framework_bedrock/_chat_client.py b/python/packages/bedrock/agent_framework_bedrock/_chat_client.py new file mode 100644 index 0000000000..0788c7e522 --- /dev/null +++ b/python/packages/bedrock/agent_framework_bedrock/_chat_client.py @@ -0,0 +1,540 @@ +# 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_observability +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_observability +@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. + + Parameters + ---------- + region: str | None + Region to send Bedrock requests to; falls back to BEDROCK_REGION. + model_id: str | None + Default model identifier; falls back to BEDROCK_CHAT_MODEL_ID. + access_key: str | None + Optional AWS access key for manual credential injection. + secret_key: str | None + Optional AWS secret key paired with ``access_key``. + session_token: str | None + Optional AWS session token for temporary credentials. + client: BaseClient | None + Preconfigured Bedrock runtime client; when omitted a boto3 session is created. + boto3_session: Boto3Session | None + Custom boto3 session used to build the runtime client if provided. + env_file_path: str | None + Optional .env file path used by ``BedrockSettings`` to load defaults. + env_file_encoding: str | None + Encoding for the optional .env file. + kwargs: Any + 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 diff --git a/python/packages/bedrock/pyproject.toml b/python/packages/bedrock/pyproject.toml new file mode 100644 index 0000000000..ea6cffda42 --- /dev/null +++ b/python/packages/bedrock/pyproject.toml @@ -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" \ No newline at end of file diff --git a/python/packages/bedrock/samples/__init__.py b/python/packages/bedrock/samples/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/python/packages/bedrock/samples/bedrock_sample.py b/python/packages/bedrock/samples/bedrock_sample.py new file mode 100644 index 0000000000..5b5e900923 --- /dev/null +++ b/python/packages/bedrock/samples/bedrock_sample.py @@ -0,0 +1,64 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio + +from collections.abc import Sequence + +from agent_framework import ( + AgentRunResponse, + ChatAgent, + FunctionCallContent, + FunctionResultContent, + Role, + TextContent, + ToolMode, + ai_function, +) +import logging + +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: + 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 "") + _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()) diff --git a/python/packages/bedrock/tests/test_bedrock_client.py b/python/packages/bedrock/tests/test_bedrock_client.py new file mode 100644 index 0000000000..4086dfa429 --- /dev/null +++ b/python/packages/bedrock/tests/test_bedrock_client.py @@ -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()) diff --git a/python/packages/bedrock/tests/test_bedrock_settings.py b/python/packages/bedrock/tests/test_bedrock_settings.py new file mode 100644 index 0000000000..a3b0894d28 --- /dev/null +++ b/python/packages/bedrock/tests/test_bedrock_settings.py @@ -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} diff --git a/python/pyproject.toml b/python/pyproject.toml index 464f7c3f61..aedfcf1bd6 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -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 } diff --git a/python/samples/amazon/bedrock_sample.py b/python/samples/amazon/bedrock_sample.py new file mode 100644 index 0000000000..42feb98ebd --- /dev/null +++ b/python/samples/amazon/bedrock_sample.py @@ -0,0 +1 @@ +"""This sample has moved to python/packages/bedrock/samples/bedrock_sample.py.""" diff --git a/python/uv.lock b/python/uv.lock index 04f3469175..c2a9def1db 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -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", @@ -273,6 +274,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.0b251211"