chore: import upstream snapshot with attribution
dotnet-build-and-test / dotnet-test-functions (push) Has been cancelled
dotnet-build-and-test / paths-filter (push) Has been cancelled
dotnet-build-and-test / dotnet-build (Debug, windows-latest, net9.0) (push) Has been cancelled
dotnet-build-and-test / dotnet-build (Release, ubuntu-latest, net10.0) (push) Has been cancelled
dotnet-build-and-test / dotnet-build (Release, ubuntu-latest, net8.0) (push) Has been cancelled
dotnet-build-and-test / dotnet-build (Release, windows-latest, net472) (push) Has been cancelled
dotnet-build-and-test / dotnet-test (Release, integration, true, ubuntu-latest, net10.0) (push) Has been cancelled
dotnet-build-and-test / dotnet-test (Release, integration, true, windows-latest, net472) (push) Has been cancelled
dotnet-build-and-test / dotnet-foundry-hosted-it (push) Has been cancelled
dotnet-build-and-test / dotnet-build-and-test-check (push) Has been cancelled
dotnet-build-and-test / Integration Test Report (push) Has been cancelled
CodeQL / Analyze (csharp) (push) Has been cancelled
CodeQL / Analyze (python) (push) Has been cancelled
dotnet-build-and-test / dotnet-test-functions (push) Has been cancelled
dotnet-build-and-test / paths-filter (push) Has been cancelled
dotnet-build-and-test / dotnet-build (Debug, windows-latest, net9.0) (push) Has been cancelled
dotnet-build-and-test / dotnet-build (Release, ubuntu-latest, net10.0) (push) Has been cancelled
dotnet-build-and-test / dotnet-build (Release, ubuntu-latest, net8.0) (push) Has been cancelled
dotnet-build-and-test / dotnet-build (Release, windows-latest, net472) (push) Has been cancelled
dotnet-build-and-test / dotnet-test (Release, integration, true, ubuntu-latest, net10.0) (push) Has been cancelled
dotnet-build-and-test / dotnet-test (Release, integration, true, windows-latest, net472) (push) Has been cancelled
dotnet-build-and-test / dotnet-foundry-hosted-it (push) Has been cancelled
dotnet-build-and-test / dotnet-build-and-test-check (push) Has been cancelled
dotnet-build-and-test / Integration Test Report (push) Has been cancelled
CodeQL / Analyze (csharp) (push) Has been cancelled
CodeQL / Analyze (python) (push) Has been cancelled
This commit is contained in:
@@ -0,0 +1,169 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from agent_framework import Embedding, GeneratedEmbeddings
|
||||
|
||||
from agent_framework_bedrock import BedrockEmbeddingClient, BedrockEmbeddingOptions
|
||||
|
||||
|
||||
class _StubBedrockEmbeddingRuntime:
|
||||
"""Stub for the Bedrock runtime client that handles invoke_model for embeddings."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[dict[str, Any]] = []
|
||||
self.meta = MagicMock(endpoint_url="https://bedrock-runtime.us-west-2.amazonaws.com")
|
||||
|
||||
def invoke_model(self, **kwargs: Any) -> dict[str, Any]:
|
||||
self.calls.append(kwargs)
|
||||
body = json.loads(kwargs.get("body", "{}"))
|
||||
# Simulate Titan embedding response
|
||||
dimensions = body.get("dimensions", 3)
|
||||
return {
|
||||
"body": MagicMock(
|
||||
read=lambda: json.dumps({
|
||||
"embedding": [0.1 * (i + 1) for i in range(dimensions)],
|
||||
"inputTextTokenCount": 5,
|
||||
}).encode()
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
async def test_bedrock_embedding_construction() -> None:
|
||||
"""Test construction with explicit parameters."""
|
||||
stub = _StubBedrockEmbeddingRuntime()
|
||||
client = BedrockEmbeddingClient(
|
||||
model="amazon.titan-embed-text-v2:0",
|
||||
region="us-west-2",
|
||||
client=stub, # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] # pyright: ignore[reportArgumentType]
|
||||
)
|
||||
assert client.model == "amazon.titan-embed-text-v2:0"
|
||||
assert client.region == "us-west-2"
|
||||
|
||||
|
||||
async def test_bedrock_embedding_construction_missing_model_raises(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Test that missing model raises an error."""
|
||||
monkeypatch.delenv("BEDROCK_EMBEDDING_MODEL", raising=False)
|
||||
from agent_framework.exceptions import SettingNotFoundError
|
||||
|
||||
with pytest.raises(SettingNotFoundError):
|
||||
BedrockEmbeddingClient(region="us-west-2")
|
||||
|
||||
|
||||
async def test_bedrock_embedding_get_embeddings() -> None:
|
||||
"""Test generating embeddings via the Bedrock invoke_model API."""
|
||||
stub = _StubBedrockEmbeddingRuntime()
|
||||
client = BedrockEmbeddingClient(
|
||||
model="amazon.titan-embed-text-v2:0",
|
||||
region="us-west-2",
|
||||
client=stub, # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] # pyright: ignore[reportArgumentType]
|
||||
)
|
||||
|
||||
result = await client.get_embeddings(["hello", "world"])
|
||||
|
||||
assert isinstance(result, GeneratedEmbeddings)
|
||||
assert len(result) == 2
|
||||
assert len(result[0].vector) == 3
|
||||
assert len(result[1].vector) == 3
|
||||
assert result[0].model == "amazon.titan-embed-text-v2:0"
|
||||
assert result.usage == {"input_token_count": 10}
|
||||
|
||||
# Two calls since Titan processes one input at a time
|
||||
assert len(stub.calls) == 2
|
||||
call_texts = {json.loads(call["body"])["inputText"] for call in stub.calls}
|
||||
assert call_texts == {"hello", "world"}
|
||||
|
||||
|
||||
async def test_bedrock_embedding_get_embeddings_empty_input() -> None:
|
||||
"""Test generating embeddings with empty input."""
|
||||
stub = _StubBedrockEmbeddingRuntime()
|
||||
client = BedrockEmbeddingClient(
|
||||
model="amazon.titan-embed-text-v2:0",
|
||||
region="us-west-2",
|
||||
client=stub, # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] # pyright: ignore[reportArgumentType]
|
||||
)
|
||||
|
||||
result = await client.get_embeddings([])
|
||||
|
||||
assert isinstance(result, GeneratedEmbeddings)
|
||||
assert len(result) == 0
|
||||
assert len(stub.calls) == 0
|
||||
|
||||
|
||||
async def test_bedrock_embedding_get_embeddings_with_options() -> None:
|
||||
"""Test generating embeddings with custom options."""
|
||||
stub = _StubBedrockEmbeddingRuntime()
|
||||
client = BedrockEmbeddingClient(
|
||||
model="amazon.titan-embed-text-v2:0",
|
||||
region="us-west-2",
|
||||
client=stub, # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] # pyright: ignore[reportArgumentType]
|
||||
)
|
||||
|
||||
options: BedrockEmbeddingOptions = {
|
||||
"dimensions": 5,
|
||||
"normalize": True,
|
||||
}
|
||||
result = await client.get_embeddings(["hello"], options=options) # ty: ignore[invalid-argument-type]
|
||||
|
||||
assert len(result) == 1
|
||||
assert len(result[0].vector) == 5
|
||||
|
||||
body = json.loads(stub.calls[0]["body"])
|
||||
assert body["dimensions"] == 5
|
||||
assert body["normalize"] is True
|
||||
|
||||
|
||||
async def test_bedrock_embedding_get_embeddings_no_model_raises() -> None:
|
||||
"""Test that missing model at call time raises ValueError."""
|
||||
stub = _StubBedrockEmbeddingRuntime()
|
||||
client = BedrockEmbeddingClient(
|
||||
model="amazon.titan-embed-text-v2:0",
|
||||
region="us-west-2",
|
||||
client=stub, # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] # pyright: ignore[reportArgumentType]
|
||||
)
|
||||
client.model = None # type: ignore[assignment] # ty: ignore[invalid-assignment]
|
||||
|
||||
with pytest.raises(ValueError, match="model is required"):
|
||||
await client.get_embeddings(["hello"])
|
||||
|
||||
|
||||
async def test_bedrock_embedding_default_region() -> None:
|
||||
"""Test that default region is us-east-1."""
|
||||
stub = _StubBedrockEmbeddingRuntime()
|
||||
client = BedrockEmbeddingClient(
|
||||
model="amazon.titan-embed-text-v2:0",
|
||||
client=stub, # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] # pyright: ignore[reportArgumentType]
|
||||
)
|
||||
assert client.region == "us-east-1"
|
||||
|
||||
|
||||
# region: Integration Tests
|
||||
|
||||
skip_if_bedrock_embedding_integration_tests_disabled = pytest.mark.skipif(
|
||||
os.getenv("BEDROCK_EMBEDDING_MODEL", "") in ("", "test-model")
|
||||
or not (os.getenv("AWS_ACCESS_KEY_ID") or os.getenv("BEDROCK_ACCESS_KEY")),
|
||||
reason="No real Bedrock embedding model or AWS credentials provided; skipping integration tests.",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.flaky
|
||||
@pytest.mark.integration
|
||||
@skip_if_bedrock_embedding_integration_tests_disabled
|
||||
async def test_bedrock_embedding_integration() -> None:
|
||||
"""Integration test for Bedrock embedding client."""
|
||||
client = BedrockEmbeddingClient()
|
||||
result = await client.get_embeddings(["Hello, world!", "How are you?"])
|
||||
|
||||
assert isinstance(result, GeneratedEmbeddings)
|
||||
assert len(result) == 2
|
||||
for embedding in result:
|
||||
assert isinstance(embedding, Embedding)
|
||||
assert isinstance(embedding.vector, list)
|
||||
assert len(embedding.vector) > 0
|
||||
assert all(isinstance(v, float) for v in embedding.vector)
|
||||
@@ -0,0 +1,236 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from agent_framework import Agent, Content, Message
|
||||
|
||||
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 _make_client() -> BedrockChatClient:
|
||||
"""Create a BedrockChatClient with a stub runtime for unit tests."""
|
||||
return BedrockChatClient(
|
||||
model="amazon.titan-text",
|
||||
region="us-west-2",
|
||||
client=_StubBedrockRuntime(), # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] # pyright: ignore[reportArgumentType]
|
||||
)
|
||||
|
||||
|
||||
def test_agent_accepts_bedrock_chat_client() -> None:
|
||||
client = _make_client()
|
||||
agent = Agent(client=client, instructions="test agent")
|
||||
assert agent.client is client
|
||||
|
||||
|
||||
async def test_get_response_invokes_bedrock_runtime() -> None:
|
||||
stub = _StubBedrockRuntime()
|
||||
client = BedrockChatClient(
|
||||
model="amazon.titan-text",
|
||||
region="us-west-2",
|
||||
client=stub, # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] # pyright: ignore[reportArgumentType]
|
||||
)
|
||||
|
||||
messages = [
|
||||
Message(role="system", contents=[Content.from_text(text="You are concise.")]),
|
||||
Message(role="user", contents=[Content.from_text(text="hello")]),
|
||||
]
|
||||
|
||||
response = await client.get_response(messages=messages, options={"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="amazon.titan-text",
|
||||
region="us-west-2",
|
||||
client=_StubBedrockRuntime(), # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] # pyright: ignore[reportArgumentType]
|
||||
)
|
||||
|
||||
messages = [Message(role="system", contents=[Content.from_text(text="Only system text")])]
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
client._prepare_options(messages, {})
|
||||
|
||||
|
||||
def test_prepare_options_tool_choice_none_omits_tool_config() -> None:
|
||||
"""When tool_choice='none', toolConfig must be omitted entirely.
|
||||
|
||||
Bedrock's Converse API only accepts 'auto', 'any', or 'tool' as valid
|
||||
toolChoice keys. Sending {"none": {}} causes a ParamValidationError.
|
||||
The fix omits toolConfig so the model won't attempt tool calls.
|
||||
|
||||
Fixes #4529.
|
||||
"""
|
||||
client = _make_client()
|
||||
messages = [Message(role="user", contents=[Content.from_text(text="hello")])]
|
||||
|
||||
# Even when tools are provided, tool_choice="none" should strip toolConfig
|
||||
options: dict[str, Any] = {
|
||||
"tool_choice": "none",
|
||||
"tools": [
|
||||
{"toolSpec": {"name": "get_weather", "description": "Get weather", "inputSchema": {"json": {}}}},
|
||||
],
|
||||
}
|
||||
|
||||
request = client._prepare_options(messages, options)
|
||||
|
||||
assert "toolConfig" not in request, (
|
||||
f"toolConfig should be omitted when tool_choice='none', got: {request.get('toolConfig')}"
|
||||
)
|
||||
|
||||
|
||||
def test_prepare_options_tool_choice_auto_includes_tool_config() -> None:
|
||||
"""When tool_choice='auto', toolConfig.toolChoice should be {'auto': {}}."""
|
||||
client = _make_client()
|
||||
messages = [Message(role="user", contents=[Content.from_text(text="hello")])]
|
||||
|
||||
options: dict[str, Any] = {
|
||||
"tool_choice": "auto",
|
||||
"tools": [
|
||||
{"toolSpec": {"name": "get_weather", "description": "Get weather", "inputSchema": {"json": {}}}},
|
||||
],
|
||||
}
|
||||
|
||||
request = client._prepare_options(messages, options)
|
||||
|
||||
assert "toolConfig" in request
|
||||
assert request["toolConfig"]["toolChoice"] == {"auto": {}}
|
||||
|
||||
|
||||
def test_prepare_options_tool_choice_required_includes_any() -> None:
|
||||
"""When tool_choice='required' (no specific function), toolChoice should be {'any': {}}."""
|
||||
client = _make_client()
|
||||
messages = [Message(role="user", contents=[Content.from_text(text="hello")])]
|
||||
|
||||
options: dict[str, Any] = {
|
||||
"tool_choice": "required",
|
||||
"tools": [
|
||||
{"toolSpec": {"name": "get_weather", "description": "Get weather", "inputSchema": {"json": {}}}},
|
||||
],
|
||||
}
|
||||
|
||||
request = client._prepare_options(messages, options)
|
||||
|
||||
assert "toolConfig" in request
|
||||
assert request["toolConfig"]["toolChoice"] == {"any": {}}
|
||||
|
||||
|
||||
def test_prepare_options_tool_choice_auto_without_tools_omits_tool_config() -> None:
|
||||
"""When tool_choice='auto' but no tools are provided, toolConfig must be omitted.
|
||||
|
||||
Without tools, setting toolChoice would cause a ParamValidationError from Bedrock.
|
||||
"""
|
||||
client = _make_client()
|
||||
messages = [Message(role="user", contents=[Content.from_text(text="hello")])]
|
||||
|
||||
options: dict[str, Any] = {
|
||||
"tool_choice": "auto",
|
||||
}
|
||||
|
||||
request = client._prepare_options(messages, options)
|
||||
|
||||
assert "toolConfig" not in request, (
|
||||
f"toolConfig should be omitted when no tools are provided, got: {request.get('toolConfig')}"
|
||||
)
|
||||
|
||||
|
||||
def test_prepare_options_tool_choice_required_without_tools_raises() -> None:
|
||||
"""When tool_choice='required' but no tools are provided, a ValueError must be raised."""
|
||||
client = _make_client()
|
||||
messages = [Message(role="user", contents=[Content.from_text(text="hello")])]
|
||||
|
||||
options: dict[str, Any] = {
|
||||
"tool_choice": "required",
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="tool_choice='required' requires at least one tool"):
|
||||
client._prepare_options(messages, options)
|
||||
|
||||
|
||||
def test_process_converse_response_preserves_non_ascii_in_json_block() -> None:
|
||||
"""Non-ASCII text in a Bedrock ``json`` content block must be preserved, not \\uXXXX-escaped.
|
||||
|
||||
The Converse API can return structured ``json`` content blocks. These are serialized to
|
||||
text via ``json.dumps``; without ``ensure_ascii=False`` CJK characters and emoji are escaped
|
||||
to ``\\uXXXX`` sequences and surface garbled to the user.
|
||||
"""
|
||||
client = _make_client()
|
||||
json_payload = {"greeting": "你好世界", "emoji": "🎉"}
|
||||
response: dict[str, Any] = {
|
||||
"modelId": "amazon.titan-text",
|
||||
"output": {
|
||||
"completionReason": "end_turn",
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": [{"json": json_payload}],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
chat_response = client._process_converse_response(response)
|
||||
|
||||
text = chat_response.messages[0].text
|
||||
assert "你好世界" in text
|
||||
assert "🎉" in text
|
||||
# Must not be escaped to Unicode code points.
|
||||
assert "\\u" not in text
|
||||
# Serialized text must remain valid JSON that round-trips to the original payload.
|
||||
assert json.loads(text) == json_payload
|
||||
|
||||
|
||||
def test_parse_usage_surfaces_cache_tokens() -> None:
|
||||
"""Bedrock Converse reports cache token counts when prompt caching is used."""
|
||||
client = _make_client()
|
||||
|
||||
details = client._parse_usage({
|
||||
"inputTokens": 10,
|
||||
"outputTokens": 5,
|
||||
"totalTokens": 15,
|
||||
"cacheReadInputTokens": 8,
|
||||
"cacheWriteInputTokens": 3,
|
||||
})
|
||||
|
||||
assert details is not None
|
||||
assert details["input_token_count"] == 10
|
||||
assert details["cache_read_input_token_count"] == 8
|
||||
assert details["cache_creation_input_token_count"] == 3
|
||||
|
||||
|
||||
def test_parse_usage_returns_none_when_no_recognized_keys() -> None:
|
||||
"""A truthy usage payload with no recognized keys yields None, not an empty mapping."""
|
||||
client = _make_client()
|
||||
|
||||
assert client._parse_usage({"unexpected": 1}) is None
|
||||
assert client._parse_usage({}) is None
|
||||
assert client._parse_usage(None) is None
|
||||
@@ -0,0 +1,136 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from agent_framework import (
|
||||
ChatOptions,
|
||||
Content,
|
||||
FunctionTool,
|
||||
Message,
|
||||
)
|
||||
from agent_framework._settings import load_settings
|
||||
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="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", "anthropic.claude-v2")
|
||||
settings = load_settings(BedrockSettings, env_prefix="BEDROCK_")
|
||||
assert settings["region"] == "us-west-2"
|
||||
assert settings["chat_model"] == "anthropic.claude-v2"
|
||||
|
||||
|
||||
def test_build_request_includes_tool_config() -> None:
|
||||
client = _build_client()
|
||||
|
||||
tool = FunctionTool(name="get_weather", description="desc", func=_dummy_weather, input_model=_WeatherArgs)
|
||||
options = {
|
||||
"tools": [tool],
|
||||
"tool_choice": {"mode": "required", "required_function_name": "get_weather"},
|
||||
}
|
||||
messages = [Message(role="user", contents=[Content.from_text(text="hi")])]
|
||||
|
||||
request = client._prepare_options(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 = [
|
||||
Message(role="user", contents=[Content.from_text(text="how's weather?")]),
|
||||
Message(
|
||||
role="assistant",
|
||||
contents=[
|
||||
Content.from_function_call(call_id="call-1", name="get_weather", arguments='{"location": "SEA"}')
|
||||
],
|
||||
),
|
||||
Message(
|
||||
role="tool",
|
||||
contents=[Content.from_function_result(call_id="call-1", result='{"answer": "72F"}')],
|
||||
),
|
||||
]
|
||||
|
||||
request = client._prepare_options(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 contents[0].type == "function_call"
|
||||
assert contents[0].name == "get_weather"
|
||||
assert contents[1].type == "text"
|
||||
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 contents[0].type == "function_result"
|
||||
assert "answer" in str(contents[0].result)
|
||||
assert contents[0].items is not None
|
||||
@@ -0,0 +1,378 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import json
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from agent_framework import Content, Message
|
||||
from botocore.exceptions import ClientError
|
||||
from pydantic import BaseModel
|
||||
|
||||
from agent_framework_bedrock import BedrockChatClient
|
||||
|
||||
# region Test models
|
||||
|
||||
|
||||
class WeatherReport(BaseModel):
|
||||
city: str
|
||||
temperature: float
|
||||
summary: str
|
||||
|
||||
|
||||
class NestedAddress(BaseModel):
|
||||
street: str
|
||||
city: str
|
||||
zip_code: str
|
||||
|
||||
|
||||
class Person(BaseModel):
|
||||
name: str
|
||||
age: int
|
||||
address: NestedAddress
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
|
||||
# region Helpers
|
||||
|
||||
|
||||
class _StubBedrockRuntime:
|
||||
"""Stub that records calls and returns a canned response."""
|
||||
|
||||
def __init__(self, response_text: str = "Bedrock says hi") -> None:
|
||||
self.calls: list[dict[str, Any]] = []
|
||||
self._response_text = response_text
|
||||
|
||||
def converse(self, **kwargs: Any) -> dict[str, Any]:
|
||||
self.calls.append(kwargs)
|
||||
return {
|
||||
"modelId": kwargs["modelId"],
|
||||
"responseId": "resp-structured",
|
||||
"usage": {"inputTokens": 10, "outputTokens": 20, "totalTokens": 30},
|
||||
"output": {
|
||||
"completionReason": "end_turn",
|
||||
"message": {
|
||||
"id": "msg-structured",
|
||||
"role": "assistant",
|
||||
"content": [{"text": self._response_text}],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _make_client(response_text: str = "Bedrock says hi") -> tuple[BedrockChatClient, _StubBedrockRuntime]:
|
||||
stub = _StubBedrockRuntime(response_text)
|
||||
client = BedrockChatClient(
|
||||
model="us.anthropic.claude-haiku-4-5-v1:0",
|
||||
region="us-east-1",
|
||||
client=stub, # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] # pyright: ignore[reportArgumentType]
|
||||
)
|
||||
return client, stub
|
||||
|
||||
|
||||
def _user_messages() -> list[Message]:
|
||||
return [Message(role="user", contents=[Content.from_text(text="Give me a weather report")])]
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
|
||||
# region Tests
|
||||
|
||||
|
||||
def test_prepare_output_config_correct_wire_shape() -> None:
|
||||
"""_prepare_output_config(WeatherReport) must produce the correct
|
||||
textFormat → structure → jsonSchema shape with type: 'json_schema'."""
|
||||
client, _ = _make_client()
|
||||
|
||||
output_config = client._prepare_output_config(WeatherReport)
|
||||
|
||||
assert output_config is not None
|
||||
text_format = output_config["textFormat"]
|
||||
assert text_format["type"] == "json_schema"
|
||||
assert "structure" in text_format
|
||||
json_schema = text_format["structure"]["jsonSchema"]
|
||||
assert json_schema["name"] == "WeatherReport"
|
||||
assert "schema" in json_schema
|
||||
|
||||
|
||||
def test_prepare_output_config_schema_is_json_string() -> None:
|
||||
"""The schema value inside jsonSchema must be a JSON string, not a dict."""
|
||||
client, _ = _make_client()
|
||||
|
||||
output_config = client._prepare_output_config(WeatherReport)
|
||||
|
||||
assert output_config is not None
|
||||
schema_value = output_config["textFormat"]["structure"]["jsonSchema"]["schema"]
|
||||
assert isinstance(schema_value, str), f"Expected str, got {type(schema_value)}"
|
||||
# Verify it's valid JSON
|
||||
parsed = json.loads(schema_value)
|
||||
assert isinstance(parsed, dict)
|
||||
assert parsed["type"] == "object"
|
||||
|
||||
|
||||
def test_additional_properties_false_set_recursively() -> None:
|
||||
"""additionalProperties: false must be set on all nested object types."""
|
||||
client, _ = _make_client()
|
||||
|
||||
output_config = client._prepare_output_config(Person)
|
||||
|
||||
assert output_config is not None
|
||||
schema_str = output_config["textFormat"]["structure"]["jsonSchema"]["schema"]
|
||||
schema = json.loads(schema_str)
|
||||
|
||||
# Top-level object
|
||||
assert schema.get("additionalProperties") is False
|
||||
|
||||
# Check $defs for NestedAddress
|
||||
defs = schema.get("$defs", {})
|
||||
assert "NestedAddress" in defs, "Expected NestedAddress to be present in $defs"
|
||||
assert defs["NestedAddress"].get("additionalProperties") is False, (
|
||||
"Expected additionalProperties=False on nested NestedAddress schema"
|
||||
)
|
||||
|
||||
|
||||
def test_no_output_config_when_response_format_none() -> None:
|
||||
"""When response_format is None, no outputConfig key should appear in the request."""
|
||||
client, stub = _make_client()
|
||||
messages = _user_messages()
|
||||
|
||||
request = client._prepare_options(messages, {"max_tokens": 100})
|
||||
|
||||
assert "outputConfig" not in request, (
|
||||
f"outputConfig should not be present when response_format is None, got: {request.get('outputConfig')}"
|
||||
)
|
||||
|
||||
|
||||
async def test_chat_response_value_populated() -> None:
|
||||
"""After a mocked response with response_format, .value should be a populated Pydantic model."""
|
||||
json_response = json.dumps({"city": "Seattle", "temperature": 72.5, "summary": "Sunny and warm"})
|
||||
client, stub = _make_client(response_text=json_response)
|
||||
messages = _user_messages()
|
||||
|
||||
response = await client.get_response(
|
||||
messages=messages,
|
||||
options={"max_tokens": 100, "response_format": WeatherReport},
|
||||
)
|
||||
|
||||
assert response.text == json_response
|
||||
assert response.value is not None
|
||||
assert isinstance(response.value, WeatherReport)
|
||||
assert response.value.city == "Seattle"
|
||||
assert response.value.temperature == 72.5
|
||||
assert response.value.summary == "Sunny and warm"
|
||||
|
||||
# Verify outputConfig was sent to the API
|
||||
assert len(stub.calls) == 1
|
||||
api_request = stub.calls[0]
|
||||
assert "outputConfig" in api_request
|
||||
assert api_request["outputConfig"]["textFormat"]["type"] == "json_schema"
|
||||
|
||||
|
||||
def test_dict_schema_response_format() -> None:
|
||||
"""_prepare_output_config should work when response_format is a dict, not just a Pydantic class."""
|
||||
client, _ = _make_client()
|
||||
|
||||
dict_schema = {
|
||||
"json_schema": {
|
||||
"name": "weather_output",
|
||||
"schema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"city": {"type": "string"},
|
||||
"temp": {"type": "number"},
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
output_config = client._prepare_output_config(dict_schema)
|
||||
|
||||
assert output_config is not None
|
||||
json_schema = output_config["textFormat"]["structure"]["jsonSchema"]
|
||||
assert json_schema["name"] == "weather_output"
|
||||
schema_parsed = json.loads(json_schema["schema"])
|
||||
assert schema_parsed["type"] == "object"
|
||||
assert "city" in schema_parsed["properties"]
|
||||
|
||||
|
||||
def test_prepare_output_config_none_returns_none() -> None:
|
||||
"""_prepare_output_config(None) must return None."""
|
||||
client, _ = _make_client()
|
||||
|
||||
result = client._prepare_output_config(None)
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
async def test_chat_response_value_populated_streaming() -> None:
|
||||
"""In streaming mode, .value should also be populated on the final response."""
|
||||
json_response = json.dumps({"city": "Portland", "temperature": 68.0, "summary": "Cloudy"})
|
||||
client, stub = _make_client(response_text=json_response)
|
||||
messages = _user_messages()
|
||||
|
||||
stream = client.get_response(
|
||||
messages=messages,
|
||||
stream=True,
|
||||
options={"max_tokens": 100, "response_format": WeatherReport},
|
||||
)
|
||||
|
||||
# Consume stream and get final response
|
||||
async for _ in stream:
|
||||
pass
|
||||
response = await stream.get_final_response()
|
||||
|
||||
assert response.value is not None
|
||||
assert isinstance(response.value, WeatherReport)
|
||||
assert response.value.city == "Portland"
|
||||
|
||||
# Verify outputConfig was sent
|
||||
assert len(stub.calls) == 1
|
||||
assert "outputConfig" in stub.calls[0]
|
||||
|
||||
|
||||
async def test_unsupported_model_validation_exception() -> None:
|
||||
"""When a model doesn't support outputConfig, a clear error should be raised."""
|
||||
|
||||
class _FailingStubBedrockRuntime:
|
||||
def converse(self, **kwargs: Any) -> dict[str, Any]:
|
||||
# Simulate botocore ClientError for ValidationException
|
||||
error_response = {"Error": {"Code": "ValidationException", "Message": "Invalid field outputConfig"}}
|
||||
raise ClientError(error_response, "Converse")
|
||||
|
||||
client = BedrockChatClient(
|
||||
model="us.anthropic.claude-v2",
|
||||
region="us-east-1",
|
||||
client=_FailingStubBedrockRuntime(), # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] # pyright: ignore[reportArgumentType]
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError) as exc:
|
||||
await client.get_response(
|
||||
messages=_user_messages(),
|
||||
options={"response_format": WeatherReport},
|
||||
)
|
||||
|
||||
assert "does not support structured output via outputConfig.textFormat" in str(exc.value)
|
||||
assert "Check the model's Bedrock Converse outputConfig/textFormat support." in str(exc.value)
|
||||
|
||||
|
||||
def test_invalid_response_format_type_raises() -> None:
|
||||
"""Non-dict, non-BaseModel response_format should raise TypeError."""
|
||||
client, _ = _make_client()
|
||||
with pytest.raises(TypeError, match="Pydantic BaseModel subclass"):
|
||||
client._prepare_output_config("not_a_valid_format")
|
||||
|
||||
|
||||
def test_mapping_response_format_accepted() -> None:
|
||||
"""A non-dict Mapping response_format must be accepted and produce
|
||||
correct outputConfig, not raise TypeError."""
|
||||
from collections.abc import MutableMapping
|
||||
|
||||
class _WrappedMapping(MutableMapping):
|
||||
def __init__(self, data):
|
||||
self._data = dict(data)
|
||||
|
||||
def __getitem__(self, key):
|
||||
return self._data[key]
|
||||
|
||||
def __setitem__(self, key, value):
|
||||
self._data[key] = value
|
||||
|
||||
def __delitem__(self, key):
|
||||
del self._data[key]
|
||||
|
||||
def __iter__(self):
|
||||
return iter(self._data)
|
||||
|
||||
def __len__(self):
|
||||
return len(self._data)
|
||||
|
||||
client, _ = _make_client()
|
||||
mapping_format = _WrappedMapping({
|
||||
"json_schema": {
|
||||
"name": "test_output",
|
||||
"schema": {
|
||||
"type": "object",
|
||||
"properties": {"result": {"type": "string"}},
|
||||
},
|
||||
}
|
||||
})
|
||||
|
||||
output_config = client._prepare_output_config(mapping_format)
|
||||
|
||||
assert output_config is not None
|
||||
json_schema = output_config["textFormat"]["structure"]["jsonSchema"]
|
||||
assert json_schema["name"] == "test_output"
|
||||
schema = json.loads(json_schema["schema"])
|
||||
assert schema.get("additionalProperties") is False
|
||||
|
||||
|
||||
def test_shape_b_dict_schema_wire_format() -> None:
|
||||
"""Dict response_format in Shape B (inner shape directly) should
|
||||
produce correct outputConfig."""
|
||||
client, _ = _make_client()
|
||||
|
||||
response_format = {
|
||||
"name": "weather_output",
|
||||
"schema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"city": {"type": "string"},
|
||||
"temperature": {"type": "number"},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
output_config = client._prepare_output_config(response_format)
|
||||
|
||||
assert output_config is not None
|
||||
text_format = output_config["textFormat"]
|
||||
assert text_format["type"] == "json_schema"
|
||||
json_schema = text_format["structure"]["jsonSchema"]
|
||||
assert json_schema["name"] == "weather_output"
|
||||
schema = json.loads(json_schema["schema"])
|
||||
assert schema.get("additionalProperties") is False
|
||||
|
||||
|
||||
def test_dict_schema_not_mutated() -> None:
|
||||
"""Caller's dict schema must not be mutated by _prepare_output_config."""
|
||||
client, _ = _make_client()
|
||||
original_schema = {
|
||||
"json_schema": {
|
||||
"name": "test",
|
||||
"schema": {
|
||||
"type": "object",
|
||||
"properties": {"a": {"type": "string"}},
|
||||
},
|
||||
}
|
||||
}
|
||||
snapshot = copy.deepcopy(original_schema)
|
||||
client._prepare_output_config(original_schema)
|
||||
assert original_schema == snapshot, "Original dict schema was mutated"
|
||||
|
||||
|
||||
async def test_non_outputconfig_validation_exception_propagates() -> None:
|
||||
"""ValidationException unrelated to outputConfig must propagate
|
||||
as raw ClientError, not be caught and reclassified."""
|
||||
client, _ = _make_client()
|
||||
error_response = {
|
||||
"Error": {
|
||||
"Code": "ValidationException",
|
||||
"Message": "Invalid message format",
|
||||
}
|
||||
}
|
||||
failing_client = MagicMock()
|
||||
failing_client.converse.side_effect = ClientError(error_response, "Converse")
|
||||
with patch.object(client, "_bedrock_client", failing_client), pytest.raises(ClientError):
|
||||
await client.get_response(
|
||||
messages=_user_messages(),
|
||||
options={"max_tokens": 100},
|
||||
)
|
||||
|
||||
|
||||
# endregion
|
||||
Reference in New Issue
Block a user