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>
2026-06-16 21:04:09 +08:00 · 2025-12-19 13:35:53 -05:00
parent defe0f1a89
commit 3b77192ad0
12 changed files with 958 additions and 0 deletions
@@ -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}