Files
agent-framework/python/packages/bedrock/tests/test_bedrock_client.py
T
Eduard van Valkenburg 0521f5bed8 Python: [BREAKING] Simplify API: ChatAgent -> Agent, ChatMessage -> Message (#3747)
* [BREAKING] Rename ChatAgent -> Agent, ChatMessage -> Message, ChatClientProtocol -> SupportsChatGetResponse

Simplify the public API by removing redundant 'Chat' prefix from core types:
- ChatAgent -> Agent
- RawChatAgent -> RawAgent
- ChatMessage -> Message
- ChatClientProtocol -> SupportsChatGetResponse

Also renamed internal WorkflowMessage (was Message in _runner_context) to avoid collision.

No backward compatibility aliases - this is a clean breaking change.

* [BREAKING] Rename Agent chat_client parameter to client

* Fix rebase issues: WorkflowMessage references and broken markdown links

* Fix formatting and lint issues from code quality checks

* Fix import ordering in workflow sample files

* fixed rebase

* Fix test failures: use WorkflowMessage and A2AMessage after ChatMessage→Message rename

- Replace Message(data=..., source_id=...) with WorkflowMessage(...) in workflow tests
- Fix isinstance check in A2A agent to use A2AMessage instead of Message
- Fix import in test_workflow_observability.py (Message→WorkflowMessage)

* Fix lint, fmt, and sample errors after ChatMessage→Message rename

- Auto-fix 70+ ruff lint issues across samples (ChatMessage→Message refs)
- Fix HostedVectorStoreContent→Content.from_hosted_vector_store in file search sample
- Fix _normalize_messages→normalize_messages in custom agent sample
- Fix context.terminate→raise MiddlewareTermination in middleware samples
- Fix with_update_hook→with_transform_hook in override middleware sample
- Add TOptions_co import back to custom_chat_client sample
- Add noqa for FastAPI File() default in chatkit sample
- Fix B023 loop variable capture in weather agent sample

* fix: update Agent constructor calls from chat_client to client in declaration-only tool tests

* fix: add register_cleanup to devui lazy-loading proxy and type stub

* fixed tests and updated new pieces

* fix agui typevar

* fix merge errors

* fix merge conflicts

* fiux merge

* Remove unused links

---------

Co-authored-by: Evan Mattson <evan.mattson@microsoft.com>
2026-02-10 23:04:32 +00:00

69 lines
2.2 KiB
Python

# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
from typing import Any
import pytest
from agent_framework import Content, Message
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"}],
},
},
}
async def test_get_response_invokes_bedrock_runtime() -> None:
stub = _StubBedrockRuntime()
client = BedrockChatClient(
model_id="amazon.titan-text",
region="us-west-2",
client=stub,
)
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_id="amazon.titan-text",
region="us-west-2",
client=_StubBedrockRuntime(),
)
messages = [Message(role="system", contents=[Content.from_text(text="Only system text")])]
with pytest.raises(ServiceInitializationError):
client._prepare_options(messages, {})