208 lines
7.5 KiB
Python
208 lines
7.5 KiB
Python
"""Tests for the BedrockTransport."""
|
|
|
|
import json
|
|
import pytest
|
|
from types import SimpleNamespace
|
|
|
|
from agent.transports import get_transport
|
|
from agent.transports.types import NormalizedResponse, ToolCall
|
|
|
|
|
|
@pytest.fixture
|
|
def transport():
|
|
# The bedrock plugin registers the transport via the plugin registry.
|
|
from agent.plugin_registries import registries
|
|
if registries.get_provider_service("bedrock", "build_converse_kwargs") is None:
|
|
from hermes_agent_bedrock import register as _bedrock_register
|
|
|
|
class _Ctx:
|
|
def register_provider_services(self, name, services):
|
|
registries.register_provider_services(name, services)
|
|
def register_provider_resolver(self, name, fn):
|
|
registries.register_provider_resolver(name, fn)
|
|
def register_transport(self, api_mode, transport_cls):
|
|
registries._transports[api_mode] = transport_cls
|
|
def register_pricing_provider(self, name, entries):
|
|
registries.register_pricing_provider(name, entries)
|
|
def register_credential_pool_hook(self, name, hook):
|
|
registries.register_credential_pool_hook(name, hook)
|
|
def register_provider_overlay(self, entry):
|
|
registries.register_provider_overlay(entry)
|
|
def __getattr__(self, name):
|
|
if name.startswith("register_"):
|
|
return lambda *a, **kw: None
|
|
raise AttributeError(name)
|
|
|
|
_bedrock_register(_Ctx())
|
|
import agent.transports as _t
|
|
_t._discovered = False
|
|
return get_transport("bedrock_converse")
|
|
|
|
|
|
class TestBedrockBasic:
|
|
|
|
def test_api_mode(self, transport):
|
|
assert transport.api_mode == "bedrock_converse"
|
|
|
|
def test_registered(self, transport):
|
|
assert transport is not None
|
|
|
|
|
|
class TestBedrockBuildKwargs:
|
|
|
|
def test_basic_kwargs(self, transport):
|
|
msgs = [{"role": "user", "content": "Hello"}]
|
|
kw = transport.build_kwargs(model="anthropic.claude-3-5-sonnet-20241022-v2:0", messages=msgs)
|
|
assert kw["modelId"] == "anthropic.claude-3-5-sonnet-20241022-v2:0"
|
|
assert kw["__bedrock_converse__"] is True
|
|
assert kw["__bedrock_region__"] == "us-east-1"
|
|
assert "messages" in kw
|
|
|
|
def test_custom_region(self, transport):
|
|
msgs = [{"role": "user", "content": "Hi"}]
|
|
kw = transport.build_kwargs(
|
|
model="anthropic.claude-3-5-sonnet-20241022-v2:0",
|
|
messages=msgs,
|
|
region="eu-west-1",
|
|
)
|
|
assert kw["__bedrock_region__"] == "eu-west-1"
|
|
|
|
def test_max_tokens(self, transport):
|
|
msgs = [{"role": "user", "content": "Hi"}]
|
|
kw = transport.build_kwargs(
|
|
model="anthropic.claude-3-5-sonnet-20241022-v2:0",
|
|
messages=msgs,
|
|
max_tokens=8192,
|
|
)
|
|
assert kw["inferenceConfig"]["maxTokens"] == 8192
|
|
|
|
|
|
class TestBedrockConvertTools:
|
|
|
|
def test_convert_tools(self, transport):
|
|
tools = [{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "terminal",
|
|
"description": "Run commands",
|
|
"parameters": {"type": "object", "properties": {"command": {"type": "string"}}},
|
|
}
|
|
}]
|
|
result = transport.convert_tools(tools)
|
|
assert len(result) == 1
|
|
assert result[0]["toolSpec"]["name"] == "terminal"
|
|
|
|
|
|
class TestBedrockValidate:
|
|
|
|
def test_none(self, transport):
|
|
assert transport.validate_response(None) is False
|
|
|
|
def test_raw_dict_valid(self, transport):
|
|
assert transport.validate_response({"output": {"message": {}}}) is True
|
|
|
|
def test_raw_dict_invalid(self, transport):
|
|
assert transport.validate_response({"error": "fail"}) is False
|
|
|
|
def test_normalized_valid(self, transport):
|
|
r = SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(content="hi"))])
|
|
assert transport.validate_response(r) is True
|
|
|
|
|
|
class TestBedrockMapFinishReason:
|
|
|
|
def test_end_turn(self, transport):
|
|
assert transport.map_finish_reason("end_turn") == "stop"
|
|
|
|
def test_tool_use(self, transport):
|
|
assert transport.map_finish_reason("tool_use") == "tool_calls"
|
|
|
|
def test_max_tokens(self, transport):
|
|
assert transport.map_finish_reason("max_tokens") == "length"
|
|
|
|
def test_guardrail(self, transport):
|
|
assert transport.map_finish_reason("guardrail_intervened") == "content_filter"
|
|
|
|
def test_unknown(self, transport):
|
|
assert transport.map_finish_reason("unknown") == "stop"
|
|
|
|
|
|
class TestBedrockNormalize:
|
|
|
|
def _make_bedrock_response(self, text="Hello", tool_calls=None, stop_reason="end_turn"):
|
|
"""Build a raw Bedrock converse response dict."""
|
|
content = []
|
|
if text:
|
|
content.append({"text": text})
|
|
if tool_calls:
|
|
for tc in tool_calls:
|
|
content.append({
|
|
"toolUse": {
|
|
"toolUseId": tc["id"],
|
|
"name": tc["name"],
|
|
"input": tc["input"],
|
|
}
|
|
})
|
|
return {
|
|
"output": {"message": {"role": "assistant", "content": content}},
|
|
"stopReason": stop_reason,
|
|
"usage": {"inputTokens": 10, "outputTokens": 5, "totalTokens": 15},
|
|
}
|
|
|
|
def test_text_response(self, transport):
|
|
raw = self._make_bedrock_response(text="Hello world")
|
|
nr = transport.normalize_response(raw)
|
|
assert isinstance(nr, NormalizedResponse)
|
|
assert nr.content == "Hello world"
|
|
assert nr.finish_reason == "stop"
|
|
|
|
def test_tool_call_response(self, transport):
|
|
raw = self._make_bedrock_response(
|
|
text=None,
|
|
tool_calls=[{"id": "tool_1", "name": "terminal", "input": {"command": "ls"}}],
|
|
stop_reason="tool_use",
|
|
)
|
|
nr = transport.normalize_response(raw)
|
|
assert nr.finish_reason == "tool_calls"
|
|
assert len(nr.tool_calls) == 1
|
|
assert nr.tool_calls[0].name == "terminal"
|
|
|
|
def test_raw_reasoning_content_response(self, transport):
|
|
raw = {
|
|
"output": {
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": [
|
|
{"reasoningContent": {"text": "Let me think..."}},
|
|
{"text": "Answer."},
|
|
],
|
|
}
|
|
},
|
|
"stopReason": "end_turn",
|
|
"usage": {"inputTokens": 10, "outputTokens": 5, "totalTokens": 15},
|
|
}
|
|
nr = transport.normalize_response(raw)
|
|
assert nr.reasoning == "Let me think..."
|
|
assert nr.content == "Answer."
|
|
|
|
def test_already_normalized_response(self, transport):
|
|
"""Test normalize_response handles already-normalized SimpleNamespace (from dispatch site)."""
|
|
pre_normalized = SimpleNamespace(
|
|
choices=[SimpleNamespace(
|
|
message=SimpleNamespace(
|
|
content="Hello from Bedrock",
|
|
tool_calls=None,
|
|
reasoning=None,
|
|
reasoning_content=None,
|
|
),
|
|
finish_reason="stop",
|
|
)],
|
|
usage=SimpleNamespace(prompt_tokens=10, completion_tokens=5, total_tokens=15),
|
|
)
|
|
nr = transport.normalize_response(pre_normalized)
|
|
assert isinstance(nr, NormalizedResponse)
|
|
assert nr.content == "Hello from Bedrock"
|
|
assert nr.finish_reason == "stop"
|
|
assert nr.usage is not None
|
|
assert nr.usage.prompt_tokens == 10
|