agent/agent-framework
2
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Python: [BREAKING] Simplify API: ChatAgent -> Agent, ChatMessage -> Message (#3747)

Browse Source 
* [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>
This commit is contained in:
Eduard van Valkenburg
2026-02-10 23:04:32 +00:00
committed by GitHub
co-authored by Evan Mattson
parent a4c9e43afb
commit 0521f5bed8
418 changed files with 5385 additions and 5389 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/a2a/agent_framework_a2a/_agent.py
+19 -19
View File
@@ -18,7 +18,6 @@ from a2a.types import (
FilePart,
FileWithBytes,
FileWithUri,
Message,
Task,
TaskIdParams,
TaskQueryParams,
@@ -34,9 +33,9 @@ from agent_framework import (
AgentResponseUpdate,
AgentThread,
BaseAgent,
ChatMessage,
Content,
ContinuationToken,
Message,
ResponseStream,
normalize_messages,
prepend_agent_framework_to_user_agent,
@@ -83,7 +82,7 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
"""Agent2Agent (A2A) protocol implementation.
Wraps an A2A Client to connect the Agent Framework with external A2A-compliant agents
via HTTP/JSON-RPC. Converts framework ChatMessages to A2A Messages on send, and converts
via HTTP/JSON-RPC. Converts framework Messages to A2A Messages on send, and converts
A2A responses (Messages/Tasks) back to framework types. Inherits BaseAgent capabilities
while managing the underlying A2A protocol communication.
@@ -209,7 +208,7 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
@overload
def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[False] = ...,
thread: AgentThread | None = None,
@@ -221,7 +220,7 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
@overload
def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[True],
thread: AgentThread | None = None,
@@ -232,7 +231,7 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
@@ -268,7 +267,7 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
response = ResponseStream(
self._map_a2a_stream(a2a_stream, background=background),
finalizer=lambda updates: AgentResponse.from_updates(list(updates)),
finalizer=AgentResponse.from_updates,
)
if stream:
return response
@@ -291,7 +290,8 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
When True, they are yielded with a continuation token.
"""
async for item in a2a_stream:
if isinstance(item, Message):
if isinstance(item, A2AMessage):
# Process A2A Message
contents = self._parse_contents_from_a2a(item.parts)
yield AgentResponseUpdate(
contents=contents,
@@ -377,10 +377,10 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
return AgentResponse.from_updates(updates)
return AgentResponse(messages=[], response_id=task.id, raw_representation=task)
def _prepare_message_for_a2a(self, message: ChatMessage) -> A2AMessage:
"""Prepare a ChatMessage for the A2A protocol.
def _prepare_message_for_a2a(self, message: Message) -> A2AMessage:
"""Prepare a Message for the A2A protocol.
Transforms Agent Framework ChatMessage objects into A2A protocol Messages by:
Transforms Agent Framework Message objects into A2A protocol Messages by:
- Converting all message contents to appropriate A2A Part types
- Mapping text content to TextPart objects
- Converting file references (URI/data/hosted_file) to FilePart objects
@@ -389,7 +389,7 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
"""
parts: list[A2APart] = []
if not message.contents:
raise ValueError("ChatMessage.contents is empty; cannot convert to A2AMessage.")
raise ValueError("Message.contents is empty; cannot convert to A2AMessage.")
# Process ALL contents
for content in message.contents:
@@ -511,9 +511,9 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
raise ValueError(f"Unknown Part kind: {inner_part.kind}")
return contents
def _parse_messages_from_task(self, task: Task) -> list[ChatMessage]:
"""Parse A2A Task artifacts into ChatMessages with ASSISTANT role."""
messages: list[ChatMessage] = []
def _parse_messages_from_task(self, task: Task) -> list[Message]:
"""Parse A2A Task artifacts into Messages with ASSISTANT role."""
messages: list[Message] = []
if task.artifacts is not None:
for artifact in task.artifacts:
@@ -523,7 +523,7 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
history_item = task.history[-1]
contents = self._parse_contents_from_a2a(history_item.parts)
messages.append(
ChatMessage(
Message(
role="assistant" if history_item.role == A2ARole.agent else "user",
contents=contents,
raw_representation=history_item,
@@ -532,10 +532,10 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
return messages
def _parse_message_from_artifact(self, artifact: Artifact) -> ChatMessage:
"""Parse A2A Artifact into ChatMessage using part contents."""
def _parse_message_from_artifact(self, artifact: Artifact) -> Message:
"""Parse A2A Artifact into Message using part contents."""
contents = self._parse_contents_from_a2a(artifact.parts)
return ChatMessage(
return Message(
role="assistant",
contents=contents,
raw_representation=artifact,
python/packages/a2a/tests/test_a2a_agent.py
+17 -17
View File
@@ -12,19 +12,19 @@ from a2a.types import (
DataPart,
FilePart,
FileWithUri,
Message,
Part,
Task,
TaskState,
TaskStatus,
TextPart,
)
from a2a.types import Message as A2AMessage
from a2a.types import Role as A2ARole
from agent_framework import (
AgentResponse,
AgentResponseUpdate,
ChatMessage,
Content,
Message,
)
from agent_framework.a2a import A2AAgent
from pytest import fixture, raises
@@ -49,7 +49,7 @@ class MockA2AClient:
text_part = Part(root=TextPart(text=text))
# Create actual Message instance
message = Message(
message = A2AMessage(
message_id=message_id, role=A2ARole.agent if role == "agent" else A2ARole.user, parts=[text_part]
)
self.responses.append(message)
@@ -281,7 +281,7 @@ def test_parse_message_from_artifact(a2a_agent: A2AAgent) -> None:
result = a2a_agent._parse_message_from_artifact(artifact)
assert isinstance(result, ChatMessage)
assert isinstance(result, Message)
assert result.role == "assistant"
assert result.text == "Artifact content"
assert result.raw_representation == artifact
@@ -324,9 +324,9 @@ def test_parse_contents_from_a2a_conversion(a2a_agent: A2AAgent) -> None:
def test_prepare_message_for_a2a_with_error_content(a2a_agent: A2AAgent) -> None:
"""Test _prepare_message_for_a2a with ErrorContent."""
# Create ChatMessage with ErrorContent
# Create Message with ErrorContent
error_content = Content.from_error(message="Test error message")
message = ChatMessage(role="user", contents=[error_content])
message = Message(role="user", contents=[error_content])
# Convert to A2A message
a2a_message = a2a_agent._prepare_message_for_a2a(message)
@@ -339,9 +339,9 @@ def test_prepare_message_for_a2a_with_error_content(a2a_agent: A2AAgent) -> None
def test_prepare_message_for_a2a_with_uri_content(a2a_agent: A2AAgent) -> None:
"""Test _prepare_message_for_a2a with UriContent."""
# Create ChatMessage with UriContent
# Create Message with UriContent
uri_content = Content.from_uri(uri="http://example.com/file.pdf", media_type="application/pdf")
message = ChatMessage(role="user", contents=[uri_content])
message = Message(role="user", contents=[uri_content])
# Convert to A2A message
a2a_message = a2a_agent._prepare_message_for_a2a(message)
@@ -355,9 +355,9 @@ def test_prepare_message_for_a2a_with_uri_content(a2a_agent: A2AAgent) -> None:
def test_prepare_message_for_a2a_with_data_content(a2a_agent: A2AAgent) -> None:
"""Test _prepare_message_for_a2a with DataContent."""
# Create ChatMessage with DataContent (base64 data URI)
# Create Message with DataContent (base64 data URI)
data_content = Content.from_uri(uri="data:text/plain;base64,SGVsbG8gV29ybGQ=", media_type="text/plain")
message = ChatMessage(role="user", contents=[data_content])
message = Message(role="user", contents=[data_content])
# Convert to A2A message
a2a_message = a2a_agent._prepare_message_for_a2a(message)
@@ -370,11 +370,11 @@ def test_prepare_message_for_a2a_with_data_content(a2a_agent: A2AAgent) -> None:
def test_prepare_message_for_a2a_empty_contents_raises_error(a2a_agent: A2AAgent) -> None:
"""Test _prepare_message_for_a2a with empty contents raises ValueError."""
# Create ChatMessage with no contents
message = ChatMessage(role="user", contents=[])
# Create Message with no contents
message = Message(role="user", contents=[])
# Should raise ValueError for empty contents
with raises(ValueError, match="ChatMessage.contents is empty"):
with raises(ValueError, match="Message.contents is empty"):
a2a_agent._prepare_message_for_a2a(message)
@@ -432,12 +432,12 @@ async def test_context_manager_no_cleanup_when_no_http_client() -> None:
def test_prepare_message_for_a2a_with_multiple_contents() -> None:
"""Test conversion of ChatMessage with multiple contents."""
"""Test conversion of Message with multiple contents."""
agent = A2AAgent(client=MagicMock(), _http_client=None)
# Create message with multiple content types
message = ChatMessage(
message = Message(
role="user",
contents=[
Content.from_text(text="Here's the analysis:"),
@@ -489,12 +489,12 @@ def test_parse_contents_from_a2a_unknown_part_kind() -> None:
def test_prepare_message_for_a2a_with_hosted_file() -> None:
"""Test conversion of ChatMessage with HostedFileContent to A2A message."""
"""Test conversion of Message with HostedFileContent to A2A message."""
agent = A2AAgent(client=MagicMock(), _http_client=None)
# Create message with hosted file content
message = ChatMessage(
message = Message(
role="user",
contents=[Content.from_hosted_file(file_id="hosted://storage/document.pdf")],
)
python/packages/ag-ui/README.md
+6 -6
View File
@@ -14,15 +14,15 @@ pip install agent-framework-ag-ui
```python
from fastapi import FastAPI
from agent_framework import ChatAgent
from agent_framework import Agent
from agent_framework.azure import AzureOpenAIChatClient
from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
# Create your agent
agent = ChatAgent(
agent = Agent(
name="my_agent",
instructions="You are a helpful assistant.",
chat_client=AzureOpenAIChatClient(
client=AzureOpenAIChatClient(
endpoint="https://your-resource.openai.azure.com/",
deployment_name="gpt-4o-mini",
api_key="your-api-key",
@@ -58,7 +58,7 @@ The `AGUIChatClient` supports:
- Streaming and non-streaming responses
- Hybrid tool execution (client-side + server-side tools)
- Automatic thread management for conversation continuity
- Integration with `ChatAgent` for client-side history management
- Integration with `Agent` for client-side history management
## Documentation
@@ -91,7 +91,7 @@ The AG-UI endpoint does not enforce authentication by default. **For production
import os
from fastapi import Depends, FastAPI, HTTPException, Security
from fastapi.security import APIKeyHeader
from agent_framework import ChatAgent
from agent_framework import Agent
from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
# Configure API key authentication
@@ -104,7 +104,7 @@ async def verify_api_key(api_key: str | None = Security(API_KEY_HEADER)) -> None
raise HTTPException(status_code=401, detail="Invalid or missing API key")
# Create agent and app
agent = ChatAgent(name="my_agent", instructions="...", chat_client=...)
agent = Agent(name="my_agent", instructions="...", client=...)
app = FastAPI()
# Register endpoint WITH authentication
python/packages/ag-ui/agent_framework_ag_ui/_client.py
+18 -20
View File
@@ -15,11 +15,11 @@ from typing import TYPE_CHECKING, Any, Generic, TypedDict, cast
import httpx
from agent_framework import (
BaseChatClient,
ChatMessage,
ChatResponse,
ChatResponseUpdate,
Content,
FunctionTool,
Message,
ResponseStream,
)
from agent_framework._middleware import ChatMiddlewareLayer
@@ -69,10 +69,10 @@ AGUIChatOptionsT = TypeVar(
)
def _apply_server_function_call_unwrap(chat_client: BaseChatClientT) -> BaseChatClientT:
def _apply_server_function_call_unwrap(client: BaseChatClientT) -> BaseChatClientT:
"""Class decorator that unwraps server-side function calls after tool handling."""
original_get_response = chat_client.get_response
original_get_response = client.get_response
@wraps(original_get_response)
def response_wrapper(
@@ -105,8 +105,8 @@ def _apply_server_function_call_unwrap(chat_client: BaseChatClientT) -> BaseChat
_unwrap_server_function_call_contents(cast(MutableSequence[Content | dict[str, Any]], update.contents))
return update
chat_client.get_response = response_wrapper # type: ignore[assignment]
return chat_client
client.get_response = response_wrapper # type: ignore[assignment]
return client
@_apply_server_function_call_unwrap
@@ -130,8 +130,8 @@ class AGUIChatClient(
This client sends exactly the messages it receives to the server. It does NOT
automatically maintain conversation history. The server must handle history via thread_id.
For stateless servers: Use ChatAgent wrapper which will send full message history on each
request. However, even with ChatAgent, the server must echo back all context for the
For stateless servers: Use Agent wrapper which will send full message history on each
request. However, even with Agent, the server must echo back all context for the
agent to maintain history across turns.
Important: Tool Handling (Hybrid Execution - matches .NET)
@@ -140,7 +140,7 @@ class AGUIChatClient(
3. When LLM calls a client tool, function invocation executes it locally
4. Both client and server tools work together (hybrid pattern)
The wrapping ChatAgent's function invocation handles client tool execution
The wrapping Agent's function invocation handles client tool execution
automatically when the server's LLM decides to call them.
Examples:
@@ -162,18 +162,18 @@ class AGUIChatClient(
metadata={"thread_id": thread_id}
)
Recommended usage with ChatAgent (client manages history):
Recommended usage with Agent (client manages history):
.. code-block:: python
from agent_framework import ChatAgent
from agent_framework import Agent
from agent_framework.ag_ui import AGUIChatClient
client = AGUIChatClient(endpoint="http://localhost:8888/")
agent = ChatAgent(name="assistant", client=client)
agent = Agent(name="assistant", client=client)
thread = await agent.get_new_thread()
# ChatAgent automatically maintains history and sends full context
# Agent automatically maintains history and sends full context
response = await agent.run("Hello!", thread=thread)
response2 = await agent.run("How are you?", thread=thread)
@@ -282,9 +282,7 @@ class AGUIChatClient(
logger = get_logger()
logger.debug(f"[AGUIChatClient] Registered server placeholder: {tool_name}")
def _extract_state_from_messages(
self, messages: Sequence[ChatMessage]
) -> tuple[list[ChatMessage], dict[str, Any] | None]:
def _extract_state_from_messages(self, messages: Sequence[Message]) -> tuple[list[Message], dict[str, Any] | None]:
"""Extract state from last message if present.
Args:
@@ -319,11 +317,11 @@ class AGUIChatClient(
return list(messages), None
def _convert_messages_to_agui_format(self, messages: list[ChatMessage]) -> list[dict[str, Any]]:
def _convert_messages_to_agui_format(self, messages: list[Message]) -> list[dict[str, Any]]:
"""Convert Agent Framework messages to AG-UI format.
Args:
messages: List of ChatMessage objects
messages: List of Message objects
Returns:
List of AG-UI formatted message dictionaries
@@ -353,7 +351,7 @@ class AGUIChatClient(
def _inner_get_response(
self,
*,
messages: Sequence[ChatMessage],
messages: Sequence[Message],
stream: bool = False,
options: Mapping[str, Any],
**kwargs: Any,
@@ -393,7 +391,7 @@ class AGUIChatClient(
async def _streaming_impl(
self,
*,
messages: Sequence[ChatMessage],
messages: Sequence[Message],
options: Mapping[str, Any],
**kwargs: Any,
) -> AsyncIterable[ChatResponseUpdate]:
@@ -415,7 +413,7 @@ class AGUIChatClient(
agui_messages = self._convert_messages_to_agui_format(messages_to_send)
# Send client tools to server so LLM knows about them
# Client tools execute via ChatAgent's function invocation wrapper
# Client tools execute via Agent's function invocation wrapper
agui_tools = convert_tools_to_agui_format(options.get("tools"))
# Build set of client tool names (matches .NET clientToolSet)
python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py
+25 -25
View File
@@ -9,8 +9,8 @@ import logging
from typing import Any, cast
from agent_framework import (
ChatMessage,
Content,
Message,
prepare_function_call_results,
)
@@ -25,9 +25,9 @@ from ._utils import (
logger = logging.getLogger(__name__)
def _sanitize_tool_history(messages: list[ChatMessage]) -> list[ChatMessage]:
def _sanitize_tool_history(messages: list[Message]) -> list[Message]:
"""Normalize tool ordering and inject synthetic results for AG-UI edge cases."""
sanitized: list[ChatMessage] = []
sanitized: list[Message] = []
pending_tool_call_ids: set[str] | None = None
pending_confirm_changes_id: str | None = None
@@ -60,7 +60,7 @@ def _sanitize_tool_history(messages: list[ChatMessage]) -> list[ChatMessage]:
]
if filtered_contents:
# Create a new message without confirm_changes to avoid mutating the input
filtered_msg = ChatMessage(role=msg.role, contents=filtered_contents)
filtered_msg = Message(role=msg.role, contents=filtered_contents)
sanitized.append(filtered_msg)
# If no contents left after filtering, don't append anything
@@ -99,7 +99,7 @@ def _sanitize_tool_history(messages: list[ChatMessage]) -> list[ChatMessage]:
if pending_confirm_changes_id and approval_accepted is not None:
logger.info(f"Injecting synthetic tool result for confirm_changes call_id={pending_confirm_changes_id}")
synthetic_result = ChatMessage(
synthetic_result = Message(
role="tool",
contents=[
Content.from_function_result(
@@ -128,7 +128,7 @@ def _sanitize_tool_history(messages: list[ChatMessage]) -> list[ChatMessage]:
logger.info(
f"Injecting synthetic tool result for confirm_changes call_id={pending_confirm_changes_id}"
)
synthetic_result = ChatMessage(
synthetic_result = Message(
role="tool",
contents=[
Content.from_function_result(
@@ -152,7 +152,7 @@ def _sanitize_tool_history(messages: list[ChatMessage]) -> list[ChatMessage]:
)
for pending_call_id in pending_tool_call_ids:
logger.info(f"Injecting synthetic tool result for pending call_id={pending_call_id}")
synthetic_result = ChatMessage(
synthetic_result = Message(
role="tool",
contents=[
Content.from_function_result(
@@ -196,10 +196,10 @@ def _sanitize_tool_history(messages: list[ChatMessage]) -> list[ChatMessage]:
return sanitized
def _deduplicate_messages(messages: list[ChatMessage]) -> list[ChatMessage]:
def _deduplicate_messages(messages: list[Message]) -> list[Message]:
"""Remove duplicate messages while preserving order."""
seen_keys: dict[Any, int] = {}
unique_messages: list[ChatMessage] = []
unique_messages: list[Message] = []
for idx, msg in enumerate(messages):
role_value = get_role_value(msg)
@@ -256,7 +256,7 @@ def _deduplicate_messages(messages: list[ChatMessage]) -> list[ChatMessage]:
def normalize_agui_input_messages(
messages: list[dict[str, Any]],
) -> tuple[list[ChatMessage], list[dict[str, Any]]]:
) -> tuple[list[Message], list[dict[str, Any]]]:
"""Normalize raw AG-UI messages into provider and snapshot formats."""
provider_messages = agui_messages_to_agent_framework(messages)
provider_messages = _sanitize_tool_history(provider_messages)
@@ -265,14 +265,14 @@ def normalize_agui_input_messages(
return provider_messages, snapshot_messages
def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[ChatMessage]:
def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Message]:
"""Convert AG-UI messages to Agent Framework format.
Args:
messages: List of AG-UI messages
Returns:
List of Agent Framework ChatMessage objects
List of Agent Framework Message objects
"""
def _update_tool_call_arguments(
@@ -367,7 +367,7 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
allowed_keys = set(original_args.keys())
return {key: value for key, value in modified_args.items() if key in allowed_keys}
result: list[ChatMessage] = []
result: list[Message] = []
for msg in messages:
# Handle standard tool result messages early (role="tool") to preserve provider invariants
# This path maps AG‑UI tool messages to function_result content with the correct tool_call_id
@@ -480,7 +480,7 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
merged_args["steps"] = merged_steps
state_args = merged_args
# Update the ChatMessage tool call with only enabled steps (for LLM context).
# Update the Message tool call with only enabled steps (for LLM context).
# The LLM should only see the steps that were actually approved/executed.
updated_args_for_llm = (
json.dumps(filtered_args)
@@ -510,14 +510,14 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
function_call=func_call_for_approval,
additional_properties={"ag_ui_state_args": state_args} if state_args else None,
)
chat_msg = ChatMessage(
chat_msg = Message(
role="user",
contents=[approval_response],
)
else:
# No matching function call found - this is likely a confirm_changes approval
# Keep the old behavior for backwards compatibility
chat_msg = ChatMessage(
chat_msg = Message(
role="user",
contents=[Content.from_text(text=approval_payload_text)],
additional_properties={"is_tool_result": True, "tool_call_id": str(tool_call_id or "")},
@@ -537,7 +537,7 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
func_result = result_content
else:
func_result = str(result_content)
chat_msg = ChatMessage(
chat_msg = Message(
role="tool",
contents=[Content.from_function_result(call_id=str(tool_call_id), result=func_result)],
)
@@ -553,7 +553,7 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
tool_call_id = msg.get("toolCallId") or msg.get("tool_call_id") or msg.get("actionExecutionId", "")
result_content = msg.get("result", msg.get("content", ""))
chat_msg = ChatMessage(
chat_msg = Message(
role="tool",
contents=[Content.from_function_result(call_id=str(tool_call_id), result=result_content)],
)
@@ -592,7 +592,7 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
arguments=arguments,
)
)
chat_msg = ChatMessage(role="assistant", contents=contents)
chat_msg = Message(role="assistant", contents=contents)
if "id" in msg:
chat_msg.message_id = msg["id"]
result.append(chat_msg)
@@ -622,14 +622,14 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
)
approval_contents.append(approval_response)
chat_msg = ChatMessage(role=role, contents=approval_contents) # type: ignore[call-overload]
chat_msg = Message(role=role, contents=approval_contents) # type: ignore[call-overload]
else:
# Regular text message
content = msg.get("content", "")
if isinstance(content, str):
chat_msg = ChatMessage(role=role, contents=[Content.from_text(text=content)]) # type: ignore[call-overload]
chat_msg = Message(role=role, contents=[Content.from_text(text=content)]) # type: ignore[call-overload]
else:
chat_msg = ChatMessage(role=role, contents=[Content.from_text(text=str(content))]) # type: ignore[call-overload]
chat_msg = Message(role=role, contents=[Content.from_text(text=str(content))]) # type: ignore[call-overload]
if "id" in msg:
chat_msg.message_id = msg["id"]
@@ -639,11 +639,11 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
return result
def agent_framework_messages_to_agui(messages: list[ChatMessage] | list[dict[str, Any]]) -> list[dict[str, Any]]:
def agent_framework_messages_to_agui(messages: list[Message] | list[dict[str, Any]]) -> list[dict[str, Any]]:
"""Convert Agent Framework messages to AG-UI format.
Args:
messages: List of Agent Framework ChatMessage objects or AG-UI dicts (already converted)
messages: List of Agent Framework Message objects or AG-UI dicts (already converted)
Returns:
List of AG-UI message dictionaries
@@ -672,7 +672,7 @@ def agent_framework_messages_to_agui(messages: list[ChatMessage] | list[dict[str
result.append(normalized_msg)
continue
# Convert ChatMessage to AG-UI format
# Convert Message to AG-UI format
role_value: str = msg.role if hasattr(msg.role, "value") else msg.role # type: ignore[assignment]
role = FRAMEWORK_TO_AGUI_ROLE.get(role_value, "user")
python/packages/ag-ui/agent_framework_ag_ui/_orchestration/_helpers.py
+4 -4
View File
@@ -13,8 +13,8 @@ import logging
from typing import Any
from agent_framework import (
ChatMessage,
Content,
Message,
)
from .._utils import get_role_value
@@ -22,7 +22,7 @@ from .._utils import get_role_value
logger = logging.getLogger(__name__)
def pending_tool_call_ids(messages: list[ChatMessage]) -> set[str]:
def pending_tool_call_ids(messages: list[Message]) -> set[str]:
"""Get IDs of tool calls without corresponding results.
Args:
@@ -42,7 +42,7 @@ def pending_tool_call_ids(messages: list[ChatMessage]) -> set[str]:
return pending_ids - resolved_ids
def is_state_context_message(message: ChatMessage) -> bool:
def is_state_context_message(message: Message) -> bool:
"""Check if a message is a state context system message.
Args:
@@ -178,7 +178,7 @@ def build_safe_metadata(thread_metadata: dict[str, Any] | None) -> dict[str, Any
return safe_metadata
def latest_approval_response(messages: list[ChatMessage]) -> Content | None:
def latest_approval_response(messages: list[Message]) -> Content | None:
"""Get the latest approval response from messages.
Args:
python/packages/ag-ui/agent_framework_ag_ui/_orchestration/_tooling.py
+8 -8
View File
@@ -39,7 +39,7 @@ def collect_server_tools(agent: SupportsAgentRun) -> list[Any]:
functions need to be included for tool execution during approval flows.
Args:
agent: Agent instance to collect tools from. Works with ChatAgent
agent: Agent instance to collect tools from. Works with Agent
or any agent with default_options and optional mcp_tools attributes.
Returns:
@@ -53,7 +53,7 @@ def collect_server_tools(agent: SupportsAgentRun) -> list[Any]:
tools_from_agent = default_options.get("tools") if isinstance(default_options, dict) else None
server_tools = list(tools_from_agent) if tools_from_agent else []
# Include functions from connected MCP tools (only available on ChatAgent)
# Include functions from connected MCP tools (only available on Agent)
mcp_tools = getattr(agent, "mcp_tools", None)
if mcp_tools:
server_tools.extend(_collect_mcp_tool_functions(mcp_tools))
@@ -70,19 +70,19 @@ def register_additional_client_tools(agent: SupportsAgentRun, client_tools: list
"""Register client tools as additional declaration-only tools to avoid server execution.
Args:
agent: Agent instance to register tools on. Works with ChatAgent
or any agent with a chat_client attribute.
agent: Agent instance to register tools on. Works with Agent
or any agent with a client attribute.
client_tools: List of client tools to register.
"""
if not client_tools:
return
chat_client = getattr(agent, "chat_client", None)
if chat_client is None:
client = getattr(agent, "client", None)
if client is None:
return
if isinstance(chat_client, BaseChatClient) and chat_client.function_invocation_configuration is not None: # type: ignore[attr-defined]
chat_client.function_invocation_configuration["additional_tools"] = client_tools # type: ignore[attr-defined]
if isinstance(client, BaseChatClient) and client.function_invocation_configuration is not None: # type: ignore[attr-defined]
client.function_invocation_configuration["additional_tools"] = client_tools # type: ignore[attr-defined]
logger.debug(f"[TOOLS] Registered {len(client_tools)} client tools as additional_tools (declaration-only)")
python/packages/ag-ui/agent_framework_ag_ui/_run.py
+15 -17
View File
@@ -28,8 +28,8 @@ from ag_ui.core import (
)
from agent_framework import (
AgentThread,
ChatMessage,
Content,
Message,
SupportsAgentRun,
prepare_function_call_results,
)
@@ -195,7 +195,7 @@ class FlowState:
def _create_state_context_message(
current_state: dict[str, Any],
state_schema: dict[str, Any],
) -> ChatMessage | None:
) -> Message | None:
"""Create a system message with current state context.
This injects the current state into the conversation so the model
@@ -206,13 +206,13 @@ def _create_state_context_message(
state_schema: The state schema (used to determine if injection is needed)
Returns:
ChatMessage with state context, or None if not needed
Message with state context, or None if not needed
"""
if not current_state or not state_schema:
return None
state_json = json.dumps(current_state, indent=2)
return ChatMessage(
return Message(
role="system",
contents=[
Content.from_text(
@@ -229,10 +229,10 @@ def _create_state_context_message(
def _inject_state_context(
messages: list[ChatMessage],
messages: list[Message],
current_state: dict[str, Any],
state_schema: dict[str, Any],
) -> list[ChatMessage]:
) -> list[Message]:
"""Inject state context message into messages if appropriate.
The state context is injected before the last user message to give
@@ -592,7 +592,7 @@ async def _resolve_approval_responses(
Args:
messages: List of messages (will be modified in place)
tools: List of available tools
agent: The agent instance (to get chat_client and config)
agent: The agent instance (to get client and config)
run_kwargs: Kwargs for tool execution
"""
fcc_todo = _collect_approval_responses(messages)
@@ -605,12 +605,10 @@ async def _resolve_approval_responses(
# Execute approved tool calls
if approved_responses and tools:
chat_client = getattr(agent, "chat_client", None)
config = normalize_function_invocation_configuration(
getattr(chat_client, "function_invocation_configuration", None)
)
client = getattr(agent, "client", None)
config = normalize_function_invocation_configuration(getattr(client, "function_invocation_configuration", None))
middleware_pipeline = FunctionMiddlewarePipeline(
*getattr(chat_client, "function_middleware", ()),
*getattr(client, "function_middleware", ()),
*run_kwargs.get("middleware", ()),
)
# Filter out AG-UI-specific kwargs that should not be passed to tool execution
@@ -672,7 +670,7 @@ def _convert_approval_results_to_tool_messages(messages: list[Any]) -> None:
This modifies the messages list in place.
Args:
messages: List of ChatMessage objects to process
messages: List of Message objects to process
"""
result: list[Any] = []
@@ -694,11 +692,11 @@ def _convert_approval_results_to_tool_messages(messages: list[Any]) -> None:
# Tool messages first (right after the preceding assistant message per OpenAI requirements)
for func_result in function_results:
result.append(ChatMessage(role="tool", contents=[func_result]))
result.append(Message(role="tool", contents=[func_result]))
# Then user message with remaining content (if any)
if other_contents:
result.append(ChatMessage(role=msg.role, contents=other_contents))
result.append(Message(role=msg.role, contents=other_contents))
messages[:] = result
@@ -793,9 +791,9 @@ async def run_agent_stream(
# Check for structured output mode (skip text content)
skip_text = False
response_format = None
from agent_framework import ChatAgent
from agent_framework import Agent
if isinstance(agent, ChatAgent):
if isinstance(agent, Agent):
response_format = agent.default_options.get("response_format")
skip_text = response_format is not None
python/packages/ag-ui/agent_framework_ag_ui_examples/README.md
+37 -37
View File
@@ -12,7 +12,7 @@ pip install agent-framework-ag-ui
### Using Example Agents with Any Chat Client
All example agents are factory functions that accept any `ChatClientProtocol`-compatible chat client:
All example agents are factory functions that accept any `SupportsChatGetResponse`-compatible chat client:
```python
from fastapi import FastAPI
@@ -38,15 +38,15 @@ add_agent_framework_fastapi_endpoint(app, weather_agent(openai_client), "/weathe
```python
from fastapi import FastAPI
from agent_framework import ChatAgent
from agent_framework import Agent
from agent_framework.azure import AzureOpenAIChatClient
from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
# Create your agent
agent = ChatAgent(
agent = Agent(
name="my_agent",
instructions="You are a helpful assistant.",
chat_client=AzureOpenAIChatClient(model_id="gpt-4o"),
client=AzureOpenAIChatClient(model_id="gpt-4o"),
)
# Create FastAPI app and add AG-UI endpoint
@@ -70,21 +70,21 @@ This integration supports all 7 AG-UI features:
## Examples
All example agents are implemented as **factory functions** that accept any chat client implementing `ChatClientProtocol`. This provides maximum flexibility to use Azure OpenAI, OpenAI, Anthropic, or any custom chat client implementation.
All example agents are implemented as **factory functions** that accept any chat client implementing `SupportsChatGetResponse`. This provides maximum flexibility to use Azure OpenAI, OpenAI, Anthropic, or any custom chat client implementation.
### Available Example Agents
Complete examples for all AG-UI features are available:
- `simple_agent(chat_client)` - Basic agentic chat (Feature 1)
- `weather_agent(chat_client)` - Backend tool rendering (Feature 2)
- `human_in_the_loop_agent(chat_client)` - Human-in-the-loop with step customization (Feature 3)
- `task_steps_agent_wrapped(chat_client)` - Agentic generative UI with step execution (Feature 4)
- `ui_generator_agent(chat_client)` - Tool-based generative UI (Feature 5)
- `recipe_agent(chat_client)` - Shared state management (Feature 6)
- `document_writer_agent(chat_client)` - Predictive state updates (Feature 7)
- `research_assistant_agent(chat_client)` - Research with progress events
- `task_planner_agent(chat_client)` - Task planning with approvals
- `simple_agent(client)` - Basic agentic chat (Feature 1)
- `weather_agent(client)` - Backend tool rendering (Feature 2)
- `human_in_the_loop_agent(client)` - Human-in-the-loop with step customization (Feature 3)
- `task_steps_agent_wrapped(client)` - Agentic generative UI with step execution (Feature 4)
- `ui_generator_agent(client)` - Tool-based generative UI (Feature 5)
- `recipe_agent(client)` - Shared state management (Feature 6)
- `document_writer_agent(client)` - Predictive state updates (Feature 7)
- `research_assistant_agent(client)` - Research with progress events
- `task_planner_agent(client)` - Task planning with approvals
### Using Example Agents
@@ -97,7 +97,7 @@ from agent_framework_ag_ui_examples.agents import (
recipe_agent,
)
# Create a chat client (use any ChatClientProtocol implementation)
# Create a chat client (use any SupportsChatGetResponse implementation)
azure_client = AzureOpenAIChatClient(model_id="gpt-4")
openai_client = OpenAIChatClient(model_id="gpt-4o")
@@ -150,16 +150,16 @@ from agent_framework_ag_ui_examples.agents import (
app = FastAPI(title="AG-UI Examples")
# Create a chat client (shared across all agents, or create individual ones)
chat_client = AzureOpenAIChatClient(model_id="gpt-4")
client = AzureOpenAIChatClient(model_id="gpt-4")
# Add all example endpoints
add_agent_framework_fastapi_endpoint(app, simple_agent(chat_client), "/agentic_chat")
add_agent_framework_fastapi_endpoint(app, weather_agent(chat_client), "/backend_tool_rendering")
add_agent_framework_fastapi_endpoint(app, human_in_the_loop_agent(chat_client), "/human_in_the_loop")
add_agent_framework_fastapi_endpoint(app, task_steps_agent_wrapped(chat_client), "/agentic_generative_ui") # type: ignore[arg-type]
add_agent_framework_fastapi_endpoint(app, ui_generator_agent(chat_client), "/tool_based_generative_ui")
add_agent_framework_fastapi_endpoint(app, recipe_agent(chat_client), "/shared_state")
add_agent_framework_fastapi_endpoint(app, document_writer_agent(chat_client), "/predictive_state_updates")
add_agent_framework_fastapi_endpoint(app, simple_agent(client), "/agentic_chat")
add_agent_framework_fastapi_endpoint(app, weather_agent(client), "/backend_tool_rendering")
add_agent_framework_fastapi_endpoint(app, human_in_the_loop_agent(client), "/human_in_the_loop")
add_agent_framework_fastapi_endpoint(app, task_steps_agent_wrapped(client), "/agentic_generative_ui") # type: ignore[arg-type]
add_agent_framework_fastapi_endpoint(app, ui_generator_agent(client), "/tool_based_generative_ui")
add_agent_framework_fastapi_endpoint(app, recipe_agent(client), "/shared_state")
add_agent_framework_fastapi_endpoint(app, document_writer_agent(client), "/predictive_state_updates")
```
## Architecture
@@ -187,8 +187,8 @@ The package uses a clean, orchestrator-based architecture:
You can create your own agent factories following the same pattern as the examples:
```python
from agent_framework import ChatAgent, tool
from agent_framework import ChatClientProtocol
from agent_framework import Agent, tool
from agent_framework import SupportsChatGetResponse
from agent_framework.ag_ui import AgentFrameworkAgent
@tool
@@ -196,19 +196,19 @@ def my_tool(param: str) -> str:
"""My custom tool."""
return f"Result: {param}"
def my_custom_agent(chat_client: ChatClientProtocol) -> AgentFrameworkAgent:
def my_custom_agent(client: SupportsChatGetResponse) -> AgentFrameworkAgent:
"""Create a custom agent with the specified chat client.
Args:
chat_client: The chat client to use for the agent
client: The chat client to use for the agent
Returns:
A configured AgentFrameworkAgent instance
"""
agent = ChatAgent(
agent = Agent(
name="my_custom_agent",
instructions="Custom instructions here",
chat_client=chat_client,
client=client,
tools=[my_tool],
)
@@ -220,8 +220,8 @@ def my_custom_agent(chat_client: ChatClientProtocol) -> AgentFrameworkAgent:
# Use it
from agent_framework.azure import AzureOpenAIChatClient
chat_client = AzureOpenAIChatClient()
agent = my_custom_agent(chat_client)
client = AzureOpenAIChatClient()
agent = my_custom_agent(client)
```
### Shared State
@@ -229,14 +229,14 @@ agent = my_custom_agent(chat_client)
State is injected as system messages and updated via predictive state updates:
```python
from agent_framework import ChatAgent
from agent_framework import Agent
from agent_framework.azure import AzureOpenAIChatClient
from agent_framework.ag_ui import AgentFrameworkAgent
# Create your agent
agent = ChatAgent(
agent = Agent(
name="recipe_agent",
chat_client=AzureOpenAIChatClient(model_id="gpt-4o"),
client=AzureOpenAIChatClient(model_id="gpt-4o"),
)
state_schema = {
@@ -266,14 +266,14 @@ wrapped_agent = AgentFrameworkAgent(
Predictive state updates automatically stream tool arguments as optimistic state updates:
```python
from agent_framework import ChatAgent
from agent_framework import Agent
from agent_framework.azure import AzureOpenAIChatClient
from agent_framework.ag_ui import AgentFrameworkAgent
# Create your agent
agent = ChatAgent(
agent = Agent(
name="document_writer",
chat_client=AzureOpenAIChatClient(model_id="gpt-4o"),
client=AzureOpenAIChatClient(model_id="gpt-4o"),
)
predict_state_config = {
python/packages/ag-ui/agent_framework_ag_ui_examples/agents/document_writer_agent.py
+5 -5
View File
@@ -4,7 +4,7 @@
from __future__ import annotations
from agent_framework import ChatAgent, ChatClientProtocol, tool
from agent_framework import Agent, SupportsChatGetResponse, tool
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -40,19 +40,19 @@ _DOCUMENT_WRITER_INSTRUCTIONS = (
)
def document_writer_agent(chat_client: ChatClientProtocol) -> AgentFrameworkAgent:
def document_writer_agent(client: SupportsChatGetResponse) -> AgentFrameworkAgent:
"""Create a document writer agent with predictive state updates.
Args:
chat_client: The chat client to use for the agent
client: The chat client to use for the agent
Returns:
A configured AgentFrameworkAgent instance with document writing capabilities
"""
agent = ChatAgent(
agent = Agent(
name="document_writer",
instructions=_DOCUMENT_WRITER_INSTRUCTIONS,
chat_client=chat_client,
client=client,
tools=[write_document],
)
python/packages/ag-ui/agent_framework_ag_ui_examples/agents/human_in_the_loop_agent.py
+6 -6
View File
@@ -5,7 +5,7 @@
from enum import Enum
from typing import Any
from agent_framework import ChatAgent, ChatClientProtocol, tool
from agent_framework import Agent, SupportsChatGetResponse, tool
from pydantic import BaseModel, Field
@@ -43,16 +43,16 @@ def generate_task_steps(steps: list[TaskStep]) -> str:
return f"Generated {len(steps)} execution steps for the task."
def human_in_the_loop_agent(chat_client: ChatClientProtocol[Any]) -> ChatAgent[Any]:
def human_in_the_loop_agent(client: SupportsChatGetResponse[Any]) -> Agent[Any]:
"""Create a human-in-the-loop agent using tool-based approach for predictive state.
Args:
chat_client: The chat client to use for the agent
client: The chat client to use for the agent
Returns:
A configured ChatAgent instance with human-in-the-loop capabilities
A configured Agent instance with human-in-the-loop capabilities
"""
return ChatAgent(
return Agent(
name="human_in_the_loop_agent",
instructions="""You are a helpful assistant that can perform any task by breaking it down into steps.
@@ -81,6 +81,6 @@ def human_in_the_loop_agent(chat_client: ChatClientProtocol[Any]) -> ChatAgent[A
After the user approves and the function executes, THEN provide a brief acknowledgment like:
"The plan has been created with X steps selected."
""",
chat_client=chat_client,
client=client,
tools=[generate_task_steps],
)
python/packages/ag-ui/agent_framework_ag_ui_examples/agents/recipe_agent.py
+5 -5
View File
@@ -7,7 +7,7 @@ from __future__ import annotations
from enum import Enum
from typing import Any
from agent_framework import ChatAgent, ChatClientProtocol, tool
from agent_framework import Agent, SupportsChatGetResponse, tool
from agent_framework.ag_ui import AgentFrameworkAgent
from pydantic import BaseModel, Field
@@ -104,19 +104,19 @@ _RECIPE_INSTRUCTIONS = """You are a helpful recipe assistant that creates and mo
"""
def recipe_agent(chat_client: ChatClientProtocol[Any]) -> AgentFrameworkAgent:
def recipe_agent(client: SupportsChatGetResponse[Any]) -> AgentFrameworkAgent:
"""Create a recipe agent with streaming state updates.
Args:
chat_client: The chat client to use for the agent
client: The chat client to use for the agent
Returns:
A configured AgentFrameworkAgent instance with recipe management
"""
agent = ChatAgent(
agent = Agent(
name="recipe_agent",
instructions=_RECIPE_INSTRUCTIONS,
chat_client=chat_client,
client=client,
tools=[update_recipe],
)
python/packages/ag-ui/agent_framework_ag_ui_examples/agents/research_assistant_agent.py
+5 -5
View File
@@ -5,7 +5,7 @@
import asyncio
from typing import Any
from agent_framework import ChatAgent, ChatClientProtocol, tool
from agent_framework import Agent, SupportsChatGetResponse, tool
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -88,19 +88,19 @@ _RESEARCH_ASSISTANT_INSTRUCTIONS = (
)
def research_assistant_agent(chat_client: ChatClientProtocol[Any]) -> AgentFrameworkAgent:
def research_assistant_agent(client: SupportsChatGetResponse[Any]) -> AgentFrameworkAgent:
"""Create a research assistant agent.
Args:
chat_client: The chat client to use for the agent
client: The chat client to use for the agent
Returns:
A configured AgentFrameworkAgent instance with research capabilities
"""
agent = ChatAgent(
agent = Agent(
name="research_assistant",
instructions=_RESEARCH_ASSISTANT_INSTRUCTIONS,
chat_client=chat_client,
client=client,
tools=[research_topic, create_presentation, analyze_data],
)
python/packages/ag-ui/agent_framework_ag_ui_examples/agents/simple_agent.py
+6 -6
View File
@@ -4,20 +4,20 @@
from typing import Any
from agent_framework import ChatAgent, ChatClientProtocol
from agent_framework import Agent, SupportsChatGetResponse
def simple_agent(chat_client: ChatClientProtocol[Any]) -> ChatAgent[Any]:
def simple_agent(client: SupportsChatGetResponse[Any]) -> Agent[Any]:
"""Create a simple chat agent.
Args:
chat_client: The chat client to use for the agent
client: The chat client to use for the agent
Returns:
A configured ChatAgent instance
A configured Agent instance
"""
return ChatAgent[Any](
return Agent[Any](
name="simple_chat_agent",
instructions="You are a helpful assistant. Be concise and friendly.",
chat_client=chat_client,
client=client,
)
python/packages/ag-ui/agent_framework_ag_ui_examples/agents/task_planner_agent.py
+5 -5
View File
@@ -4,7 +4,7 @@
from typing import Any
from agent_framework import ChatAgent, ChatClientProtocol, tool
from agent_framework import Agent, SupportsChatGetResponse, tool
from agent_framework.ag_ui import AgentFrameworkAgent
@@ -61,19 +61,19 @@ _TASK_PLANNER_INSTRUCTIONS = (
)
def task_planner_agent(chat_client: ChatClientProtocol[Any]) -> AgentFrameworkAgent:
def task_planner_agent(client: SupportsChatGetResponse[Any]) -> AgentFrameworkAgent:
"""Create a task planner agent with user approval for actions.
Args:
chat_client: The chat client to use for the agent
client: The chat client to use for the agent
Returns:
A configured AgentFrameworkAgent instance with task planning capabilities
"""
agent = ChatAgent(
agent = Agent(
name="task_planner",
instructions=_TASK_PLANNER_INSTRUCTIONS,
chat_client=chat_client,
client=client,
tools=[create_calendar_event, send_email, book_meeting_room],
)
python/packages/ag-ui/agent_framework_ag_ui_examples/agents/task_steps_agent.py
+15 -15
View File
@@ -20,7 +20,7 @@ from ag_ui.core import (
TextMessageStartEvent,
ToolCallStartEvent,
)
from agent_framework import ChatAgent, ChatClientProtocol, ChatMessage, Content, tool
from agent_framework import Agent, Content, Message, SupportsChatGetResponse, tool
from agent_framework.ag_ui import AgentFrameworkAgent
from pydantic import BaseModel, Field
@@ -54,16 +54,16 @@ def generate_task_steps(steps: list[TaskStep]) -> str:
return "Steps generated."
def _create_task_steps_agent(chat_client: ChatClientProtocol[Any]) -> AgentFrameworkAgent:
def _create_task_steps_agent(client: SupportsChatGetResponse[Any]) -> AgentFrameworkAgent:
"""Create the task steps agent using tool-based approach for streaming.
Args:
chat_client: The chat client to use for the agent
client: The chat client to use for the agent
Returns:
A configured AgentFrameworkAgent instance
"""
agent = ChatAgent[Any](
agent = Agent[Any](
name="task_steps_agent",
instructions="""You are a helpful assistant that breaks down tasks into actionable steps.
@@ -83,7 +83,7 @@ def _create_task_steps_agent(chat_client: ChatClientProtocol[Any]) -> AgentFrame
- "Installing platform"
- "Adding finishing touches"
""",
chat_client=chat_client,
client=client,
tools=[generate_task_steps],
)
@@ -220,30 +220,30 @@ class TaskStepsAgentWithExecution:
# Get the underlying chat agent and client
chat_agent = self._base_agent.agent # type: ignore
chat_client = chat_agent.chat_client # type: ignore
client = chat_agent.client # type: ignore
# Build messages for summary call
original_messages = input_data.get("messages", [])
# Convert to ChatMessage objects if needed
messages: list[ChatMessage] = []
# Convert to Message objects if needed
messages: list[Message] = []
for msg in original_messages:
if isinstance(msg, dict):
content_str = msg.get("content", "")
if isinstance(content_str, str):
messages.append(
ChatMessage(
Message(
role=msg.get("role", "user"),
contents=[Content.from_text(text=content_str)],
)
)
elif isinstance(msg, ChatMessage):
elif isinstance(msg, Message):
messages.append(msg)
# Add completion message
messages.append(
ChatMessage(
Message(
role="user",
contents=[
Content.from_text(
@@ -270,7 +270,7 @@ class TaskStepsAgentWithExecution:
# Stream completion
accumulated_text = ""
async for chunk in chat_client.get_response(messages=messages, stream=True):
async for chunk in client.get_response(messages=messages, stream=True):
# chunk is ChatResponseUpdate
if hasattr(chunk, "text") and chunk.text:
accumulated_text += chunk.text
@@ -332,14 +332,14 @@ class TaskStepsAgentWithExecution:
yield run_finished_event
def task_steps_agent_wrapped(chat_client: ChatClientProtocol[Any]) -> TaskStepsAgentWithExecution:
def task_steps_agent_wrapped(client: SupportsChatGetResponse[Any]) -> TaskStepsAgentWithExecution:
"""Create a task steps agent with execution simulation.
Args:
chat_client: The chat client to use for the agent
client: The chat client to use for the agent
Returns:
A wrapped agent instance with step execution simulation
"""
base_agent = _create_task_steps_agent(chat_client)
base_agent = _create_task_steps_agent(client)
return TaskStepsAgentWithExecution(base_agent)
python/packages/ag-ui/agent_framework_ag_ui_examples/agents/ui_generator_agent.py
+5 -5
View File
@@ -7,7 +7,7 @@ from __future__ import annotations
import sys
from typing import TYPE_CHECKING, Any, TypedDict
from agent_framework import ChatAgent, ChatClientProtocol, FunctionTool
from agent_framework import Agent, FunctionTool, SupportsChatGetResponse
from agent_framework.ag_ui import AgentFrameworkAgent
if sys.version_info >= (3, 13):
@@ -168,19 +168,19 @@ _UI_GENERATOR_INSTRUCTIONS = """You MUST use the provided tools to generate cont
OptionsT = TypeVar("OptionsT", bound=TypedDict, default="ChatOptions") # type: ignore[valid-type]
def ui_generator_agent(chat_client: ChatClientProtocol[OptionsT]) -> AgentFrameworkAgent:
def ui_generator_agent(client: SupportsChatGetResponse[OptionsT]) -> AgentFrameworkAgent:
"""Create a UI generator agent with custom React component rendering.
Args:
chat_client: The chat client to use for the agent
client: The chat client to use for the agent
Returns:
A configured AgentFrameworkAgent instance with UI generation capabilities
"""
agent = ChatAgent(
agent = Agent(
name="ui_generator",
instructions=_UI_GENERATOR_INSTRUCTIONS,
chat_client=chat_client,
client=client,
tools=[generate_haiku, create_chart, display_timeline, show_comparison_table],
# Force tool usage - the LLM MUST call a tool, cannot respond with plain text
default_options={"tool_choice": "required"}, # type: ignore
python/packages/ag-ui/agent_framework_ag_ui_examples/agents/weather_agent.py
+6 -6
View File
@@ -6,7 +6,7 @@ from __future__ import annotations
from typing import Any
from agent_framework import ChatAgent, ChatClientProtocol, tool
from agent_framework import Agent, SupportsChatGetResponse, tool
@tool
@@ -59,16 +59,16 @@ def get_forecast(location: str, days: int = 3) -> str:
return f"{days}-day forecast for {location}:\n" + "\n".join(forecast)
def weather_agent(chat_client: ChatClientProtocol[Any]) -> ChatAgent[Any]:
def weather_agent(client: SupportsChatGetResponse[Any]) -> Agent[Any]:
"""Create a weather agent with get_weather and get_forecast tools.
Args:
chat_client: The chat client to use for the agent
client: The chat client to use for the agent
Returns:
A configured ChatAgent instance with weather tools
A configured Agent instance with weather tools
"""
return ChatAgent[Any](
return Agent[Any](
name="weather_agent",
instructions=(
"You are a helpful weather assistant. "
@@ -76,6 +76,6 @@ def weather_agent(chat_client: ChatClientProtocol[Any]) -> ChatAgent[Any]:
"Always provide friendly and informative responses. "
"First return the weather result, and then return details about the forecast."
),
chat_client=chat_client,
client=client,
tools=[get_weather, get_forecast],
)
python/packages/ag-ui/agent_framework_ag_ui_examples/server/api/backend_tool_rendering.py
+3 -3
View File
@@ -4,7 +4,7 @@
from typing import Any, cast
from agent_framework._clients import ChatClientProtocol
from agent_framework._clients import SupportsChatGetResponse
from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
from agent_framework.azure import AzureOpenAIChatClient
from fastapi import FastAPI
@@ -19,10 +19,10 @@ def register_backend_tool_rendering(app: FastAPI) -> None:
app: The FastAPI application.
"""
# Create a chat client and call the factory function
chat_client = cast(ChatClientProtocol[Any], AzureOpenAIChatClient())
client = cast(SupportsChatGetResponse[Any], AzureOpenAIChatClient())
add_agent_framework_fastapi_endpoint(
app,
weather_agent(chat_client),
weather_agent(client),
"/backend_tool_rendering",
)
python/packages/ag-ui/agent_framework_ag_ui_examples/server/main.py
+10 -10
View File
@@ -10,7 +10,7 @@ from typing import cast
import uvicorn
from agent_framework import ChatOptions
from agent_framework._clients import ChatClientProtocol
from agent_framework._clients import SupportsChatGetResponse
from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
from agent_framework.anthropic import AnthropicClient
from agent_framework.azure import AzureOpenAIChatClient
@@ -67,43 +67,43 @@ app.add_middleware(
# Create a shared chat client for all agents
# You can use different chat clients for different agents if needed
# Set CHAT_CLIENT=anthropic to use Anthropic, defaults to Azure OpenAI
chat_client: ChatClientProtocol[ChatOptions] = cast(
ChatClientProtocol[ChatOptions],
client: SupportsChatGetResponse[ChatOptions] = cast(
SupportsChatGetResponse[ChatOptions],
AnthropicClient() if os.getenv("CHAT_CLIENT", "").lower() == "anthropic" else AzureOpenAIChatClient(),
)
# Agentic Chat - basic chat agent
add_agent_framework_fastapi_endpoint(
app=app,
agent=simple_agent(chat_client),
agent=simple_agent(client),
path="/agentic_chat",
)
# Backend Tool Rendering - agent with tools
add_agent_framework_fastapi_endpoint(
app=app,
agent=weather_agent(chat_client),
agent=weather_agent(client),
path="/backend_tool_rendering",
)
# Shared State - recipe agent with structured output
add_agent_framework_fastapi_endpoint(
app=app,
agent=recipe_agent(chat_client),
agent=recipe_agent(client),
path="/shared_state",
)
# Predictive State Updates - document writer with predictive state
add_agent_framework_fastapi_endpoint(
app=app,
agent=document_writer_agent(chat_client),
agent=document_writer_agent(client),
path="/predictive_state_updates",
)
# Human in the Loop - human-in-the-loop agent with step customization
add_agent_framework_fastapi_endpoint(
app=app,
agent=human_in_the_loop_agent(chat_client),
agent=human_in_the_loop_agent(client),
path="/human_in_the_loop",
state_schema={"steps": {"type": "array"}},
predict_state_config={"steps": {"tool": "generate_task_steps", "tool_argument": "steps"}},
@@ -112,14 +112,14 @@ add_agent_framework_fastapi_endpoint(
# Agentic Generative UI - task steps agent with streaming state updates
add_agent_framework_fastapi_endpoint(
app=app,
agent=task_steps_agent_wrapped(chat_client), # type: ignore[arg-type]
agent=task_steps_agent_wrapped(client), # type: ignore[arg-type]
path="/agentic_generative_ui",
)
# Tool-based Generative UI - UI generator with frontend-rendered tools
add_agent_framework_fastapi_endpoint(
app=app,
agent=ui_generator_agent(chat_client),
agent=ui_generator_agent(client),
path="/tool_based_generative_ui",
)
python/packages/ag-ui/getting_started/README.md
+14 -14
View File
@@ -35,9 +35,9 @@ python client_advanced.py
**Note:** This example shows direct `AGUIChatClient` usage. Tool execution and conversation continuity depend on server-side configuration and capabilities.
### ChatAgent Integration (`client_with_agent.py`)
### Agent Integration (`client_with_agent.py`)
Best practice example using `ChatAgent` wrapper with **AgentThread**
Best practice example using `Agent` wrapper with **AgentThread**
- **AgentThread** maintains conversation state
- Client-side conversation history management via `thread.message_store`
- **Hybrid tool execution**: client-side + server-side tools simultaneously
@@ -77,7 +77,7 @@ The AG-UI protocol supports two approaches to conversation history:
- Full message history sent with each request
- Works with any AG-UI server (stateful or stateless)
The `ChatAgent` wrapper (used in client_with_agent.py) collects messages from local storage and sends the full history to `AGUIChatClient`, which then forwards everything to the server.
The `Agent` wrapper (used in client_with_agent.py) collects messages from local storage and sends the full history to `AGUIChatClient`, which then forwards everything to the server.
### Tool/Function Calling
@@ -91,14 +91,14 @@ Client defines: Server defines:
User: "What's the weather in SF and what time is it?"
↓
ChatAgent sends: full history + tool definitions for get_weather, read_sensors
Agent sends: full history + tool definitions for get_weather, read_sensors
↓
Server LLM decides: "I need get_weather('SF') and get_current_time()"
↓
Server executes get_current_time() → "2025-11-11 14:30:00 UTC"
Server sends function call request → get_weather('SF')
↓
ChatAgent intercepts get_weather call → executes locally
Agent intercepts get_weather call → executes locally
↓
Client sends result → "Sunny, 72°F"
↓
@@ -110,7 +110,7 @@ Client receives final response
**How it works:**
1. **Client-Side Tools** (`client_with_agent.py`):
- Tools defined in ChatAgent's `tools` parameter execute locally
- Tools defined in Agent's `tools` parameter execute locally
- Tool metadata (name, description, schema) sent to server for planning
- When server requests client tool → client intercepts → executes locally → sends result
@@ -126,7 +126,7 @@ Client receives final response
- Client tools execute client-side
**Direct AGUIChatClient Usage** (client_advanced.py):
Even without ChatAgent wrapper, client-side tools work:
Even without Agent wrapper, client-side tools work:
- Tools passed in ChatOptions execute locally
- Server can also have its own tools
- Hybrid execution works automatically
@@ -184,7 +184,7 @@ Create a file named `server.py`:
import os
from agent_framework import ChatAgent
from agent_framework import Agent
from agent_framework.azure import AzureOpenAIChatClient
from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
from fastapi import FastAPI
@@ -202,10 +202,10 @@ if not api_key:
raise ValueError("AZURE_OPENAI_API_KEY environment variable is required")
# Create the AI agent
agent = ChatAgent(
agent = Agent(
name="AGUIAssistant",
instructions="You are a helpful assistant.",
chat_client=AzureOpenAIChatClient(
client=AzureOpenAIChatClient(
endpoint=endpoint,
deployment_name=deployment_name,
api_key=api_key,
@@ -227,7 +227,7 @@ if __name__ == "__main__":
### Key Concepts
- **`add_agent_framework_fastapi_endpoint`**: Registers the AG-UI endpoint with automatic request/response handling and SSE streaming
- **`ChatAgent`**: The agent that will handle incoming requests
- **`Agent`**: The agent that will handle incoming requests
- **FastAPI Integration**: Uses FastAPI's native async support for streaming responses
- **Instructions**: The agent is created with default instructions, which can be overridden by client messages
- **Configuration**: `AzureOpenAIChatClient` can read from environment variables (`AZURE_OPENAI_ENDPOINT`, `AZURE_OPENAI_CHAT_DEPLOYMENT_NAME`, `AZURE_OPENAI_API_KEY`) or accept parameters directly
@@ -236,10 +236,10 @@ if __name__ == "__main__":
```python
# No need to read environment variables manually
agent = ChatAgent(
agent = Agent(
name="AGUIAssistant",
instructions="You are a helpful assistant.",
chat_client=AzureOpenAIChatClient(), # Reads from environment automatically
client=AzureOpenAIChatClient(), # Reads from environment automatically
)
```
@@ -354,7 +354,7 @@ if __name__ == "__main__":
- **Thread Management**: Pass `thread_id` in metadata to maintain conversation context across requests
- **Streaming Responses**: Use `get_response(..., stream=True)` for real-time streaming or `get_response(..., stream=False)` for non-streaming
- **Context Manager**: Use `async with` for automatic cleanup of HTTP connections
- **Standard Interface**: Works with all Agent Framework patterns (ChatAgent, tools, etc.)
- **Standard Interface**: Works with all Agent Framework patterns (Agent, tools, etc.)
- **Hybrid Tool Execution**: Supports both client-side and server-side tools executing together in the same conversation
### Configure and Run the Client
python/packages/ag-ui/getting_started/client_advanced.py
+4 -4
View File
@@ -114,15 +114,15 @@ async def non_streaming_example(client: AGUIChatClient, thread_id: str | None =
async def tool_example(client: AGUIChatClient, thread_id: str | None = None):
"""Demonstrate sending tool definitions to the server.
IMPORTANT: When using AGUIChatClient directly (without ChatAgent wrapper):
IMPORTANT: When using AGUIChatClient directly (without Agent wrapper):
- Tools are sent as DEFINITIONS only
- No automatic client-side execution (no function invocation middleware)
- Server must have matching tool implementations to execute them
For CLIENT-SIDE tool execution (like .NET AGUIClient sample):
- Use ChatAgent wrapper with tools
- Use Agent wrapper with tools
- See client_with_agent.py for the hybrid pattern
- ChatAgent middleware intercepts and executes client tools locally
- Agent middleware intercepts and executes client tools locally
- Server can have its own tools that execute server-side
- Both client and server tools work together in same conversation
@@ -186,7 +186,7 @@ async def conversation_example(client: AGUIChatClient):
# Check if context was maintained
if "alice" not in response2.text.lower():
print("\n[Note: Server may not maintain thread context - consider using ChatAgent for history management]")
print("\n[Note: Server may not maintain thread context - consider using Agent for history management]")
# Third turn
print("\nUser: Can you also tell me what 10 * 5 is?\n")
python/packages/ag-ui/getting_started/client_with_agent.py
+8 -8
View File
@@ -1,6 +1,6 @@
# Copyright (c) Microsoft. All rights reserved.
"""Example showing ChatAgent with AGUIChatClient for hybrid tool execution.
"""Example showing Agent with AGUIChatClient for hybrid tool execution.
This demonstrates the HYBRID pattern matching .NET AGUIClient implementation:
@@ -24,7 +24,7 @@ import asyncio
import logging
import os
from agent_framework import ChatAgent, tool
from agent_framework import Agent, tool
from agent_framework.ag_ui import AGUIChatClient
# Enable debug logging
@@ -55,7 +55,7 @@ def get_weather(location: str) -> str:
async def main():
"""Demonstrate ChatAgent + AGUIChatClient hybrid tool execution.
"""Demonstrate Agent + AGUIChatClient hybrid tool execution.
This matches the .NET pattern from Program.cs where:
- AIAgent agent = chatClient.CreateAIAgent(tools: [...])
@@ -63,14 +63,14 @@ async def main():
- RunStreamingAsync(messages, thread)
Python equivalent:
- agent = ChatAgent(chat_client=AGUIChatClient(...), tools=[...])
- agent = Agent(client=AGUIChatClient(...), tools=[...])
- thread = agent.get_new_thread() # Creates thread with message_store
- agent.run(message, stream=True, thread=thread) # Thread accumulates history
"""
server_url = os.environ.get("AGUI_SERVER_URL", "http://127.0.0.1:5100/")
print("=" * 70)
print("ChatAgent + AGUIChatClient: Hybrid Tool Execution")
print("Agent + AGUIChatClient: Hybrid Tool Execution")
print("=" * 70)
print(f"\nServer: {server_url}")
print("\nThis example demonstrates:")
@@ -82,11 +82,11 @@ async def main():
try:
# Create remote client in async context manager
async with AGUIChatClient(endpoint=server_url) as remote_client:
# Wrap in ChatAgent for conversation history management
agent = ChatAgent(
# Wrap in Agent for conversation history management
agent = Agent(
name="remote_assistant",
instructions="You are a helpful assistant. Remember user information across the conversation.",
chat_client=remote_client,
client=remote_client,
tools=[get_weather],
)
python/packages/ag-ui/getting_started/server.py
+3 -3
View File
@@ -7,7 +7,7 @@ from __future__ import annotations
import logging
import os
from agent_framework import ChatAgent, tool
from agent_framework import Agent, tool
from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
from agent_framework.azure import AzureOpenAIChatClient
from dotenv import load_dotenv
@@ -116,10 +116,10 @@ def get_time_zone(location: str) -> str:
# The client will send get_weather tool metadata so the LLM knows about it,
# and the function invocation mixin on AGUIChatClient will execute it client-side.
# This matches the .NET AG-UI hybrid execution pattern.
agent = ChatAgent(
agent = Agent(
name="AGUIAssistant",
instructions="You are a helpful assistant. Use get_weather for weather and get_time_zone for time zones.",
chat_client=AzureOpenAIChatClient(
client=AzureOpenAIChatClient(
endpoint=endpoint,
deployment_name=deployment_name,
),
python/packages/ag-ui/tests/ag_ui/conftest.py
+17 -17
View File
@@ -13,13 +13,13 @@ from agent_framework import (
AgentResponseUpdate,
AgentThread,
BaseChatClient,
ChatClientProtocol,
ChatMessage,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
Content,
Message,
SupportsAgentRun,
SupportsChatGetResponse,
)
from agent_framework._clients import OptionsCoT
from agent_framework._middleware import ChatMiddlewareLayer
@@ -43,7 +43,7 @@ class StreamingChatClientStub(
BaseChatClient[OptionsCoT],
Generic[OptionsCoT],
):
"""Typed streaming stub that satisfies ChatClientProtocol."""
"""Typed streaming stub that satisfies SupportsChatGetResponse."""
def __init__(self, stream_fn: StreamFn, response_fn: ResponseFn | None = None) -> None:
super().__init__(function_middleware=[])
@@ -55,7 +55,7 @@ class StreamingChatClientStub(
@overload
def get_response(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage],
messages: str | Message | Sequence[str | Message],
*,
stream: Literal[False] = ...,
options: ChatOptions[Any],
@@ -65,7 +65,7 @@ class StreamingChatClientStub(
@overload
def get_response(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage],
messages: str | Message | Sequence[str | Message],
*,
stream: Literal[False] = ...,
options: OptionsCoT | ChatOptions[None] | None = ...,
@@ -75,7 +75,7 @@ class StreamingChatClientStub(
@overload
def get_response(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage],
messages: str | Message | Sequence[str | Message],
*,
stream: Literal[True],
options: OptionsCoT | ChatOptions[Any] | None = ...,
@@ -84,7 +84,7 @@ class StreamingChatClientStub(
def get_response(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage],
messages: str | Message | Sequence[str | Message],
*,
stream: bool = False,
options: OptionsCoT | ChatOptions[Any] | None = None,
@@ -106,7 +106,7 @@ class StreamingChatClientStub(
def _inner_get_response(
self,
*,
messages: Sequence[ChatMessage],
messages: Sequence[Message],
stream: bool = False,
options: Mapping[str, Any],
**kwargs: Any,
@@ -121,7 +121,7 @@ class StreamingChatClientStub(
return self._get_response_impl(messages, options, **kwargs)
async def _get_response_impl(
self, messages: Sequence[ChatMessage], options: Mapping[str, Any], **kwargs: Any
self, messages: Sequence[Message], options: Mapping[str, Any], **kwargs: Any
) -> ChatResponse:
"""Non-streaming implementation."""
if self._response_fn is not None:
@@ -132,7 +132,7 @@ class StreamingChatClientStub(
contents.extend(update.contents)
return ChatResponse(
messages=[ChatMessage(role="assistant", contents=contents)],
messages=[Message(role="assistant", contents=contents)],
response_id="stub-response",
)
@@ -141,7 +141,7 @@ def stream_from_updates(updates: list[ChatResponseUpdate]) -> StreamFn:
"""Create a stream function that yields from a static list of updates."""
async def _stream(
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
for update in updates:
yield update
@@ -159,7 +159,7 @@ class StubAgent(SupportsAgentRun):
agent_id: str = "stub-agent",
agent_name: str | None = "stub-agent",
default_options: Any | None = None,
chat_client: Any | None = None,
client: Any | None = None,
) -> None:
self.id = agent_id
self.name = agent_name
@@ -168,14 +168,14 @@ class StubAgent(SupportsAgentRun):
self.default_options: dict[str, Any] = (
default_options if isinstance(default_options, dict) else {"tools": None, "response_format": None}
)
self.chat_client = chat_client or SimpleNamespace(function_invocation_configuration=None)
self.client = client or SimpleNamespace(function_invocation_configuration=None)
self.messages_received: list[Any] = []
self.tools_received: list[Any] | None = None
@overload
def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[False] = ...,
thread: AgentThread | None = None,
@@ -185,7 +185,7 @@ class StubAgent(SupportsAgentRun):
@overload
def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[True],
thread: AgentThread | None = None,
@@ -194,7 +194,7 @@ class StubAgent(SupportsAgentRun):
def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
@@ -226,7 +226,7 @@ class StubAgent(SupportsAgentRun):
@pytest.fixture
def streaming_chat_client_stub() -> type[ChatClientProtocol]:
def streaming_chat_client_stub() -> type[SupportsChatGetResponse]:
"""Return the StreamingChatClientStub class for creating test instances."""
return StreamingChatClientStub # type: ignore[return-value]
python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py
+18 -20
View File
@@ -7,11 +7,11 @@ from collections.abc import AsyncGenerator, Awaitable, MutableSequence
from typing import Any
from agent_framework import (
ChatMessage,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
Content,
Message,
ResponseStream,
tool,
)
@@ -29,13 +29,11 @@ class TestableAGUIChatClient(AGUIChatClient):
"""Expose http service for monkeypatching."""
return self._http_service
def extract_state_from_messages(
self, messages: list[ChatMessage]
) -> tuple[list[ChatMessage], dict[str, Any] | None]:
def extract_state_from_messages(self, messages: list[Message]) -> tuple[list[Message], dict[str, Any] | None]:
"""Expose state extraction helper."""
return self._extract_state_from_messages(messages)
def convert_messages_to_agui_format(self, messages: list[ChatMessage]) -> list[dict[str, Any]]:
def convert_messages_to_agui_format(self, messages: list[Message]) -> list[dict[str, Any]]:
"""Expose message conversion helper."""
return self._convert_messages_to_agui_format(messages)
@@ -44,7 +42,7 @@ class TestableAGUIChatClient(AGUIChatClient):
return self._get_thread_id(options)
def inner_get_response(
self, *, messages: MutableSequence[ChatMessage], options: dict[str, Any], stream: bool = False
self, *, messages: MutableSequence[Message], options: dict[str, Any], stream: bool = False
) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
"""Proxy to protected response call."""
return self._inner_get_response(messages=messages, options=options, stream=stream)
@@ -69,8 +67,8 @@ class TestAGUIChatClient:
"""Test state extraction when no state is present."""
client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
messages = [
ChatMessage(role="user", text="Hello"),
ChatMessage(role="assistant", text="Hi there"),
Message(role="user", text="Hello"),
Message(role="assistant", text="Hi there"),
]
result_messages, state = client.extract_state_from_messages(messages)
@@ -89,8 +87,8 @@ class TestAGUIChatClient:
state_b64 = base64.b64encode(state_json.encode("utf-8")).decode("utf-8")
messages = [
ChatMessage(role="user", text="Hello"),
ChatMessage(
Message(role="user", text="Hello"),
Message(
role="user",
contents=[Content.from_uri(uri=f"data:application/json;base64,{state_b64}")],
),
@@ -112,7 +110,7 @@ class TestAGUIChatClient:
state_b64 = base64.b64encode(invalid_json.encode("utf-8")).decode("utf-8")
messages = [
ChatMessage(
Message(
role="user",
contents=[Content.from_uri(uri=f"data:application/json;base64,{state_b64}")],
),
@@ -127,8 +125,8 @@ class TestAGUIChatClient:
"""Test message conversion to AG-UI format."""
client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
messages = [
ChatMessage(role="user", text="What is the weather?"),
ChatMessage(role="assistant", text="Let me check.", message_id="msg_123"),
Message(role="user", text="What is the weather?"),
Message(role="assistant", text="Let me check.", message_id="msg_123"),
]
agui_messages = client.convert_messages_to_agui_format(messages)
@@ -175,7 +173,7 @@ class TestAGUIChatClient:
client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
messages = [ChatMessage(role="user", text="Test message")]
messages = [Message(role="user", text="Test message")]
chat_options = ChatOptions()
updates: list[ChatResponseUpdate] = []
@@ -208,7 +206,7 @@ class TestAGUIChatClient:
client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
messages = [ChatMessage(role="user", text="Test message")]
messages = [Message(role="user", text="Test message")]
chat_options = {}
response = await client.inner_get_response(messages=messages, options=chat_options)
@@ -251,7 +249,7 @@ class TestAGUIChatClient:
client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
messages = [ChatMessage(role="user", text="Test with tools")]
messages = [Message(role="user", text="Test with tools")]
chat_options = ChatOptions(tools=[test_tool])
response = await client.inner_get_response(messages=messages, options=chat_options)
@@ -275,7 +273,7 @@ class TestAGUIChatClient:
client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
messages = [ChatMessage(role="user", text="Test server tool execution")]
messages = [Message(role="user", text="Test server tool execution")]
updates: list[ChatResponseUpdate] = []
async for update in client.get_response(messages, stream=True):
@@ -317,7 +315,7 @@ class TestAGUIChatClient:
client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
messages = [ChatMessage(role="user", text="Test server tool execution")]
messages = [Message(role="user", text="Test server tool execution")]
async for _ in client.get_response(
messages, stream=True, options={"tool_choice": "auto", "tools": [client_tool]}
@@ -333,8 +331,8 @@ class TestAGUIChatClient:
state_b64 = base64.b64encode(state_json.encode("utf-8")).decode("utf-8")
messages = [
ChatMessage(role="user", text="Hello"),
ChatMessage(
Message(role="user", text="Hello"),
Message(
role="user",
contents=[Content.from_uri(uri=f"data:application/json;base64,{state_b64}")],
),
python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py
+55 -55
View File
@@ -7,7 +7,7 @@ from collections.abc import AsyncIterator, MutableSequence
from typing import Any
import pytest
from agent_framework import ChatAgent, ChatMessage, ChatOptions, ChatResponseUpdate, Content
from agent_framework import Agent, ChatOptions, ChatResponseUpdate, Content, Message
from pydantic import BaseModel
@@ -16,12 +16,12 @@ async def test_agent_initialization_basic(streaming_chat_client_stub):
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
agent = ChatAgent[ChatOptions](
chat_client=streaming_chat_client_stub(stream_fn),
agent = Agent[ChatOptions](
client=streaming_chat_client_stub(stream_fn),
name="test_agent",
instructions="Test",
)
@@ -38,11 +38,11 @@ async def test_agent_initialization_with_state_schema(streaming_chat_client_stub
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
state_schema: dict[str, dict[str, Any]] = {"document": {"type": "string"}}
wrapper = AgentFrameworkAgent(agent=agent, state_schema=state_schema)
@@ -54,11 +54,11 @@ async def test_agent_initialization_with_predict_state_config(streaming_chat_cli
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
predict_config = {"document": {"tool": "write_doc", "tool_argument": "content"}}
wrapper = AgentFrameworkAgent(agent=agent, predict_state_config=predict_config)
@@ -70,7 +70,7 @@ async def test_agent_initialization_with_pydantic_state_schema(streaming_chat_cl
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
@@ -78,7 +78,7 @@ async def test_agent_initialization_with_pydantic_state_schema(streaming_chat_cl
document: str
tags: list[str] = []
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
wrapper_class_schema = AgentFrameworkAgent(agent=agent, state_schema=MyState)
wrapper_instance_schema = AgentFrameworkAgent(agent=agent, state_schema=MyState(document="hi"))
@@ -93,11 +93,11 @@ async def test_run_started_event_emission(streaming_chat_client_stub):
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent)
input_data = {"messages": [{"role": "user", "content": "Hi"}]}
@@ -117,11 +117,11 @@ async def test_predict_state_custom_event_emission(streaming_chat_client_stub):
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
predict_config = {
"document": {"tool": "write_doc", "tool_argument": "content"},
"summary": {"tool": "summarize", "tool_argument": "text"},
@@ -149,11 +149,11 @@ async def test_initial_state_snapshot_with_schema(streaming_chat_client_stub):
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
state_schema = {"document": {"type": "string"}}
wrapper = AgentFrameworkAgent(agent=agent, state_schema=state_schema)
@@ -179,11 +179,11 @@ async def test_state_initialization_object_type(streaming_chat_client_stub):
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
state_schema: dict[str, dict[str, Any]] = {"recipe": {"type": "object", "properties": {}}}
wrapper = AgentFrameworkAgent(agent=agent, state_schema=state_schema)
@@ -206,11 +206,11 @@ async def test_state_initialization_array_type(streaming_chat_client_stub):
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
state_schema: dict[str, dict[str, Any]] = {"steps": {"type": "array", "items": {}}}
wrapper = AgentFrameworkAgent(agent=agent, state_schema=state_schema)
@@ -233,11 +233,11 @@ async def test_run_finished_event_emission(streaming_chat_client_stub):
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent)
input_data = {"messages": [{"role": "user", "content": "Hi"}]}
@@ -255,11 +255,11 @@ async def test_tool_result_confirm_changes_accepted(streaming_chat_client_stub):
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text="Document updated")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
wrapper = AgentFrameworkAgent(
agent=agent,
state_schema={"document": {"type": "string"}},
@@ -302,11 +302,11 @@ async def test_tool_result_confirm_changes_rejected(streaming_chat_client_stub):
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text="OK")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent)
# Simulate tool result message with rejection
@@ -336,11 +336,11 @@ async def test_tool_result_function_approval_accepted(streaming_chat_client_stub
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text="OK")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent)
# Simulate tool result with multiple steps
@@ -382,11 +382,11 @@ async def test_tool_result_function_approval_rejected(streaming_chat_client_stub
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text="OK")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent)
# Simulate tool result rejection with steps
@@ -425,13 +425,13 @@ async def test_thread_metadata_tracking(streaming_chat_client_stub):
captured_options: dict[str, Any] = {}
async def stream_fn(
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
# Capture options to verify internal keys are NOT passed to chat client
captured_options.update(options)
yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent)
input_data = {
@@ -445,7 +445,7 @@ async def test_thread_metadata_tracking(streaming_chat_client_stub):
events.append(event)
# AG-UI internal metadata should be stored in thread.metadata
thread = agent.chat_client.last_thread
thread = agent.client.last_thread
thread_metadata = thread.metadata if thread and hasattr(thread, "metadata") else {}
assert thread_metadata.get("ag_ui_thread_id") == "test_thread_123"
assert thread_metadata.get("ag_ui_run_id") == "test_run_456"
@@ -467,13 +467,13 @@ async def test_state_context_injection(streaming_chat_client_stub):
captured_options: dict[str, Any] = {}
async def stream_fn(
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
# Capture options to verify internal keys are NOT passed to chat client
captured_options.update(options)
yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
wrapper = AgentFrameworkAgent(
agent=agent,
state_schema={"document": {"type": "string"}},
@@ -489,7 +489,7 @@ async def test_state_context_injection(streaming_chat_client_stub):
events.append(event)
# Current state should be stored in thread.metadata
thread = agent.chat_client.last_thread
thread = agent.client.last_thread
thread_metadata = thread.metadata if thread and hasattr(thread, "metadata") else {}
current_state = thread_metadata.get("current_state")
if isinstance(current_state, str):
@@ -506,11 +506,11 @@ async def test_no_messages_provided(streaming_chat_client_stub):
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent)
input_data: dict[str, Any] = {"messages": []}
@@ -530,11 +530,11 @@ async def test_message_end_event_emission(streaming_chat_client_stub):
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text="Hello world")])
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent)
input_data: dict[str, Any] = {"messages": [{"role": "user", "content": "Hi"}]}
@@ -558,13 +558,13 @@ async def test_error_handling_with_exception(streaming_chat_client_stub):
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
if False:
yield ChatResponseUpdate(contents=[])
raise RuntimeError("Simulated failure")
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent)
input_data: dict[str, Any] = {"messages": [{"role": "user", "content": "Hi"}]}
@@ -579,13 +579,13 @@ async def test_json_decode_error_in_tool_result(streaming_chat_client_stub):
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
if False:
yield ChatResponseUpdate(contents=[])
raise AssertionError("ChatClient should not be called with orphaned tool result")
agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent)
# Send invalid JSON as tool result without preceding tool call
@@ -618,13 +618,13 @@ async def test_agent_with_use_service_thread_is_false(streaming_chat_client_stub
request_service_thread_id: str | None = None
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
messages: MutableSequence[Message], chat_options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(
contents=[Content.from_text(text="Response")], response_id="resp_67890", conversation_id="conv_12345"
)
agent = ChatAgent(chat_client=streaming_chat_client_stub(stream_fn))
agent = Agent(client=streaming_chat_client_stub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent, use_service_thread=False)
input_data = {"messages": [{"role": "user", "content": "Hi"}], "thread_id": "conv_123456"}
@@ -642,7 +642,7 @@ async def test_agent_with_use_service_thread_is_true(streaming_chat_client_stub)
request_service_thread_id: str | None = None
async def stream_fn(
messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
messages: MutableSequence[Message], chat_options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
nonlocal request_service_thread_id
thread = kwargs.get("thread")
@@ -651,7 +651,7 @@ async def test_agent_with_use_service_thread_is_true(streaming_chat_client_stub)
contents=[Content.from_text(text="Response")], response_id="resp_67890", conversation_id="conv_12345"
)
agent = ChatAgent(chat_client=streaming_chat_client_stub(stream_fn))
agent = Agent(client=streaming_chat_client_stub(stream_fn))
wrapper = AgentFrameworkAgent(agent=agent, use_service_thread=True)
input_data = {"messages": [{"role": "user", "content": "Hi"}], "thread_id": "conv_123456"}
@@ -659,7 +659,7 @@ async def test_agent_with_use_service_thread_is_true(streaming_chat_client_stub)
events: list[Any] = []
async for event in wrapper.run_agent(input_data):
events.append(event)
request_service_thread_id = agent.chat_client.last_service_thread_id
request_service_thread_id = agent.client.last_service_thread_id
assert request_service_thread_id == "conv_123456" # type: ignore[attr-defined] (service_thread_id should be set)
@@ -679,15 +679,15 @@ async def test_function_approval_mode_executes_tool(streaming_chat_client_stub):
return "2025/12/01 12:00:00"
async def stream_fn(
messages: MutableSequence[ChatMessage], options: ChatOptions, **kwargs: Any
messages: MutableSequence[Message], options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
# Capture the messages received by the chat client
messages_received.clear()
messages_received.extend(messages)
yield ChatResponseUpdate(contents=[Content.from_text(text="Processing completed")])
agent = ChatAgent(
chat_client=streaming_chat_client_stub(stream_fn),
agent = Agent(
client=streaming_chat_client_stub(stream_fn),
name="test_agent",
instructions="Test",
tools=[get_datetime],
@@ -770,17 +770,17 @@ async def test_function_approval_mode_rejection(streaming_chat_client_stub):
return "All data deleted"
async def stream_fn(
messages: MutableSequence[ChatMessage], options: ChatOptions, **kwargs: Any
messages: MutableSequence[Message], options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
# Capture the messages received by the chat client
messages_received.clear()
messages_received.extend(messages)
yield ChatResponseUpdate(contents=[Content.from_text(text="Operation cancelled")])
agent = ChatAgent(
agent = Agent(
name="test_agent",
instructions="Test",
chat_client=streaming_chat_client_stub(stream_fn),
client=streaming_chat_client_stub(stream_fn),
tools=[delete_all_data],
)
wrapper = AgentFrameworkAgent(agent=agent)
python/packages/ag-ui/tests/ag_ui/test_endpoint.py
+24 -24
View File
@@ -5,7 +5,7 @@
import json
import pytest
from agent_framework import ChatAgent, ChatResponseUpdate, Content
from agent_framework import Agent, ChatResponseUpdate, Content
from fastapi import FastAPI, Header, HTTPException
from fastapi.params import Depends
from fastapi.testclient import TestClient
@@ -28,7 +28,7 @@ def build_chat_client(streaming_chat_client_stub, stream_from_updates_fixture):
async def test_add_endpoint_with_agent_protocol(build_chat_client):
"""Test adding endpoint with raw SupportsAgentRun."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
add_agent_framework_fastapi_endpoint(app, agent, path="/test-agent")
@@ -42,7 +42,7 @@ async def test_add_endpoint_with_agent_protocol(build_chat_client):
async def test_add_endpoint_with_wrapped_agent(build_chat_client):
"""Test adding endpoint with pre-wrapped AgentFrameworkAgent."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
wrapped_agent = AgentFrameworkAgent(agent=agent, name="wrapped")
add_agent_framework_fastapi_endpoint(app, wrapped_agent, path="/wrapped-agent")
@@ -57,7 +57,7 @@ async def test_add_endpoint_with_wrapped_agent(build_chat_client):
async def test_endpoint_with_state_schema(build_chat_client):
"""Test endpoint with state_schema parameter."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
state_schema = {"document": {"type": "string"}}
add_agent_framework_fastapi_endpoint(app, agent, path="/stateful", state_schema=state_schema)
@@ -73,7 +73,7 @@ async def test_endpoint_with_state_schema(build_chat_client):
async def test_endpoint_with_default_state_seed(build_chat_client):
"""Test endpoint seeds default state when client omits it."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
state_schema = {"proverbs": {"type": "array"}}
default_state = {"proverbs": ["Keep the original."]}
@@ -100,7 +100,7 @@ async def test_endpoint_with_default_state_seed(build_chat_client):
async def test_endpoint_with_predict_state_config(build_chat_client):
"""Test endpoint with predict_state_config parameter."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
predict_config = {"document": {"tool": "write_doc", "tool_argument": "content"}}
add_agent_framework_fastapi_endpoint(app, agent, path="/predictive", predict_state_config=predict_config)
@@ -114,7 +114,7 @@ async def test_endpoint_with_predict_state_config(build_chat_client):
async def test_endpoint_request_logging(build_chat_client):
"""Test that endpoint logs request details."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
add_agent_framework_fastapi_endpoint(app, agent, path="/logged")
@@ -134,7 +134,7 @@ async def test_endpoint_request_logging(build_chat_client):
async def test_endpoint_event_streaming(build_chat_client):
"""Test that endpoint streams events correctly."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client("Streamed response"))
agent = Agent(name="test", instructions="Test agent", client=build_chat_client("Streamed response"))
add_agent_framework_fastapi_endpoint(app, agent, path="/stream")
@@ -168,7 +168,7 @@ async def test_endpoint_event_streaming(build_chat_client):
async def test_endpoint_error_handling(build_chat_client):
"""Test endpoint error handling during request parsing."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
add_agent_framework_fastapi_endpoint(app, agent, path="/failing")
@@ -184,8 +184,8 @@ async def test_endpoint_error_handling(build_chat_client):
async def test_endpoint_multiple_paths(build_chat_client):
"""Test adding multiple endpoints with different paths."""
app = FastAPI()
agent1 = ChatAgent(name="agent1", instructions="First agent", chat_client=build_chat_client("Response 1"))
agent2 = ChatAgent(name="agent2", instructions="Second agent", chat_client=build_chat_client("Response 2"))
agent1 = Agent(name="agent1", instructions="First agent", client=build_chat_client("Response 1"))
agent2 = Agent(name="agent2", instructions="Second agent", client=build_chat_client("Response 2"))
add_agent_framework_fastapi_endpoint(app, agent1, path="/agent1")
add_agent_framework_fastapi_endpoint(app, agent2, path="/agent2")
@@ -202,7 +202,7 @@ async def test_endpoint_multiple_paths(build_chat_client):
async def test_endpoint_default_path(build_chat_client):
"""Test endpoint with default path."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
add_agent_framework_fastapi_endpoint(app, agent)
@@ -215,7 +215,7 @@ async def test_endpoint_default_path(build_chat_client):
async def test_endpoint_response_headers(build_chat_client):
"""Test that endpoint sets correct response headers."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
add_agent_framework_fastapi_endpoint(app, agent, path="/headers")
@@ -231,7 +231,7 @@ async def test_endpoint_response_headers(build_chat_client):
async def test_endpoint_empty_messages(build_chat_client):
"""Test endpoint with empty messages list."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
add_agent_framework_fastapi_endpoint(app, agent, path="/empty")
@@ -244,7 +244,7 @@ async def test_endpoint_empty_messages(build_chat_client):
async def test_endpoint_complex_input(build_chat_client):
"""Test endpoint with complex input data."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
add_agent_framework_fastapi_endpoint(app, agent, path="/complex")
@@ -269,7 +269,7 @@ async def test_endpoint_complex_input(build_chat_client):
async def test_endpoint_openapi_schema(build_chat_client):
"""Test that endpoint generates proper OpenAPI schema with request model."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
add_agent_framework_fastapi_endpoint(app, agent, path="/schema-test")
@@ -313,7 +313,7 @@ async def test_endpoint_openapi_schema(build_chat_client):
async def test_endpoint_default_tags(build_chat_client):
"""Test that endpoint uses default 'AG-UI' tag."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
add_agent_framework_fastapi_endpoint(app, agent, path="/default-tags")
@@ -331,7 +331,7 @@ async def test_endpoint_default_tags(build_chat_client):
async def test_endpoint_custom_tags(build_chat_client):
"""Test that endpoint accepts custom tags."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
add_agent_framework_fastapi_endpoint(app, agent, path="/custom-tags", tags=["Custom", "Agent"])
@@ -349,7 +349,7 @@ async def test_endpoint_custom_tags(build_chat_client):
async def test_endpoint_missing_required_field(build_chat_client):
"""Test that endpoint validates required fields with Pydantic."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
add_agent_framework_fastapi_endpoint(app, agent, path="/validation")
@@ -368,7 +368,7 @@ async def test_endpoint_internal_error_handling(build_chat_client):
from unittest.mock import patch
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
# Use default_state to trigger the code path that can raise an exception
add_agent_framework_fastapi_endpoint(app, agent, path="/error-test", default_state={"key": "value"})
@@ -387,7 +387,7 @@ async def test_endpoint_internal_error_handling(build_chat_client):
async def test_endpoint_with_dependencies_blocks_unauthorized(build_chat_client):
"""Test that endpoint blocks requests when authentication dependency fails."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
async def require_api_key(x_api_key: str | None = Header(None)):
if x_api_key != "secret-key":
@@ -406,7 +406,7 @@ async def test_endpoint_with_dependencies_blocks_unauthorized(build_chat_client)
async def test_endpoint_with_dependencies_allows_authorized(build_chat_client):
"""Test that endpoint allows requests when authentication dependency passes."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
async def require_api_key(x_api_key: str | None = Header(None)):
if x_api_key != "secret-key":
@@ -429,7 +429,7 @@ async def test_endpoint_with_dependencies_allows_authorized(build_chat_client):
async def test_endpoint_with_multiple_dependencies(build_chat_client):
"""Test that endpoint supports multiple dependencies."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
execution_order: list[str] = []
@@ -457,7 +457,7 @@ async def test_endpoint_with_multiple_dependencies(build_chat_client):
async def test_endpoint_without_dependencies_is_accessible(build_chat_client):
"""Test that endpoint without dependencies remains accessible (backward compatibility)."""
app = FastAPI()
agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
# No dependencies parameter - should be accessible without auth
add_agent_framework_fastapi_endpoint(app, agent, path="/open")
python/packages/ag-ui/tests/ag_ui/test_helpers.py
+15 -15
View File
@@ -2,7 +2,7 @@
"""Tests for orchestration helper functions."""
from agent_framework import ChatMessage, Content
from agent_framework import Content, Message
from agent_framework_ag_ui._orchestration._helpers import (
approval_steps,
@@ -29,8 +29,8 @@ class TestPendingToolCallIds:
def test_no_tool_calls(self):
"""Returns empty set when no tool calls in messages."""
messages = [
ChatMessage(role="user", contents=[Content.from_text("Hello")]),
ChatMessage(role="assistant", contents=[Content.from_text("Hi there")]),
Message(role="user", contents=[Content.from_text("Hello")]),
Message(role="assistant", contents=[Content.from_text("Hi there")]),
]
result = pending_tool_call_ids(messages)
assert result == set()
@@ -38,7 +38,7 @@ class TestPendingToolCallIds:
def test_pending_tool_call(self):
"""Returns pending tool call ID when no result exists."""
messages = [
ChatMessage(
Message(
role="assistant",
contents=[Content.from_function_call(call_id="call_123", name="get_weather", arguments="{}")],
),
@@ -49,11 +49,11 @@ class TestPendingToolCallIds:
def test_resolved_tool_call(self):
"""Returns empty set when tool call has result."""
messages = [
ChatMessage(
Message(
role="assistant",
contents=[Content.from_function_call(call_id="call_123", name="get_weather", arguments="{}")],
),
ChatMessage(
Message(
role="tool",
contents=[Content.from_function_result(call_id="call_123", result="sunny")],
),
@@ -64,7 +64,7 @@ class TestPendingToolCallIds:
def test_multiple_tool_calls_some_resolved(self):
"""Returns only unresolved tool call IDs."""
messages = [
ChatMessage(
Message(
role="assistant",
contents=[
Content.from_function_call(call_id="call_1", name="tool_a", arguments="{}"),
@@ -72,11 +72,11 @@ class TestPendingToolCallIds:
Content.from_function_call(call_id="call_3", name="tool_c", arguments="{}"),
],
),
ChatMessage(
Message(
role="tool",
contents=[Content.from_function_result(call_id="call_1", result="result_a")],
),
ChatMessage(
Message(
role="tool",
contents=[Content.from_function_result(call_id="call_3", result="result_c")],
),
@@ -90,7 +90,7 @@ class TestIsStateContextMessage:
def test_state_context_message(self):
"""Returns True for state context message."""
message = ChatMessage(
message = Message(
role="system",
contents=[Content.from_text("Current state of the application: {}")],
)
@@ -98,7 +98,7 @@ class TestIsStateContextMessage:
def test_non_system_message(self):
"""Returns False for non-system message."""
message = ChatMessage(
message = Message(
role="user",
contents=[Content.from_text("Current state of the application: {}")],
)
@@ -106,7 +106,7 @@ class TestIsStateContextMessage:
def test_system_message_without_state_prefix(self):
"""Returns False for system message without state prefix."""
message = ChatMessage(
message = Message(
role="system",
contents=[Content.from_text("You are a helpful assistant.")],
)
@@ -114,7 +114,7 @@ class TestIsStateContextMessage:
def test_empty_contents(self):
"""Returns False for message with empty contents."""
message = ChatMessage(role="system", contents=[])
message = Message(role="system", contents=[])
assert is_state_context_message(message) is False
@@ -342,7 +342,7 @@ class TestLatestApprovalResponse:
def test_no_approval_response(self):
"""Returns None when no approval response in last message."""
messages = [
ChatMessage(role="assistant", contents=[Content.from_text("Hello")]),
Message(role="assistant", contents=[Content.from_text("Hello")]),
]
result = latest_approval_response(messages)
assert result is None
@@ -357,7 +357,7 @@ class TestLatestApprovalResponse:
function_call=fc,
)
messages = [
ChatMessage(role="user", contents=[approval_content]),
Message(role="user", contents=[approval_content]),
]
result = latest_approval_response(messages)
assert result is approval_content
python/packages/ag-ui/tests/ag_ui/test_message_adapters.py
+13 -13
View File
@@ -5,7 +5,7 @@
import json
import pytest
from agent_framework import ChatMessage, Content
from agent_framework import Content, Message
from agent_framework_ag_ui._message_adapters import (
agent_framework_messages_to_agui,
@@ -24,7 +24,7 @@ def sample_agui_message():
@pytest.fixture
def sample_agent_framework_message():
"""Create a sample Agent Framework message."""
return ChatMessage(role="user", contents=[Content.from_text(text="Hello")], message_id="msg-123")
return Message(role="user", contents=[Content.from_text(text="Hello")], message_id="msg-123")
def test_agui_to_agent_framework_basic(sample_agui_message):
@@ -100,7 +100,7 @@ def test_agui_tool_result_to_agent_framework():
def test_agui_tool_approval_updates_tool_call_arguments():
"""Tool approval updates matching tool call arguments for snapshots and agent context.
The LLM context (ChatMessage) should contain only enabled steps, so the LLM
The LLM context (Message) should contain only enabled steps, so the LLM
generates responses based on what was actually approved/executed.
The raw messages (for MESSAGES_SNAPSHOT) should contain all steps with status,
@@ -446,7 +446,7 @@ def test_agui_with_tool_calls_to_agent_framework():
def test_agent_framework_to_agui_with_tool_calls():
"""Test converting Agent Framework message with tool calls to AG-UI."""
msg = ChatMessage(
msg = Message(
role="assistant",
contents=[
Content.from_text(text="Calling tool"),
@@ -471,7 +471,7 @@ def test_agent_framework_to_agui_with_tool_calls():
def test_agent_framework_to_agui_multiple_text_contents():
"""Test concatenating multiple text contents."""
msg = ChatMessage(
msg = Message(
role="assistant",
contents=[Content.from_text(text="Part 1 "), Content.from_text(text="Part 2")],
)
@@ -484,7 +484,7 @@ def test_agent_framework_to_agui_multiple_text_contents():
def test_agent_framework_to_agui_no_message_id():
"""Test message without message_id - should auto-generate ID."""
msg = ChatMessage(role="user", contents=[Content.from_text(text="Hello")])
msg = Message(role="user", contents=[Content.from_text(text="Hello")])
messages = agent_framework_messages_to_agui([msg])
@@ -496,7 +496,7 @@ def test_agent_framework_to_agui_no_message_id():
def test_agent_framework_to_agui_system_role():
"""Test system role conversion."""
msg = ChatMessage(role="system", contents=[Content.from_text(text="System")])
msg = Message(role="system", contents=[Content.from_text(text="System")])
messages = agent_framework_messages_to_agui([msg])
@@ -541,7 +541,7 @@ def test_extract_text_from_custom_contents():
def test_agent_framework_to_agui_function_result_dict():
"""Test converting FunctionResultContent with dict result to AG-UI."""
msg = ChatMessage(
msg = Message(
role="tool",
contents=[Content.from_function_result(call_id="call-123", result={"key": "value", "count": 42})],
message_id="msg-789",
@@ -558,7 +558,7 @@ def test_agent_framework_to_agui_function_result_dict():
def test_agent_framework_to_agui_function_result_none():
"""Test converting FunctionResultContent with None result to AG-UI."""
msg = ChatMessage(
msg = Message(
role="tool",
contents=[Content.from_function_result(call_id="call-123", result=None)],
message_id="msg-789",
@@ -574,7 +574,7 @@ def test_agent_framework_to_agui_function_result_none():
def test_agent_framework_to_agui_function_result_string():
"""Test converting FunctionResultContent with string result to AG-UI."""
msg = ChatMessage(
msg = Message(
role="tool",
contents=[Content.from_function_result(call_id="call-123", result="plain text result")],
message_id="msg-789",
@@ -589,7 +589,7 @@ def test_agent_framework_to_agui_function_result_string():
def test_agent_framework_to_agui_function_result_empty_list():
"""Test converting FunctionResultContent with empty list result to AG-UI."""
msg = ChatMessage(
msg = Message(
role="tool",
contents=[Content.from_function_result(call_id="call-123", result=[])],
message_id="msg-789",
@@ -611,7 +611,7 @@ def test_agent_framework_to_agui_function_result_single_text_content():
class MockTextContent:
text: str
msg = ChatMessage(
msg = Message(
role="tool",
contents=[Content.from_function_result(call_id="call-123", result=[MockTextContent("Hello from MCP!")])],
message_id="msg-789",
@@ -633,7 +633,7 @@ def test_agent_framework_to_agui_function_result_multiple_text_contents():
class MockTextContent:
text: str
msg = ChatMessage(
msg = Message(
role="tool",
contents=[
Content.from_function_result(
python/packages/ag-ui/tests/ag_ui/test_message_hygiene.py
+16 -16
View File
@@ -1,6 +1,6 @@
# Copyright (c) Microsoft. All rights reserved.
from agent_framework import ChatMessage, Content
from agent_framework import Content, Message
from agent_framework_ag_ui._message_adapters import _deduplicate_messages, _sanitize_tool_history
@@ -13,7 +13,7 @@ def test_sanitize_tool_history_filters_out_confirm_changes_only_message() -> Non
tool for the approval UI flow that shouldn't be sent to the LLM.
"""
messages = [
ChatMessage(
Message(
role="assistant",
contents=[
Content.from_function_call(
@@ -23,7 +23,7 @@ def test_sanitize_tool_history_filters_out_confirm_changes_only_message() -> Non
)
],
),
ChatMessage(
Message(
role="user",
contents=[Content.from_text(text='{"accepted": true}')],
),
@@ -44,11 +44,11 @@ def test_sanitize_tool_history_filters_out_confirm_changes_only_message() -> Non
def test_deduplicate_messages_prefers_non_empty_tool_results() -> None:
messages = [
ChatMessage(
Message(
role="tool",
contents=[Content.from_function_result(call_id="call1", result="")],
),
ChatMessage(
Message(
role="tool",
contents=[Content.from_function_result(call_id="call1", result="result data")],
),
@@ -71,13 +71,13 @@ def test_convert_approval_results_to_tool_messages() -> None:
# Simulate what happens after _resolve_approval_responses:
# A user message contains function_result content (the executed tool result)
messages = [
ChatMessage(
Message(
role="assistant",
contents=[
Content.from_function_call(call_id="call_123", name="my_mcp_tool", arguments="{}"),
],
),
ChatMessage(
Message(
role="user",
contents=[
Content.from_function_result(call_id="call_123", result="tool execution result"),
@@ -109,13 +109,13 @@ def test_convert_approval_results_preserves_other_user_content() -> None:
from agent_framework_ag_ui._run import _convert_approval_results_to_tool_messages
messages = [
ChatMessage(
Message(
role="assistant",
contents=[
Content.from_function_call(call_id="call_123", name="my_tool", arguments="{}"),
],
),
ChatMessage(
Message(
role="user",
contents=[
Content.from_text(text="User also said something"),
@@ -152,12 +152,12 @@ def test_sanitize_tool_history_filters_confirm_changes_keeps_other_tools() -> No
"""
messages = [
# User asks something
ChatMessage(
Message(
role="user",
contents=[Content.from_text(text="What time is it?")],
),
# Assistant calls MCP tool + confirm_changes
ChatMessage(
Message(
role="assistant",
contents=[
Content.from_function_call(call_id="call_1", name="get_datetime", arguments="{}"),
@@ -165,12 +165,12 @@ def test_sanitize_tool_history_filters_confirm_changes_keeps_other_tools() -> No
],
),
# Tool result for the actual MCP tool
ChatMessage(
Message(
role="tool",
contents=[Content.from_function_result(call_id="call_1", result="2024-01-01 12:00:00")],
),
# User asks something else
ChatMessage(
Message(
role="user",
contents=[Content.from_text(text="What's the date?")],
),
@@ -204,12 +204,12 @@ def test_sanitize_tool_history_filters_confirm_changes_from_assistant_messages()
respond with "Here's your 5-step plan" instead of "Here's your 2-step plan".
"""
messages = [
ChatMessage(
Message(
role="user",
contents=[Content.from_text(text="Build a robot")],
),
# Assistant message with both generate_task_steps and confirm_changes
ChatMessage(
Message(
role="assistant",
contents=[
Content.from_function_call(
@@ -225,7 +225,7 @@ def test_sanitize_tool_history_filters_confirm_changes_from_assistant_messages()
],
),
# Approval response
ChatMessage(
Message(
role="user",
contents=[
Content.from_function_approval_response(
python/packages/ag-ui/tests/ag_ui/test_run.py
+8 -8
View File
@@ -6,7 +6,7 @@ from ag_ui.core import (
TextMessageEndEvent,
TextMessageStartEvent,
)
from agent_framework import ChatMessage, Content
from agent_framework import Content, Message
from agent_framework_ag_ui._run import (
FlowState,
@@ -212,7 +212,7 @@ class TestInjectStateContext:
def test_no_state_message(self):
"""Returns original messages when no state context needed."""
messages = [ChatMessage(role="user", contents=[Content.from_text("Hello")])]
messages = [Message(role="user", contents=[Content.from_text("Hello")])]
result = _inject_state_context(messages, {}, {})
assert result == messages
@@ -224,8 +224,8 @@ class TestInjectStateContext:
def test_last_message_not_user(self):
"""Returns original messages when last message is not from user."""
messages = [
ChatMessage(role="user", contents=[Content.from_text("Hello")]),
ChatMessage(role="assistant", contents=[Content.from_text("Hi")]),
Message(role="user", contents=[Content.from_text("Hello")]),
Message(role="assistant", contents=[Content.from_text("Hi")]),
]
state = {"key": "value"}
schema = {"properties": {"key": {"type": "string"}}}
@@ -237,8 +237,8 @@ class TestInjectStateContext:
"""Injects state context before last user message."""
messages = [
ChatMessage(role="system", contents=[Content.from_text("You are helpful")]),
ChatMessage(role="user", contents=[Content.from_text("Hello")]),
Message(role="system", contents=[Content.from_text("You are helpful")]),
Message(role="user", contents=[Content.from_text("Hello")]),
]
state = {"document": "content"}
schema = {"properties": {"document": {"type": "string"}}}
@@ -405,7 +405,7 @@ def test_extract_approved_state_updates_no_handler():
"""Test _extract_approved_state_updates returns empty with no handler."""
from agent_framework_ag_ui._run import _extract_approved_state_updates
messages = [ChatMessage(role="user", contents=[Content.from_text("Hello")])]
messages = [Message(role="user", contents=[Content.from_text("Hello")])]
result = _extract_approved_state_updates(messages, None)
assert result == {}
@@ -416,7 +416,7 @@ def test_extract_approved_state_updates_no_approval():
from agent_framework_ag_ui._run import _extract_approved_state_updates
handler = PredictiveStateHandler(predict_state_config={"doc": {"tool": "write", "tool_argument": "content"}})
messages = [ChatMessage(role="user", contents=[Content.from_text("Hello")])]
messages = [Message(role="user", contents=[Content.from_text("Hello")])]
result = _extract_approved_state_updates(messages, handler)
assert result == {}
python/packages/ag-ui/tests/ag_ui/test_structured_output.py
+15 -15
View File
@@ -6,7 +6,7 @@ import json
from collections.abc import AsyncIterator, MutableSequence
from typing import Any
from agent_framework import ChatAgent, ChatMessage, ChatOptions, ChatResponseUpdate, Content
from agent_framework import Agent, ChatOptions, ChatResponseUpdate, Content, Message
from pydantic import BaseModel
@@ -35,13 +35,13 @@ async def test_structured_output_with_recipe(streaming_chat_client_stub, stream_
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], options: ChatOptions, **kwargs: Any
messages: MutableSequence[Message], options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(
contents=[Content.from_text(text='{"recipe": {"name": "Pasta"}, "message": "Here is your recipe"}')]
)
agent = ChatAgent(name="test", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
agent = Agent(name="test", instructions="Test", client=streaming_chat_client_stub(stream_fn))
agent.default_options = ChatOptions(response_format=RecipeOutput)
wrapper = AgentFrameworkAgent(
@@ -73,7 +73,7 @@ async def test_structured_output_with_steps(streaming_chat_client_stub, stream_f
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], options: ChatOptions, **kwargs: Any
messages: MutableSequence[Message], options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
steps_data = {
"steps": [
@@ -83,7 +83,7 @@ async def test_structured_output_with_steps(streaming_chat_client_stub, stream_f
}
yield ChatResponseUpdate(contents=[Content.from_text(text=json.dumps(steps_data))])
agent = ChatAgent(name="test", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
agent = Agent(name="test", instructions="Test", client=streaming_chat_client_stub(stream_fn))
agent.default_options = ChatOptions(response_format=StepsOutput)
wrapper = AgentFrameworkAgent(
@@ -116,8 +116,8 @@ async def test_structured_output_with_no_schema_match(streaming_chat_client_stub
ChatResponseUpdate(contents=[Content.from_text(text='{"data": {"key": "value"}}')]),
]
agent = ChatAgent(
name="test", instructions="Test", chat_client=streaming_chat_client_stub(stream_from_updates_fixture(updates))
agent = Agent(
name="test", instructions="Test", client=streaming_chat_client_stub(stream_from_updates_fixture(updates))
)
agent.default_options = ChatOptions(response_format=GenericOutput)
@@ -149,11 +149,11 @@ async def test_structured_output_without_schema(streaming_chat_client_stub, stre
info: str
async def stream_fn(
messages: MutableSequence[ChatMessage], options: ChatOptions, **kwargs: Any
messages: MutableSequence[Message], options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text='{"data": {"key": "value"}, "info": "processed"}')])
agent = ChatAgent(name="test", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
agent = Agent(name="test", instructions="Test", client=streaming_chat_client_stub(stream_fn))
agent.default_options = ChatOptions(response_format=DataOutput)
wrapper = AgentFrameworkAgent(
@@ -182,10 +182,10 @@ async def test_no_structured_output_when_no_response_format(streaming_chat_clien
updates = [ChatResponseUpdate(contents=[Content.from_text(text="Regular text")])]
agent = ChatAgent(
agent = Agent(
name="test",
instructions="Test",
chat_client=streaming_chat_client_stub(stream_from_updates_fixture(updates)),
client=streaming_chat_client_stub(stream_from_updates_fixture(updates)),
)
# No response_format set
@@ -208,12 +208,12 @@ async def test_structured_output_with_message_field(streaming_chat_client_stub,
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], options: ChatOptions, **kwargs: Any
messages: MutableSequence[Message], options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
output_data = {"recipe": {"name": "Salad"}, "message": "Fresh salad recipe ready"}
yield ChatResponseUpdate(contents=[Content.from_text(text=json.dumps(output_data))])
agent = ChatAgent(name="test", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
agent = Agent(name="test", instructions="Test", client=streaming_chat_client_stub(stream_fn))
agent.default_options = ChatOptions(response_format=RecipeOutput)
wrapper = AgentFrameworkAgent(
@@ -243,12 +243,12 @@ async def test_empty_updates_no_structured_processing(streaming_chat_client_stub
from agent_framework.ag_ui import AgentFrameworkAgent
async def stream_fn(
messages: MutableSequence[ChatMessage], options: ChatOptions, **kwargs: Any
messages: MutableSequence[Message], options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
if False:
yield ChatResponseUpdate(contents=[])
agent = ChatAgent(name="test", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
agent = Agent(name="test", instructions="Test", client=streaming_chat_client_stub(stream_fn))
agent.default_options = ChatOptions(response_format=RecipeOutput)
wrapper = AgentFrameworkAgent(agent=agent)
python/packages/ag-ui/tests/ag_ui/test_tooling.py
+7 -7
View File
@@ -2,7 +2,7 @@
from unittest.mock import MagicMock
from agent_framework import ChatAgent, tool
from agent_framework import Agent, tool
from agent_framework_ag_ui._orchestration._tooling import (
collect_server_tools,
@@ -31,14 +31,14 @@ def regular_tool() -> str:
return "result"
def _create_chat_agent_with_tool(tool_name: str = "regular_tool") -> ChatAgent:
"""Create a ChatAgent with a mocked chat client and a simple tool.
def _create_chat_agent_with_tool(tool_name: str = "regular_tool") -> Agent:
"""Create a Agent with a mocked chat client and a simple tool.
Note: tool_name parameter is kept for API compatibility but the tool
will always be named 'regular_tool' since tool uses the function name.
"""
mock_chat_client = MagicMock()
return ChatAgent(chat_client=mock_chat_client, tools=[regular_tool])
return Agent(client=mock_chat_client, tools=[regular_tool])
def test_merge_tools_filters_duplicates() -> None:
@@ -59,7 +59,7 @@ def test_register_additional_client_tools_assigns_when_configured() -> None:
mock_chat_client = MagicMock(spec=BaseChatClient)
mock_chat_client.function_invocation_configuration = normalize_function_invocation_configuration(None)
agent = ChatAgent(chat_client=mock_chat_client)
agent = Agent(client=mock_chat_client)
tools = [DummyTool("x")]
register_additional_client_tools(agent, tools)
@@ -148,14 +148,14 @@ def test_collect_server_tools_no_default_options() -> None:
def test_register_additional_client_tools_no_tools() -> None:
"""register_additional_client_tools does nothing with None tools."""
mock_chat_client = MagicMock()
agent = ChatAgent(chat_client=mock_chat_client)
agent = Agent(client=mock_chat_client)
# Should not raise
register_additional_client_tools(agent, None)
def test_register_additional_client_tools_no_chat_client() -> None:
"""register_additional_client_tools does nothing when agent has no chat_client."""
"""register_additional_client_tools does nothing when agent has no client."""
from agent_framework_ag_ui._orchestration._tooling import register_additional_client_tools
class MockAgent:
python/packages/ag-ui/tests/ag_ui/test_utils.py
+2 -2
View File
@@ -404,11 +404,11 @@ def test_safe_json_parse_with_none():
def test_get_role_value_with_enum():
"""Test get_role_value with enum role."""
from agent_framework import ChatMessage, Content
from agent_framework import Content, Message
from agent_framework_ag_ui._utils import get_role_value
message = ChatMessage(role="user", contents=[Content.from_text("test")])
message = Message(role="user", contents=[Content.from_text("test")])
result = get_role_value(message)
assert result == "user"
python/packages/anthropic/agent_framework_anthropic/_chat_client.py
+10 -10
View File
@@ -11,7 +11,6 @@ from agent_framework import (
Annotation,
BaseChatClient,
ChatAndFunctionMiddlewareTypes,
ChatMessage,
ChatMiddlewareLayer,
ChatOptions,
ChatResponse,
@@ -24,6 +23,7 @@ from agent_framework import (
HostedCodeInterpreterTool,
HostedMCPTool,
HostedWebSearchTool,
Message,
ResponseStream,
TextSpanRegion,
UsageDetails,
@@ -356,7 +356,7 @@ class AnthropicClient(
def _inner_get_response(
self,
*,
messages: Sequence[ChatMessage],
messages: Sequence[Message],
options: Mapping[str, Any],
stream: bool = False,
**kwargs: Any,
@@ -385,7 +385,7 @@ class AnthropicClient(
def _prepare_options(
self,
messages: Sequence[ChatMessage],
messages: Sequence[Message],
options: Mapping[str, Any],
**kwargs: Any,
) -> dict[str, Any]:
@@ -430,7 +430,7 @@ class AnthropicClient(
run_options["messages"] = self._prepare_messages_for_anthropic(messages)
# system message - first system message is passed as instructions
if messages and isinstance(messages[0], ChatMessage) and messages[0].role == "system":
if messages and isinstance(messages[0], Message) and messages[0].role == "system":
run_options["system"] = messages[0].text
# betas
@@ -516,22 +516,22 @@ class AnthropicClient(
"schema": schema,
}
def _prepare_messages_for_anthropic(self, messages: Sequence[ChatMessage]) -> list[dict[str, Any]]:
def _prepare_messages_for_anthropic(self, messages: Sequence[Message]) -> list[dict[str, Any]]:
"""Prepare a list of ChatMessages for the Anthropic client.
This skips the first message if it is a system message,
as Anthropic expects system instructions as a separate parameter.
"""
# first system message is passed as instructions
if messages and isinstance(messages[0], ChatMessage) and messages[0].role == "system":
if messages and isinstance(messages[0], Message) and messages[0].role == "system":
return [self._prepare_message_for_anthropic(msg) for msg in messages[1:]]
return [self._prepare_message_for_anthropic(msg) for msg in messages]
def _prepare_message_for_anthropic(self, message: ChatMessage) -> dict[str, Any]:
"""Prepare a ChatMessage for the Anthropic client.
def _prepare_message_for_anthropic(self, message: Message) -> dict[str, Any]:
"""Prepare a Message for the Anthropic client.
Args:
message: The ChatMessage to convert.
message: The Message to convert.
Returns:
A dictionary representing the message in Anthropic format.
@@ -693,7 +693,7 @@ class AnthropicClient(
return ChatResponse(
response_id=message.id,
messages=[
ChatMessage(
Message(
role="assistant",
contents=self._parse_contents_from_anthropic(message.content),
raw_representation=message,
python/packages/anthropic/tests/test_anthropic_client.py
+108 -108
View File
@@ -6,14 +6,14 @@ from unittest.mock import MagicMock, patch
import pytest
from agent_framework import (
ChatClientProtocol,
ChatMessage,
ChatOptions,
ChatResponseUpdate,
Content,
HostedCodeInterpreterTool,
HostedMCPTool,
HostedWebSearchTool,
Message,
SupportsChatGetResponse,
tool,
)
from agent_framework.exceptions import ServiceInitializationError
@@ -98,11 +98,11 @@ def test_anthropic_settings_missing_api_key(anthropic_unit_test_env: dict[str, s
def test_anthropic_client_init_with_client(mock_anthropic_client: MagicMock) -> None:
"""Test AnthropicClient initialization with existing anthropic_client."""
chat_client = create_test_anthropic_client(mock_anthropic_client, model_id="claude-3-5-sonnet-20241022")
client = create_test_anthropic_client(mock_anthropic_client, model_id="claude-3-5-sonnet-20241022")
assert chat_client.anthropic_client is mock_anthropic_client
assert chat_client.model_id == "claude-3-5-sonnet-20241022"
assert isinstance(chat_client, ChatClientProtocol)
assert client.anthropic_client is mock_anthropic_client
assert client.model_id == "claude-3-5-sonnet-20241022"
assert isinstance(client, SupportsChatGetResponse)
def test_anthropic_client_init_auto_create_client(anthropic_unit_test_env: dict[str, str]) -> None:
@@ -138,8 +138,8 @@ def test_anthropic_client_init_validation_error() -> None:
def test_anthropic_client_service_url(mock_anthropic_client: MagicMock) -> None:
"""Test service_url method."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
assert chat_client.service_url() == "https://api.anthropic.com"
client = create_test_anthropic_client(mock_anthropic_client)
assert client.service_url() == "https://api.anthropic.com"
# Message Conversion Tests
@@ -147,10 +147,10 @@ def test_anthropic_client_service_url(mock_anthropic_client: MagicMock) -> None:
def test_prepare_message_for_anthropic_text(mock_anthropic_client: MagicMock) -> None:
"""Test converting text message to Anthropic format."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
message = ChatMessage(role="user", text="Hello, world!")
client = create_test_anthropic_client(mock_anthropic_client)
message = Message(role="user", text="Hello, world!")
result = chat_client._prepare_message_for_anthropic(message)
result = client._prepare_message_for_anthropic(message)
assert result["role"] == "user"
assert len(result["content"]) == 1
@@ -160,8 +160,8 @@ def test_prepare_message_for_anthropic_text(mock_anthropic_client: MagicMock) ->
def test_prepare_message_for_anthropic_function_call(mock_anthropic_client: MagicMock) -> None:
"""Test converting function call message to Anthropic format."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
message = ChatMessage(
client = create_test_anthropic_client(mock_anthropic_client)
message = Message(
role="assistant",
contents=[
Content.from_function_call(
@@ -172,7 +172,7 @@ def test_prepare_message_for_anthropic_function_call(mock_anthropic_client: Magi
],
)
result = chat_client._prepare_message_for_anthropic(message)
result = client._prepare_message_for_anthropic(message)
assert result["role"] == "assistant"
assert len(result["content"]) == 1
@@ -184,8 +184,8 @@ def test_prepare_message_for_anthropic_function_call(mock_anthropic_client: Magi
def test_prepare_message_for_anthropic_function_result(mock_anthropic_client: MagicMock) -> None:
"""Test converting function result message to Anthropic format."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
message = ChatMessage(
client = create_test_anthropic_client(mock_anthropic_client)
message = Message(
role="tool",
contents=[
Content.from_function_result(
@@ -195,7 +195,7 @@ def test_prepare_message_for_anthropic_function_result(mock_anthropic_client: Ma
],
)
result = chat_client._prepare_message_for_anthropic(message)
result = client._prepare_message_for_anthropic(message)
assert result["role"] == "user"
assert len(result["content"]) == 1
@@ -209,13 +209,13 @@ def test_prepare_message_for_anthropic_function_result(mock_anthropic_client: Ma
def test_prepare_message_for_anthropic_text_reasoning(mock_anthropic_client: MagicMock) -> None:
"""Test converting text reasoning message to Anthropic format."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
message = ChatMessage(
client = create_test_anthropic_client(mock_anthropic_client)
message = Message(
role="assistant",
contents=[Content.from_text_reasoning(text="Let me think about this...")],
)
result = chat_client._prepare_message_for_anthropic(message)
result = client._prepare_message_for_anthropic(message)
assert result["role"] == "assistant"
assert len(result["content"]) == 1
@@ -225,13 +225,13 @@ def test_prepare_message_for_anthropic_text_reasoning(mock_anthropic_client: Mag
def test_prepare_messages_for_anthropic_with_system(mock_anthropic_client: MagicMock) -> None:
"""Test converting messages list with system message."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
client = create_test_anthropic_client(mock_anthropic_client)
messages = [
ChatMessage(role="system", text="You are a helpful assistant."),
ChatMessage(role="user", text="Hello!"),
Message(role="system", text="You are a helpful assistant."),
Message(role="user", text="Hello!"),
]
result = chat_client._prepare_messages_for_anthropic(messages)
result = client._prepare_messages_for_anthropic(messages)
# System message should be skipped
assert len(result) == 1
@@ -241,13 +241,13 @@ def test_prepare_messages_for_anthropic_with_system(mock_anthropic_client: Magic
def test_prepare_messages_for_anthropic_without_system(mock_anthropic_client: MagicMock) -> None:
"""Test converting messages list without system message."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
client = create_test_anthropic_client(mock_anthropic_client)
messages = [
ChatMessage(role="user", text="Hello!"),
ChatMessage(role="assistant", text="Hi there!"),
Message(role="user", text="Hello!"),
Message(role="assistant", text="Hi there!"),
]
result = chat_client._prepare_messages_for_anthropic(messages)
result = client._prepare_messages_for_anthropic(messages)
assert len(result) == 2
assert result[0]["role"] == "user"
@@ -259,7 +259,7 @@ def test_prepare_messages_for_anthropic_without_system(mock_anthropic_client: Ma
def test_prepare_tools_for_anthropic_tool(mock_anthropic_client: MagicMock) -> None:
"""Test converting FunctionTool to Anthropic format."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
client = create_test_anthropic_client(mock_anthropic_client)
@tool(approval_mode="never_require")
def get_weather(location: Annotated[str, Field(description="Location to get weather for")]) -> str:
@@ -267,7 +267,7 @@ def test_prepare_tools_for_anthropic_tool(mock_anthropic_client: MagicMock) -> N
return f"Weather for {location}"
chat_options = ChatOptions(tools=[get_weather])
result = chat_client._prepare_tools_for_anthropic(chat_options)
result = client._prepare_tools_for_anthropic(chat_options)
assert result is not None
assert "tools" in result
@@ -279,10 +279,10 @@ def test_prepare_tools_for_anthropic_tool(mock_anthropic_client: MagicMock) -> N
def test_prepare_tools_for_anthropic_web_search(mock_anthropic_client: MagicMock) -> None:
"""Test converting HostedWebSearchTool to Anthropic format."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
client = create_test_anthropic_client(mock_anthropic_client)
chat_options = ChatOptions(tools=[HostedWebSearchTool()])
result = chat_client._prepare_tools_for_anthropic(chat_options)
result = client._prepare_tools_for_anthropic(chat_options)
assert result is not None
assert "tools" in result
@@ -293,10 +293,10 @@ def test_prepare_tools_for_anthropic_web_search(mock_anthropic_client: MagicMock
def test_prepare_tools_for_anthropic_code_interpreter(mock_anthropic_client: MagicMock) -> None:
"""Test converting HostedCodeInterpreterTool to Anthropic format."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
client = create_test_anthropic_client(mock_anthropic_client)
chat_options = ChatOptions(tools=[HostedCodeInterpreterTool()])
result = chat_client._prepare_tools_for_anthropic(chat_options)
result = client._prepare_tools_for_anthropic(chat_options)
assert result is not None
assert "tools" in result
@@ -307,10 +307,10 @@ def test_prepare_tools_for_anthropic_code_interpreter(mock_anthropic_client: Mag
def test_prepare_tools_for_anthropic_mcp_tool(mock_anthropic_client: MagicMock) -> None:
"""Test converting HostedMCPTool to Anthropic format."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
client = create_test_anthropic_client(mock_anthropic_client)
chat_options = ChatOptions(tools=[HostedMCPTool(name="test-mcp", url="https://example.com/mcp")])
result = chat_client._prepare_tools_for_anthropic(chat_options)
result = client._prepare_tools_for_anthropic(chat_options)
assert result is not None
assert "mcp_servers" in result
@@ -322,7 +322,7 @@ def test_prepare_tools_for_anthropic_mcp_tool(mock_anthropic_client: MagicMock)
def test_prepare_tools_for_anthropic_mcp_with_auth(mock_anthropic_client: MagicMock) -> None:
"""Test converting HostedMCPTool with authorization headers."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
client = create_test_anthropic_client(mock_anthropic_client)
chat_options = ChatOptions(
tools=[
HostedMCPTool(
@@ -333,7 +333,7 @@ def test_prepare_tools_for_anthropic_mcp_with_auth(mock_anthropic_client: MagicM
]
)
result = chat_client._prepare_tools_for_anthropic(chat_options)
result = client._prepare_tools_for_anthropic(chat_options)
assert result is not None
assert "mcp_servers" in result
@@ -344,10 +344,10 @@ def test_prepare_tools_for_anthropic_mcp_with_auth(mock_anthropic_client: MagicM
def test_prepare_tools_for_anthropic_dict_tool(mock_anthropic_client: MagicMock) -> None:
"""Test converting dict tool to Anthropic format."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
client = create_test_anthropic_client(mock_anthropic_client)
chat_options = ChatOptions(tools=[{"type": "custom", "name": "custom_tool", "description": "A custom tool"}])
result = chat_client._prepare_tools_for_anthropic(chat_options)
result = client._prepare_tools_for_anthropic(chat_options)
assert result is not None
assert "tools" in result
@@ -357,10 +357,10 @@ def test_prepare_tools_for_anthropic_dict_tool(mock_anthropic_client: MagicMock)
def test_prepare_tools_for_anthropic_none(mock_anthropic_client: MagicMock) -> None:
"""Test converting None tools."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
client = create_test_anthropic_client(mock_anthropic_client)
chat_options = ChatOptions()
result = chat_client._prepare_tools_for_anthropic(chat_options)
result = client._prepare_tools_for_anthropic(chat_options)
assert result is None
@@ -370,14 +370,14 @@ def test_prepare_tools_for_anthropic_none(mock_anthropic_client: MagicMock) -> N
async def test_prepare_options_basic(mock_anthropic_client: MagicMock) -> None:
"""Test _prepare_options with basic ChatOptions."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
client = create_test_anthropic_client(mock_anthropic_client)
messages = [ChatMessage(role="user", text="Hello")]
messages = [Message(role="user", text="Hello")]
chat_options = ChatOptions(max_tokens=100, temperature=0.7)
run_options = chat_client._prepare_options(messages, chat_options)
run_options = client._prepare_options(messages, chat_options)
assert run_options["model"] == chat_client.model_id
assert run_options["model"] == client.model_id
assert run_options["max_tokens"] == 100
assert run_options["temperature"] == 0.7
assert "messages" in run_options
@@ -385,15 +385,15 @@ async def test_prepare_options_basic(mock_anthropic_client: MagicMock) -> None:
async def test_prepare_options_with_system_message(mock_anthropic_client: MagicMock) -> None:
"""Test _prepare_options with system message."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
client = create_test_anthropic_client(mock_anthropic_client)
messages = [
ChatMessage(role="system", text="You are helpful."),
ChatMessage(role="user", text="Hello"),
Message(role="system", text="You are helpful."),
Message(role="user", text="Hello"),
]
chat_options = ChatOptions()
run_options = chat_client._prepare_options(messages, chat_options)
run_options = client._prepare_options(messages, chat_options)
assert run_options["system"] == "You are helpful."
assert len(run_options["messages"]) == 1 # System message not in messages list
@@ -401,25 +401,25 @@ async def test_prepare_options_with_system_message(mock_anthropic_client: MagicM
async def test_prepare_options_with_tool_choice_auto(mock_anthropic_client: MagicMock) -> None:
"""Test _prepare_options with auto tool choice."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
client = create_test_anthropic_client(mock_anthropic_client)
messages = [ChatMessage(role="user", text="Hello")]
messages = [Message(role="user", text="Hello")]
chat_options = ChatOptions(tool_choice="auto")
run_options = chat_client._prepare_options(messages, chat_options)
run_options = client._prepare_options(messages, chat_options)
assert run_options["tool_choice"]["type"] == "auto"
async def test_prepare_options_with_tool_choice_required(mock_anthropic_client: MagicMock) -> None:
"""Test _prepare_options with required tool choice."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
client = create_test_anthropic_client(mock_anthropic_client)
messages = [ChatMessage(role="user", text="Hello")]
messages = [Message(role="user", text="Hello")]
# For required with specific function, need to pass as dict
chat_options = ChatOptions(tool_choice={"mode": "required", "required_function_name": "get_weather"})
run_options = chat_client._prepare_options(messages, chat_options)
run_options = client._prepare_options(messages, chat_options)
assert run_options["tool_choice"]["type"] == "tool"
assert run_options["tool_choice"]["name"] == "get_weather"
@@ -427,29 +427,29 @@ async def test_prepare_options_with_tool_choice_required(mock_anthropic_client:
async def test_prepare_options_with_tool_choice_none(mock_anthropic_client: MagicMock) -> None:
"""Test _prepare_options with none tool choice."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
client = create_test_anthropic_client(mock_anthropic_client)
messages = [ChatMessage(role="user", text="Hello")]
messages = [Message(role="user", text="Hello")]
chat_options = ChatOptions(tool_choice="none")
run_options = chat_client._prepare_options(messages, chat_options)
run_options = client._prepare_options(messages, chat_options)
assert run_options["tool_choice"]["type"] == "none"
async def test_prepare_options_with_tools(mock_anthropic_client: MagicMock) -> None:
"""Test _prepare_options with tools."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
client = create_test_anthropic_client(mock_anthropic_client)
@tool(approval_mode="never_require")
def get_weather(location: str) -> str:
"""Get weather for a location."""
return f"Weather for {location}"
messages = [ChatMessage(role="user", text="Hello")]
messages = [Message(role="user", text="Hello")]
chat_options = ChatOptions(tools=[get_weather])
run_options = chat_client._prepare_options(messages, chat_options)
run_options = client._prepare_options(messages, chat_options)
assert "tools" in run_options
assert len(run_options["tools"]) == 1
@@ -457,24 +457,24 @@ async def test_prepare_options_with_tools(mock_anthropic_client: MagicMock) -> N
async def test_prepare_options_with_stop_sequences(mock_anthropic_client: MagicMock) -> None:
"""Test _prepare_options with stop sequences."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
client = create_test_anthropic_client(mock_anthropic_client)
messages = [ChatMessage(role="user", text="Hello")]
messages = [Message(role="user", text="Hello")]
chat_options = ChatOptions(stop=["STOP", "END"])
run_options = chat_client._prepare_options(messages, chat_options)
run_options = client._prepare_options(messages, chat_options)
assert run_options["stop_sequences"] == ["STOP", "END"]
async def test_prepare_options_with_top_p(mock_anthropic_client: MagicMock) -> None:
"""Test _prepare_options with top_p."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
client = create_test_anthropic_client(mock_anthropic_client)
messages = [ChatMessage(role="user", text="Hello")]
messages = [Message(role="user", text="Hello")]
chat_options = ChatOptions(top_p=0.9)
run_options = chat_client._prepare_options(messages, chat_options)
run_options = client._prepare_options(messages, chat_options)
assert run_options["top_p"] == 0.9
@@ -485,9 +485,9 @@ async def test_prepare_options_filters_internal_kwargs(mock_anthropic_client: Ma
Internal kwargs like _function_middleware_pipeline, thread, and middleware
should be filtered out before being passed to the Anthropic API.
"""
chat_client = create_test_anthropic_client(mock_anthropic_client)
client = create_test_anthropic_client(mock_anthropic_client)
messages = [ChatMessage(role="user", text="Hello")]
messages = [Message(role="user", text="Hello")]
chat_options: ChatOptions = {}
# Simulate internal kwargs that get passed through the middleware pipeline
@@ -499,7 +499,7 @@ async def test_prepare_options_filters_internal_kwargs(mock_anthropic_client: Ma
"middleware": [object()],
}
run_options = chat_client._prepare_options(messages, chat_options, **internal_kwargs)
run_options = client._prepare_options(messages, chat_options, **internal_kwargs)
# Internal kwargs should be filtered out
assert "_function_middleware_pipeline" not in run_options
@@ -514,7 +514,7 @@ async def test_prepare_options_filters_internal_kwargs(mock_anthropic_client: Ma
def test_process_message_basic(mock_anthropic_client: MagicMock) -> None:
"""Test _process_message with basic text response."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
client = create_test_anthropic_client(mock_anthropic_client)
mock_message = MagicMock(spec=BetaMessage)
mock_message.id = "msg_123"
@@ -523,7 +523,7 @@ def test_process_message_basic(mock_anthropic_client: MagicMock) -> None:
mock_message.usage = BetaUsage(input_tokens=10, output_tokens=5)
mock_message.stop_reason = "end_turn"
response = chat_client._process_message(mock_message, {})
response = client._process_message(mock_message, {})
assert response.response_id == "msg_123"
assert response.model_id == "claude-3-5-sonnet-20241022"
@@ -540,7 +540,7 @@ def test_process_message_basic(mock_anthropic_client: MagicMock) -> None:
def test_process_message_with_tool_use(mock_anthropic_client: MagicMock) -> None:
"""Test _process_message with tool use."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
client = create_test_anthropic_client(mock_anthropic_client)
mock_message = MagicMock(spec=BetaMessage)
mock_message.id = "msg_123"
@@ -556,7 +556,7 @@ def test_process_message_with_tool_use(mock_anthropic_client: MagicMock) -> None
mock_message.usage = BetaUsage(input_tokens=10, output_tokens=5)
mock_message.stop_reason = "tool_use"
response = chat_client._process_message(mock_message, {})
response = client._process_message(mock_message, {})
assert len(response.messages[0].contents) == 1
assert response.messages[0].contents[0].type == "function_call"
@@ -567,10 +567,10 @@ def test_process_message_with_tool_use(mock_anthropic_client: MagicMock) -> None
def test_parse_usage_from_anthropic_basic(mock_anthropic_client: MagicMock) -> None:
"""Test _parse_usage_from_anthropic with basic usage."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
client = create_test_anthropic_client(mock_anthropic_client)
usage = BetaUsage(input_tokens=10, output_tokens=5)
result = chat_client._parse_usage_from_anthropic(usage)
result = client._parse_usage_from_anthropic(usage)
assert result is not None
assert result["input_token_count"] == 10
@@ -579,19 +579,19 @@ def test_parse_usage_from_anthropic_basic(mock_anthropic_client: MagicMock) -> N
def test_parse_usage_from_anthropic_none(mock_anthropic_client: MagicMock) -> None:
"""Test _parse_usage_from_anthropic with None usage."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
client = create_test_anthropic_client(mock_anthropic_client)
result = chat_client._parse_usage_from_anthropic(None)
result = client._parse_usage_from_anthropic(None)
assert result is None
def test_parse_contents_from_anthropic_text(mock_anthropic_client: MagicMock) -> None:
"""Test _parse_contents_from_anthropic with text content."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
client = create_test_anthropic_client(mock_anthropic_client)
content = [BetaTextBlock(type="text", text="Hello!")]
result = chat_client._parse_contents_from_anthropic(content)
result = client._parse_contents_from_anthropic(content)
assert len(result) == 1
assert result[0].type == "text"
@@ -600,7 +600,7 @@ def test_parse_contents_from_anthropic_text(mock_anthropic_client: MagicMock) ->
def test_parse_contents_from_anthropic_tool_use(mock_anthropic_client: MagicMock) -> None:
"""Test _parse_contents_from_anthropic with tool use."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
client = create_test_anthropic_client(mock_anthropic_client)
content = [
BetaToolUseBlock(
@@ -610,7 +610,7 @@ def test_parse_contents_from_anthropic_tool_use(mock_anthropic_client: MagicMock
input={"location": "SF"},
)
]
result = chat_client._parse_contents_from_anthropic(content)
result = client._parse_contents_from_anthropic(content)
assert len(result) == 1
assert result[0].type == "function_call"
@@ -625,7 +625,7 @@ def test_parse_contents_from_anthropic_input_json_delta_no_duplicate_name(mock_a
and subsequent input_json_delta events should have name="" to prevent
ag-ui from emitting duplicate ToolCallStartEvents.
"""
chat_client = create_test_anthropic_client(mock_anthropic_client)
client = create_test_anthropic_client(mock_anthropic_client)
# First, simulate a tool_use event that sets _last_call_id_name
tool_use_content = MagicMock()
@@ -634,7 +634,7 @@ def test_parse_contents_from_anthropic_input_json_delta_no_duplicate_name(mock_a
tool_use_content.name = "get_weather"
tool_use_content.input = {}
result = chat_client._parse_contents_from_anthropic([tool_use_content])
result = client._parse_contents_from_anthropic([tool_use_content])
assert len(result) == 1
assert result[0].type == "function_call"
assert result[0].call_id == "call_123"
@@ -645,7 +645,7 @@ def test_parse_contents_from_anthropic_input_json_delta_no_duplicate_name(mock_a
delta_content_1.type = "input_json_delta"
delta_content_1.partial_json = '{"location":'
result = chat_client._parse_contents_from_anthropic([delta_content_1])
result = client._parse_contents_from_anthropic([delta_content_1])
assert len(result) == 1
assert result[0].type == "function_call"
assert result[0].call_id == "call_123"
@@ -657,7 +657,7 @@ def test_parse_contents_from_anthropic_input_json_delta_no_duplicate_name(mock_a
delta_content_2.type = "input_json_delta"
delta_content_2.partial_json = '"San Francisco"}'
result = chat_client._parse_contents_from_anthropic([delta_content_2])
result = client._parse_contents_from_anthropic([delta_content_2])
assert len(result) == 1
assert result[0].type == "function_call"
assert result[0].call_id == "call_123"
@@ -670,13 +670,13 @@ def test_parse_contents_from_anthropic_input_json_delta_no_duplicate_name(mock_a
def test_process_stream_event_simple(mock_anthropic_client: MagicMock) -> None:
"""Test _process_stream_event with simple mock event."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
client = create_test_anthropic_client(mock_anthropic_client)
# Test with a basic mock event - the actual implementation will handle real events
mock_event = MagicMock()
mock_event.type = "message_stop"
result = chat_client._process_stream_event(mock_event)
result = client._process_stream_event(mock_event)
# message_stop events return None
assert result is None
@@ -684,7 +684,7 @@ def test_process_stream_event_simple(mock_anthropic_client: MagicMock) -> None:
async def test_inner_get_response(mock_anthropic_client: MagicMock) -> None:
"""Test _inner_get_response method."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
client = create_test_anthropic_client(mock_anthropic_client)
# Create a mock message response
mock_message = MagicMock(spec=BetaMessage)
@@ -696,10 +696,10 @@ async def test_inner_get_response(mock_anthropic_client: MagicMock) -> None:
mock_anthropic_client.beta.messages.create.return_value = mock_message
messages = [ChatMessage(role="user", text="Hi")]
messages = [Message(role="user", text="Hi")]
chat_options = ChatOptions(max_tokens=10)
response = await chat_client._inner_get_response( # type: ignore[attr-defined]
response = await client._inner_get_response( # type: ignore[attr-defined]
messages=messages, options=chat_options
)
@@ -710,7 +710,7 @@ async def test_inner_get_response(mock_anthropic_client: MagicMock) -> None:
async def test_inner_get_response_streaming(mock_anthropic_client: MagicMock) -> None:
"""Test _inner_get_response method with streaming."""
chat_client = create_test_anthropic_client(mock_anthropic_client)
client = create_test_anthropic_client(mock_anthropic_client)
# Create mock streaming response
async def mock_stream():
@@ -720,11 +720,11 @@ async def test_inner_get_response_streaming(mock_anthropic_client: MagicMock) ->
mock_anthropic_client.beta.messages.create.return_value = mock_stream()
messages = [ChatMessage(role="user", text="Hi")]
messages = [Message(role="user", text="Hi")]
chat_options = ChatOptions(max_tokens=10)
chunks: list[ChatResponseUpdate] = []
async for chunk in chat_client._inner_get_response( # type: ignore[attr-defined]
async for chunk in client._inner_get_response( # type: ignore[attr-defined]
messages=messages, options=chat_options, stream=True
):
if chunk:
@@ -751,7 +751,7 @@ async def test_anthropic_client_integration_basic_chat() -> None:
"""Integration test for basic chat completion."""
client = AnthropicClient()
messages = [ChatMessage(role="user", text="Say 'Hello, World!' and nothing else.")]
messages = [Message(role="user", text="Say 'Hello, World!' and nothing else.")]
response = await client.get_response(messages=messages, options={"max_tokens": 50})
@@ -768,7 +768,7 @@ async def test_anthropic_client_integration_streaming_chat() -> None:
"""Integration test for streaming chat completion."""
client = AnthropicClient()
messages = [ChatMessage(role="user", text="Count from 1 to 5.")]
messages = [Message(role="user", text="Count from 1 to 5.")]
chunks = []
async for chunk in client.get_response(messages=messages, stream=True, options={"max_tokens": 50}):
@@ -784,7 +784,7 @@ async def test_anthropic_client_integration_function_calling() -> None:
"""Integration test for function calling."""
client = AnthropicClient()
messages = [ChatMessage(role="user", text="What's the weather in San Francisco?")]
messages = [Message(role="user", text="What's the weather in San Francisco?")]
tools = [get_weather]
response = await client.get_response(
@@ -804,7 +804,7 @@ async def test_anthropic_client_integration_hosted_tools() -> None:
"""Integration test for hosted tools."""
client = AnthropicClient()
messages = [ChatMessage(role="user", text="What tools do you have available?")]
messages = [Message(role="user", text="What tools do you have available?")]
tools = [
HostedWebSearchTool(),
HostedCodeInterpreterTool(),
@@ -831,8 +831,8 @@ async def test_anthropic_client_integration_with_system_message() -> None:
client = AnthropicClient()
messages = [
ChatMessage(role="system", text="You are a pirate. Always respond like a pirate."),
ChatMessage(role="user", text="Hello!"),
Message(role="system", text="You are a pirate. Always respond like a pirate."),
Message(role="user", text="Hello!"),
]
response = await client.get_response(messages=messages, options={"max_tokens": 50})
@@ -847,7 +847,7 @@ async def test_anthropic_client_integration_temperature_control() -> None:
"""Integration test with temperature control."""
client = AnthropicClient()
messages = [ChatMessage(role="user", text="Say hello.")]
messages = [Message(role="user", text="Say hello.")]
response = await client.get_response(
messages=messages,
@@ -865,11 +865,11 @@ async def test_anthropic_client_integration_ordering() -> None:
client = AnthropicClient()
messages = [
ChatMessage(role="user", text="Say hello."),
ChatMessage(role="user", text="Then say goodbye."),
ChatMessage(role="assistant", text="Thank you for chatting!"),
ChatMessage(role="assistant", text="Let me know if I can help."),
ChatMessage(role="user", text="Just testing things."),
Message(role="user", text="Say hello."),
Message(role="user", text="Then say goodbye."),
Message(role="assistant", text="Thank you for chatting!"),
Message(role="assistant", text="Let me know if I can help."),
Message(role="user", text="Just testing things."),
]
response = await client.get_response(messages=messages)
@@ -890,7 +890,7 @@ async def test_anthropic_client_integration_images() -> None:
image_bytes = img_file.read()
messages = [
ChatMessage(
Message(
role="user",
contents=[
Content.from_text(text="Describe this image"),
python/packages/azure-ai-search/AGENTS.md
+1 -1
View File
@@ -16,7 +16,7 @@ provider = AzureAISearchContextProvider(
endpoint="https://your-search.search.windows.net",
index_name="your-index",
)
agent = ChatAgent(..., context_provider=provider)
agent = Agent(..., context_provider=provider)
```
## Import Path
python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py
+4 -4
View File
@@ -13,7 +13,7 @@ import sys
from collections.abc import Awaitable, Callable
from typing import TYPE_CHECKING, Any, ClassVar, Literal
from agent_framework import AGENT_FRAMEWORK_USER_AGENT, ChatMessage
from agent_framework import AGENT_FRAMEWORK_USER_AGENT, Message
from agent_framework._logging import get_logger
from agent_framework._sessions import AgentSession, BaseContextProvider, SessionContext
from agent_framework.exceptions import ServiceInitializationError
@@ -343,8 +343,8 @@ class _AzureAISearchContextProvider(BaseContextProvider):
if not search_result_parts:
return
context_messages = [ChatMessage(role="user", text=self.context_prompt)]
context_messages.extend([ChatMessage(role="user", text=part) for part in search_result_parts])
context_messages = [Message(role="user", text=self.context_prompt)]
context_messages.extend([Message(role="user", text=part) for part in search_result_parts])
context.extend_messages(self.source_id, context_messages)
# -- Internal methods (ported from AzureAISearchContextProvider) -----------
@@ -546,7 +546,7 @@ class _AzureAISearchContextProvider(BaseContextProvider):
user_agent=AGENT_FRAMEWORK_USER_AGENT,
)
async def _agentic_search(self, messages: list[ChatMessage]) -> list[str]:
async def _agentic_search(self, messages: list[Message]) -> list[str]:
"""Perform agentic retrieval with multi-hop reasoning."""
await self._ensure_knowledge_base()
python/packages/azure-ai-search/agent_framework_azure_ai_search/_search_provider.py
+6 -6
View File
@@ -7,7 +7,7 @@ import sys
from collections.abc import Awaitable, Callable, MutableSequence
from typing import TYPE_CHECKING, Any, ClassVar, Literal
from agent_framework import AGENT_FRAMEWORK_USER_AGENT, ChatMessage, Context, ContextProvider
from agent_framework import AGENT_FRAMEWORK_USER_AGENT, Context, ContextProvider, Message
from agent_framework._logging import get_logger
from agent_framework._pydantic import AFBaseSettings
from agent_framework.exceptions import ServiceInitializationError
@@ -511,7 +511,7 @@ class AzureAISearchContextProvider(ContextProvider):
@override
async def invoking(
self,
messages: ChatMessage | MutableSequence[ChatMessage],
messages: Message | MutableSequence[Message],
**kwargs: Any,
) -> Context:
"""Retrieve relevant context from Azure AI Search before model invocation.
@@ -524,7 +524,7 @@ class AzureAISearchContextProvider(ContextProvider):
Context object with retrieved documents as messages.
"""
# Convert to list and filter to USER/ASSISTANT messages with text only
messages_list = [messages] if isinstance(messages, ChatMessage) else list(messages)
messages_list = [messages] if isinstance(messages, Message) else list(messages)
def get_role_value(role: str | Any) -> str:
return role.value if hasattr(role, "value") else str(role)
@@ -553,8 +553,8 @@ class AzureAISearchContextProvider(ContextProvider):
return Context()
# Create context messages: first message with prompt, then one message per result part
context_messages = [ChatMessage(role="user", text=self.context_prompt)]
context_messages.extend([ChatMessage(role="user", text=part) for part in search_result_parts])
context_messages = [Message(role="user", text=self.context_prompt)]
context_messages.extend([Message(role="user", text=part) for part in search_result_parts])
return Context(messages=context_messages)
@@ -875,7 +875,7 @@ class AzureAISearchContextProvider(ContextProvider):
user_agent=AGENT_FRAMEWORK_USER_AGENT,
)
async def _agentic_search(self, messages: list[ChatMessage]) -> list[str]:
async def _agentic_search(self, messages: list[Message]) -> list[str]:
"""Perform agentic retrieval with multi-hop reasoning using Knowledge Bases.
This mode uses query planning and is slightly slower than semantic search,
python/packages/azure-ai-search/tests/test_aisearch_new_context_provider.py
+7 -7
View File
@@ -5,7 +5,7 @@ import os
from unittest.mock import AsyncMock, patch
import pytest
from agent_framework import ChatMessage
from agent_framework import Message
from agent_framework._sessions import AgentSession, SessionContext
from agent_framework.exceptions import ServiceInitializationError
@@ -179,7 +179,7 @@ class TestBeforeRunSemantic:
session = AgentSession(session_id="test-session")
ctx = SessionContext(
input_messages=[ChatMessage(role="user", contents=["test query"])],
input_messages=[Message(role="user", contents=["test query"])],
session_id="s1",
)
await provider.before_run(agent=None, session=session, context=ctx, state=session.state) # type: ignore[arg-type]
@@ -206,7 +206,7 @@ class TestBeforeRunSemantic:
session = AgentSession(session_id="test-session")
ctx = SessionContext(
input_messages=[ChatMessage(role="user", contents=["test query"])],
input_messages=[Message(role="user", contents=["test query"])],
session_id="s1",
)
await provider.before_run(agent=None, session=session, context=ctx, state=session.state) # type: ignore[arg-type]
@@ -221,7 +221,7 @@ class TestBeforeRunSemantic:
session = AgentSession(session_id="test-session")
ctx = SessionContext(
input_messages=[ChatMessage(role="user", contents=["test query"])],
input_messages=[Message(role="user", contents=["test query"])],
session_id="s1",
)
await provider.before_run(agent=None, session=session, context=ctx, state=session.state) # type: ignore[arg-type]
@@ -243,8 +243,8 @@ class TestBeforeRunFiltering:
session = AgentSession(session_id="test-session")
ctx = SessionContext(
input_messages=[
ChatMessage(role="system", contents=["system prompt"]),
ChatMessage(role="user", contents=["actual question"]),
Message(role="system", contents=["system prompt"]),
Message(role="user", contents=["actual question"]),
],
session_id="s1",
)
@@ -262,7 +262,7 @@ class TestBeforeRunFiltering:
session = AgentSession(session_id="test-session")
ctx = SessionContext(
input_messages=[ChatMessage(role="system", contents=["system prompt"])],
input_messages=[Message(role="system", contents=["system prompt"])],
session_id="s1",
)
await provider.before_run(agent=None, session=session, context=ctx, state=session.state) # type: ignore[arg-type]
python/packages/azure-ai-search/tests/test_search_provider.py
+17 -19
View File
@@ -5,7 +5,7 @@ import os
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from agent_framework import ChatMessage, Context
from agent_framework import Context, Message
from agent_framework.azure import AzureAISearchContextProvider, AzureAISearchSettings
from agent_framework.exceptions import ServiceInitializationError
from azure.core.credentials import AzureKeyCredential
@@ -36,10 +36,10 @@ def mock_index_client() -> AsyncMock:
@pytest.fixture
def sample_messages() -> list[ChatMessage]:
def sample_messages() -> list[Message]:
"""Create sample chat messages for testing."""
return [
ChatMessage(role="user", text="What is in the documents?"),
Message(role="user", text="What is in the documents?"),
]
@@ -276,9 +276,7 @@ class TestSemanticSearch:
@pytest.mark.asyncio
@patch("agent_framework_azure_ai_search._search_provider.SearchClient")
async def test_semantic_search_basic(
self, mock_search_class: MagicMock, sample_messages: list[ChatMessage]
) -> None:
async def test_semantic_search_basic(self, mock_search_class: MagicMock, sample_messages: list[Message]) -> None:
"""Test basic semantic search without vector search."""
# Setup mock
mock_search_client = AsyncMock()
@@ -318,7 +316,7 @@ class TestSemanticSearch:
)
# Empty message
context = await provider.invoking([ChatMessage(role="user", text="")])
context = await provider.invoking([Message(role="user", text="")])
assert isinstance(context, Context)
assert len(context.messages) == 0
@@ -326,7 +324,7 @@ class TestSemanticSearch:
@pytest.mark.asyncio
@patch("agent_framework_azure_ai_search._search_provider.SearchClient")
async def test_semantic_search_with_vector_query(
self, mock_search_class: MagicMock, sample_messages: list[ChatMessage]
self, mock_search_class: MagicMock, sample_messages: list[Message]
) -> None:
"""Test semantic search with vector query."""
# Setup mock
@@ -520,10 +518,10 @@ class TestMessageFiltering:
# Mix of message types
messages = [
ChatMessage(role="system", text="System message"),
ChatMessage(role="user", text="User message"),
ChatMessage(role="assistant", text="Assistant message"),
ChatMessage(role="tool", text="Tool message"),
Message(role="system", text="System message"),
Message(role="user", text="User message"),
Message(role="assistant", text="Assistant message"),
Message(role="tool", text="Tool message"),
]
context = await provider.invoking(messages)
@@ -548,9 +546,9 @@ class TestMessageFiltering:
# Messages with empty/whitespace text
messages = [
ChatMessage(role="user", text=""),
ChatMessage(role="user", text=" "),
ChatMessage(role="user", text=""), # ChatMessage with None text becomes empty string
Message(role="user", text=""),
Message(role="user", text=" "),
Message(role="user", text=""), # Message with None text becomes empty string
]
context = await provider.invoking(messages)
@@ -581,7 +579,7 @@ class TestCitations:
mode="semantic",
)
context = await provider.invoking([ChatMessage(role="user", text="test query")])
context = await provider.invoking([Message(role="user", text="test query")])
# Check that citation is included
assert isinstance(context, Context)
@@ -603,7 +601,7 @@ class TestAgenticSearch:
mock_search_class: MagicMock,
mock_index_class: MagicMock,
mock_retrieval_class: MagicMock,
sample_messages: list[ChatMessage],
sample_messages: list[Message],
) -> None:
"""Test basic agentic search with Knowledge Base retrieval."""
# Setup search client mock
@@ -660,7 +658,7 @@ class TestAgenticSearch:
mock_search_class: MagicMock,
mock_index_class: MagicMock,
mock_retrieval_class: MagicMock,
sample_messages: list[ChatMessage],
sample_messages: list[Message],
) -> None:
"""Test agentic search when no results are returned."""
# Setup mocks
@@ -705,7 +703,7 @@ class TestAgenticSearch:
mock_search_class: MagicMock,
mock_index_class: MagicMock,
mock_retrieval_class: MagicMock,
sample_messages: list[ChatMessage],
sample_messages: list[Message],
) -> None:
"""Test agentic search with medium reasoning effort."""
# Setup mocks
python/packages/azure-ai/agent_framework_azure_ai/_agent_provider.py
+29 -31
View File
@@ -4,11 +4,11 @@ from __future__ import annotations
import sys
from collections.abc import Callable, MutableMapping, Sequence
from typing import TYPE_CHECKING, Any, Generic, cast
from typing import Any, Generic, cast
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
ChatAgent,
Agent,
ContextProvider,
FunctionTool,
MiddlewareTypes,
@@ -18,16 +18,14 @@ from agent_framework import (
from agent_framework._mcp import MCPTool
from agent_framework.exceptions import ServiceInitializationError
from azure.ai.agents.aio import AgentsClient
from azure.ai.agents.models import Agent, ResponseFormatJsonSchema, ResponseFormatJsonSchemaType
from azure.ai.agents.models import Agent as AzureAgent
from azure.ai.agents.models import ResponseFormatJsonSchema, ResponseFormatJsonSchemaType
from azure.core.credentials_async import AsyncTokenCredential
from pydantic import BaseModel, ValidationError
from ._chat_client import AzureAIAgentClient
from ._chat_client import AzureAIAgentClient, AzureAIAgentOptions
from ._shared import AzureAISettings, from_azure_ai_agent_tools, to_azure_ai_agent_tools
if TYPE_CHECKING:
from ._chat_client import AzureAIAgentOptions
if sys.version_info >= (3, 13):
from typing import Self, TypeVar # type: ignore # pragma: no cover
else:
@@ -38,7 +36,7 @@ else:
from typing_extensions import TypedDict # type: ignore # pragma: no cover
# Type variable for options - allows typed ChatAgent[OptionsCoT] returns
# Type variable for options - allows typed Agent[TOptions] returns
# Default matches AzureAIAgentClient's default options type
OptionsCoT = TypeVar(
"OptionsCoT",
@@ -51,7 +49,7 @@ OptionsCoT = TypeVar(
class AzureAIAgentsProvider(Generic[OptionsCoT]):
"""Provider for Azure AI Agent Service V1 (Persistent Agents API).
This provider enables creating, retrieving, and wrapping Azure AI agents as ChatAgent
This provider enables creating, retrieving, and wrapping Azure AI agents as Agent
instances. It manages the underlying AgentsClient lifecycle and provides a high-level
interface for agent operations.
@@ -179,11 +177,11 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]):
default_options: OptionsCoT | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
) -> ChatAgent[OptionsCoT]:
"""Create a new agent on the Azure AI service and return a ChatAgent.
) -> Agent[OptionsCoT]:
"""Create a new agent on the Azure AI service and return a Agent.
This method creates a persistent agent on the Azure AI service with the specified
configuration and returns a local ChatAgent instance for interaction.
configuration and returns a local Agent instance for interaction.
Args:
name: The name for the agent.
@@ -200,7 +198,7 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]):
context_provider: Context provider to include during agent invocation.
Returns:
ChatAgent: A ChatAgent instance configured with the created agent.
Agent: A Agent instance configured with the created agent.
Raises:
ServiceInitializationError: If model deployment name is not available.
@@ -240,7 +238,7 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]):
args["response_format"] = self._create_response_format_config(response_format)
# Normalize and convert tools
# Local MCP tools (MCPTool) are handled by ChatAgent at runtime, not stored on the Azure agent
# Local MCP tools (MCPTool) are handled by Agent at runtime, not stored on the Azure agent
normalized_tools = normalize_tools(tools)
if normalized_tools:
# Only convert non-MCP tools to Azure AI format
@@ -255,7 +253,7 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]):
# Create the agent on the service
created_agent = await self._agents_client.create_agent(**args)
# Create ChatAgent wrapper
# Create Agent wrapper
return self._to_chat_agent_from_agent(
created_agent,
normalized_tools,
@@ -276,11 +274,11 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]):
default_options: OptionsCoT | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
) -> ChatAgent[OptionsCoT]:
"""Retrieve an existing agent from the service and return a ChatAgent.
) -> Agent[OptionsCoT]:
"""Retrieve an existing agent from the service and return a Agent.
This method fetches an agent by ID from the Azure AI service
and returns a local ChatAgent instance for interaction.
and returns a local Agent instance for interaction.
Args:
id: The ID of the agent to retrieve from the service.
@@ -294,7 +292,7 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]):
context_provider: Context provider to include during agent invocation.
Returns:
ChatAgent: A ChatAgent instance configured with the retrieved agent.
Agent: A Agent instance configured with the retrieved agent.
Raises:
ServiceInitializationError: If required function tools are not provided.
@@ -323,7 +321,7 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]):
def as_agent(
self,
agent: Agent,
agent: AzureAgent,
tools: ToolProtocol
| Callable[..., Any]
| MutableMapping[str, Any]
@@ -332,8 +330,8 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]):
default_options: OptionsCoT | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
) -> ChatAgent[OptionsCoT]:
"""Wrap an existing Agent SDK object as a ChatAgent without making HTTP calls.
) -> Agent[OptionsCoT]:
"""Wrap an existing Agent SDK object as a Agent without making HTTP calls.
Use this method when you already have an Agent object from a previous
SDK operation and want to use it with the Agent Framework.
@@ -348,7 +346,7 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]):
context_provider: Context provider to include during agent invocation.
Returns:
ChatAgent: A ChatAgent instance configured with the agent.
Agent: A Agent instance configured with the agent.
Raises:
ServiceInitializationError: If required function tools are not provided.
@@ -363,7 +361,7 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]):
instructions="...",
)
# Wrap as ChatAgent
# Wrap as Agent
chat_agent = provider.as_agent(sdk_agent)
"""
# Validate function tools
@@ -380,13 +378,13 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]):
def _to_chat_agent_from_agent(
self,
agent: Agent,
agent: AzureAgent,
provided_tools: Sequence[ToolProtocol | MutableMapping[str, Any]] | None = None,
default_options: OptionsCoT | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
) -> ChatAgent[OptionsCoT]:
"""Create a ChatAgent from an Agent SDK object.
) -> Agent[OptionsCoT]:
"""Create a Agent from an Agent SDK object.
Args:
agent: The Agent SDK object.
@@ -408,8 +406,8 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]):
# Merge tools: convert agent's hosted tools + user-provided function tools
merged_tools = self._merge_tools(agent.tools, provided_tools)
return ChatAgent( # type: ignore[return-value]
chat_client=client,
return Agent( # type: ignore[return-value]
client=client,
id=agent.id,
name=agent.name,
description=agent.description,
@@ -433,7 +431,7 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]):
provided_tools: User-provided tools (Agent Framework format).
Returns:
Combined list of tools for the ChatAgent.
Combined list of tools for the Agent.
"""
merged: list[ToolProtocol | dict[str, Any]] = []
@@ -452,7 +450,7 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]):
if provided_tools:
for provided_tool in provided_tools:
# FunctionTool - has implementation for function calling
# MCPTool - ChatAgent handles MCP connection and tool discovery at runtime
# MCPTool - Agent handles MCP connection and tool discovery at runtime
if isinstance(provided_tool, (FunctionTool, MCPTool)):
merged.append(provided_tool) # type: ignore[reportUnknownArgumentType]
python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py
+15 -13
View File
@@ -12,11 +12,10 @@ from typing import Any, ClassVar, Generic, TypedDict
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
Agent,
Annotation,
BaseChatClient,
ChatAgent,
ChatAndFunctionMiddlewareTypes,
ChatMessage,
ChatMessageStoreProtocol,
ChatMiddlewareLayer,
ChatOptions,
@@ -31,6 +30,7 @@ from agent_framework import (
HostedFileSearchTool,
HostedMCPTool,
HostedWebSearchTool,
Message,
MiddlewareTypes,
ResponseStream,
Role,
@@ -44,7 +44,9 @@ from agent_framework.exceptions import ServiceInitializationError, ServiceInvali
from agent_framework.observability import ChatTelemetryLayer
from azure.ai.agents.aio import AgentsClient
from azure.ai.agents.models import (
Agent,
Agent as AzureAgent,
)
from azure.ai.agents.models import (
AgentsNamedToolChoice,
AgentsNamedToolChoiceType,
AgentsToolChoiceOptionMode,
@@ -346,7 +348,7 @@ class AzureAIAgentClient(
self.should_cleanup_agent = should_cleanup_agent # Track whether we should delete the agent
self._agent_created = False # Track whether agent was created inside this class
self._should_close_client = should_close_client # Track whether we should close client connection
self._agent_definition: Agent | None = None # Cached definition for existing agent
self._agent_definition: AzureAgent | None = None # Cached definition for existing agent
async def __aenter__(self) -> Self:
"""Async context manager entry."""
@@ -365,7 +367,7 @@ class AzureAIAgentClient(
def _inner_get_response(
self,
*,
messages: Sequence[ChatMessage],
messages: Sequence[Message],
options: Mapping[str, Any],
stream: bool = False,
**kwargs: Any,
@@ -898,7 +900,7 @@ class AzureAIAgentClient(
self.agent_id = None
self._agent_created = False
async def _load_agent_definition_if_needed(self) -> Agent | None:
async def _load_agent_definition_if_needed(self) -> AzureAgent | None:
"""Load and cache agent details if not already loaded."""
if self._agent_definition is None and self.agent_id is not None:
self._agent_definition = await self.agents_client.get_agent(self.agent_id)
@@ -906,7 +908,7 @@ class AzureAIAgentClient(
async def _prepare_options(
self,
messages: Sequence[ChatMessage],
messages: Sequence[Message],
options: Mapping[str, Any],
**kwargs: Any,
) -> tuple[dict[str, Any], list[Content] | None]:
@@ -1020,7 +1022,7 @@ class AzureAIAgentClient(
async def _prepare_tool_definitions_and_resources(
self,
options: Mapping[str, Any],
agent_definition: Agent | None,
agent_definition: AzureAgent | None,
run_options: dict[str, Any],
) -> list[ToolDefinition | dict[str, Any]]:
"""Prepare tool definitions and resources for the run options."""
@@ -1084,7 +1086,7 @@ class AzureAIAgentClient(
return mcp_resources
def _prepare_messages(
self, messages: Sequence[ChatMessage]
self, messages: Sequence[Message]
) -> tuple[
list[ThreadMessageOptions] | None,
list[str],
@@ -1301,10 +1303,10 @@ class AzureAIAgentClient(
context_provider: ContextProvider | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
**kwargs: Any,
) -> ChatAgent[AzureAIAgentOptionsT]:
"""Convert this chat client to a ChatAgent.
) -> Agent[AzureAIAgentOptionsT]:
"""Convert this chat client to a Agent.
This method creates a ChatAgent instance with this client pre-configured.
This method creates a Agent instance with this client pre-configured.
It does NOT create an agent on the Azure AI service - the actual agent
will be created on the server during the first invocation (run).
@@ -1324,7 +1326,7 @@ class AzureAIAgentClient(
kwargs: Any additional keyword arguments.
Returns:
A ChatAgent instance configured with this chat client.
A Agent instance configured with this chat client.
"""
return super().as_agent(
id=id,
python/packages/azure-ai/agent_framework_azure_ai/_client.py
+10 -10
View File
@@ -8,15 +8,15 @@ from typing import Any, ClassVar, Generic, TypedDict, TypeVar, cast
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
ChatAgent,
Agent,
ChatAndFunctionMiddlewareTypes,
ChatMessage,
ChatMessageStoreProtocol,
ChatMiddlewareLayer,
ContextProvider,
FunctionInvocationConfiguration,
FunctionInvocationLayer,
HostedMCPTool,
Message,
MiddlewareTypes,
ToolProtocol,
get_logger,
@@ -329,7 +329,7 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[
if self.agent_name is None:
raise ServiceInitializationError(
"Agent name is required. Provide 'agent_name' when initializing AzureAIClient "
"or 'name' when initializing ChatAgent."
"or 'name' when initializing Agent."
)
# If no agent_version is provided, either use latest version or create a new agent:
@@ -396,7 +396,7 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[
@override
async def _prepare_options(
self,
messages: Sequence[ChatMessage],
messages: Sequence[Message],
options: Mapping[str, Any],
**kwargs: Any,
) -> dict[str, Any]:
@@ -489,9 +489,9 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[
"""Get the current conversation ID from chat options or kwargs."""
return options.get("conversation_id") or kwargs.get("conversation_id") or self.conversation_id
def _prepare_messages_for_azure_ai(self, messages: Sequence[ChatMessage]) -> tuple[list[ChatMessage], str | None]:
def _prepare_messages_for_azure_ai(self, messages: Sequence[Message]) -> tuple[list[Message], str | None]:
"""Prepare input from messages and convert system/developer messages to instructions."""
result: list[ChatMessage] = []
result: list[Message] = []
instructions_list: list[str] = []
instructions: str | None = None
@@ -575,10 +575,10 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[
context_provider: ContextProvider | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
**kwargs: Any,
) -> ChatAgent[AzureAIClientOptionsT]:
"""Convert this chat client to a ChatAgent.
) -> Agent[AzureAIClientOptionsT]:
"""Convert this chat client to a Agent.
This method creates a ChatAgent instance with this client pre-configured.
This method creates a Agent instance with this client pre-configured.
It does NOT create an agent on the Azure AI service - the actual agent
will be created on the server during the first invocation (run).
@@ -598,7 +598,7 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[
kwargs: Any additional keyword arguments.
Returns:
A ChatAgent instance configured with this chat client.
A Agent instance configured with this chat client.
"""
return super().as_agent(
id=id,
python/packages/azure-ai/agent_framework_azure_ai/_project_provider.py
+20 -20
View File
@@ -8,7 +8,7 @@ from typing import Any, Generic
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
ChatAgent,
Agent,
ContextProvider,
FunctionTool,
MiddlewareTypes,
@@ -47,7 +47,7 @@ else:
logger = get_logger("agent_framework.azure")
# Type variable for options - allows typed ChatAgent[OptionsT] returns
# Type variable for options - allows typed Agent[OptionsT] returns
# Default matches AzureAIClient's default options type
OptionsCoT = TypeVar(
"OptionsCoT",
@@ -170,8 +170,8 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
default_options: OptionsCoT | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
) -> ChatAgent[OptionsCoT]:
"""Create a new agent on the Azure AI service and return a local ChatAgent wrapper.
) -> Agent[OptionsCoT]:
"""Create a new agent on the Azure AI service and return a local Agent wrapper.
Args:
name: The name of the agent to create.
@@ -186,7 +186,7 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
context_provider: Context provider to include during agent invocation.
Returns:
ChatAgent: A ChatAgent instance configured with the created agent.
Agent: A Agent instance configured with the created agent.
Raises:
ServiceInitializationError: If required parameters are missing.
@@ -272,8 +272,8 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
default_options: OptionsCoT | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
) -> ChatAgent[OptionsCoT]:
"""Retrieve an existing agent from the Azure AI service and return a local ChatAgent wrapper.
) -> Agent[OptionsCoT]:
"""Retrieve an existing agent from the Azure AI service and return a local Agent wrapper.
You must provide either name or reference. Use `as_agent()` if you already have
AgentVersionDetails and want to avoid an async call.
@@ -288,7 +288,7 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
context_provider: Context provider to include during agent invocation.
Returns:
ChatAgent: A ChatAgent instance configured with the retrieved agent.
Agent: A Agent instance configured with the retrieved agent.
Raises:
ValueError: If no identifier is provided or required tools are missing.
@@ -308,7 +308,7 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
raise ValueError("Either name or reference must be provided to get an agent.")
if not isinstance(existing_agent.definition, PromptAgentDefinition):
raise ValueError("Agent definition must be PromptAgentDefinition to get a ChatAgent.")
raise ValueError("Agent definition must be PromptAgentDefinition to get a Agent.")
# Validate that required function tools are provided
self._validate_function_tools(existing_agent.definition.tools, tools)
@@ -332,8 +332,8 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
default_options: OptionsCoT | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
) -> ChatAgent[OptionsCoT]:
"""Wrap an SDK agent version object into a ChatAgent without making HTTP calls.
) -> Agent[OptionsCoT]:
"""Wrap an SDK agent version object into a Agent without making HTTP calls.
Use this when you already have an AgentVersionDetails from a previous API call.
@@ -346,13 +346,13 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
context_provider: Context provider to include during agent invocation.
Returns:
ChatAgent: A ChatAgent instance configured with the agent version.
Agent: A Agent instance configured with the agent version.
Raises:
ValueError: If the agent definition is not a PromptAgentDefinition or required tools are missing.
"""
if not isinstance(details.definition, PromptAgentDefinition):
raise ValueError("Agent definition must be PromptAgentDefinition to create a ChatAgent.")
raise ValueError("Agent definition must be PromptAgentDefinition to create a Agent.")
# Validate that required function tools are provided
self._validate_function_tools(details.definition.tools, tools)
@@ -372,8 +372,8 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
default_options: OptionsCoT | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
) -> ChatAgent[OptionsCoT]:
"""Create a ChatAgent from an AgentVersionDetails.
) -> Agent[OptionsCoT]:
"""Create a Agent from an AgentVersionDetails.
Args:
details: The AgentVersionDetails containing the agent definition.
@@ -385,7 +385,7 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
context_provider: Context provider to include during agent invocation.
"""
if not isinstance(details.definition, PromptAgentDefinition):
raise ValueError("Agent definition must be PromptAgentDefinition to get a ChatAgent.")
raise ValueError("Agent definition must be PromptAgentDefinition to get a Agent.")
client = AzureAIClient(
project_client=self._project_client,
@@ -400,8 +400,8 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
# but function tools need the actual implementations from provided_tools
merged_tools = self._merge_tools(details.definition.tools, provided_tools)
return ChatAgent( # type: ignore[return-value]
chat_client=client,
return Agent( # type: ignore[return-value]
client=client,
id=details.id,
name=details.name,
description=details.description,
@@ -425,7 +425,7 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
provided_tools: User-provided tools (Agent Framework format), including function implementations.
Returns:
Combined list of tools for the ChatAgent.
Combined list of tools for the Agent.
"""
merged: list[ToolProtocol | dict[str, Any]] = []
@@ -442,7 +442,7 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
if provided_tools:
for provided_tool in provided_tools:
# FunctionTool - has implementation for function calling
# MCPTool - ChatAgent handles MCP connection and tool discovery at runtime
# MCPTool - Agent handles MCP connection and tool discovery at runtime
if isinstance(provided_tool, (FunctionTool, MCPTool)):
merged.append(provided_tool) # type: ignore[reportUnknownArgumentType]
python/packages/azure-ai/tests/test_agent_provider.py
+24 -22
View File
@@ -6,7 +6,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from agent_framework import (
ChatAgent,
Agent,
Content,
HostedCodeInterpreterTool,
HostedFileSearchTool,
@@ -16,7 +16,9 @@ from agent_framework import (
)
from agent_framework.exceptions import ServiceInitializationError
from azure.ai.agents.models import (
Agent,
Agent as AzureAgent,
)
from azure.ai.agents.models import (
CodeInterpreterToolDefinition,
)
from azure.identity.aio import AzureCliCredential
@@ -156,7 +158,7 @@ async def test_create_agent_basic(
mock_agents_client: MagicMock,
) -> None:
"""Test creating a basic agent."""
mock_agent = MagicMock(spec=Agent)
mock_agent = MagicMock(spec=AzureAgent)
mock_agent.id = "test-agent-id"
mock_agent.name = "TestAgent"
mock_agent.description = "A test agent"
@@ -175,7 +177,7 @@ async def test_create_agent_basic(
description="A test agent",
)
assert isinstance(agent, ChatAgent)
assert isinstance(agent, Agent)
assert agent.name == "TestAgent"
assert agent.id == "test-agent-id"
mock_agents_client.create_agent.assert_called_once()
@@ -186,7 +188,7 @@ async def test_create_agent_with_model(
mock_agents_client: MagicMock,
) -> None:
"""Test creating an agent with explicit model."""
mock_agent = MagicMock(spec=Agent)
mock_agent = MagicMock(spec=AzureAgent)
mock_agent.id = "test-agent-id"
mock_agent.name = "TestAgent"
mock_agent.description = None
@@ -210,7 +212,7 @@ async def test_create_agent_with_tools(
mock_agents_client: MagicMock,
) -> None:
"""Test creating an agent with tools."""
mock_agent = MagicMock(spec=Agent)
mock_agent = MagicMock(spec=AzureAgent)
mock_agent.id = "test-agent-id"
mock_agent.name = "TestAgent"
mock_agent.description = None
@@ -245,7 +247,7 @@ async def test_create_agent_with_response_format(
temperature: float
description: str
mock_agent = MagicMock(spec=Agent)
mock_agent = MagicMock(spec=AzureAgent)
mock_agent.id = "test-agent-id"
mock_agent.name = "TestAgent"
mock_agent.description = None
@@ -297,7 +299,7 @@ async def test_get_agent_by_id(
mock_agents_client: MagicMock,
) -> None:
"""Test getting an agent by ID."""
mock_agent = MagicMock(spec=Agent)
mock_agent = MagicMock(spec=AzureAgent)
mock_agent.id = "existing-agent-id"
mock_agent.name = "ExistingAgent"
mock_agent.description = "An existing agent"
@@ -312,7 +314,7 @@ async def test_get_agent_by_id(
agent = await provider.get_agent("existing-agent-id")
assert isinstance(agent, ChatAgent)
assert isinstance(agent, Agent)
assert agent.id == "existing-agent-id"
mock_agents_client.get_agent.assert_called_once_with("existing-agent-id")
@@ -327,7 +329,7 @@ async def test_get_agent_with_function_tools(
mock_function_tool.function = MagicMock()
mock_function_tool.function.name = "get_weather"
mock_agent = MagicMock(spec=Agent)
mock_agent = MagicMock(spec=AzureAgent)
mock_agent.id = "agent-with-tools"
mock_agent.name = "AgentWithTools"
mock_agent.description = None
@@ -356,7 +358,7 @@ async def test_get_agent_with_provided_function_tools(
mock_function_tool.function = MagicMock()
mock_function_tool.function.name = "get_weather"
mock_agent = MagicMock(spec=Agent)
mock_agent = MagicMock(spec=AzureAgent)
mock_agent.id = "agent-with-tools"
mock_agent.name = "AgentWithTools"
mock_agent.description = None
@@ -376,7 +378,7 @@ async def test_get_agent_with_provided_function_tools(
agent = await provider.get_agent("agent-with-tools", tools=get_weather)
assert isinstance(agent, ChatAgent)
assert isinstance(agent, Agent)
assert agent.id == "agent-with-tools"
@@ -391,7 +393,7 @@ def test_as_agent_wraps_without_http(
mock_agents_client: MagicMock,
) -> None:
"""Test as_agent wraps Agent object without making HTTP calls."""
mock_agent = MagicMock(spec=Agent)
mock_agent = MagicMock(spec=AzureAgent)
mock_agent.id = "wrap-agent-id"
mock_agent.name = "WrapAgent"
mock_agent.description = "Wrapped agent"
@@ -405,7 +407,7 @@ def test_as_agent_wraps_without_http(
agent = provider.as_agent(mock_agent)
assert isinstance(agent, ChatAgent)
assert isinstance(agent, Agent)
assert agent.id == "wrap-agent-id"
assert agent.name == "WrapAgent"
# Ensure no HTTP calls were made
@@ -423,7 +425,7 @@ def test_as_agent_with_function_tools_validates(
mock_function_tool.function = MagicMock()
mock_function_tool.function.name = "my_function"
mock_agent = MagicMock(spec=Agent)
mock_agent = MagicMock(spec=AzureAgent)
mock_agent.id = "agent-id"
mock_agent.name = "Agent"
mock_agent.description = None
@@ -449,7 +451,7 @@ def test_as_agent_with_hosted_tools(
mock_code_interpreter = MagicMock()
mock_code_interpreter.type = "code_interpreter"
mock_agent = MagicMock(spec=Agent)
mock_agent = MagicMock(spec=AzureAgent)
mock_agent.id = "agent-id"
mock_agent.name = "Agent"
mock_agent.description = None
@@ -463,7 +465,7 @@ def test_as_agent_with_hosted_tools(
agent = provider.as_agent(mock_agent)
assert isinstance(agent, ChatAgent)
assert isinstance(agent, Agent)
# Should have HostedCodeInterpreterTool in the default_options tools
assert any(isinstance(t, HostedCodeInterpreterTool) for t in (agent.default_options.get("tools") or [])) # type: ignore
@@ -483,7 +485,7 @@ def test_as_agent_with_dict_function_tools_validates(
},
}
mock_agent = MagicMock(spec=Agent)
mock_agent = MagicMock(spec=AzureAgent)
mock_agent.id = "agent-id"
mock_agent.name = "Agent"
mock_agent.description = None
@@ -515,7 +517,7 @@ def test_as_agent_with_dict_function_tools_provided(
},
}
mock_agent = MagicMock(spec=Agent)
mock_agent = MagicMock(spec=AzureAgent)
mock_agent.id = "agent-id"
mock_agent.name = "Agent"
mock_agent.description = None
@@ -534,7 +536,7 @@ def test_as_agent_with_dict_function_tools_provided(
agent = provider.as_agent(mock_agent, tools=dict_based_function)
assert isinstance(agent, ChatAgent)
assert isinstance(agent, Agent)
assert agent.id == "agent-id"
@@ -810,7 +812,7 @@ async def test_integration_create_agent() -> None:
)
try:
assert isinstance(agent, ChatAgent)
assert isinstance(agent, Agent)
assert agent.name == "IntegrationTestAgent"
assert agent.id is not None
finally:
@@ -837,7 +839,7 @@ async def test_integration_get_agent() -> None:
# Then get it using the provider
agent = await provider.get_agent(created.id)
assert isinstance(agent, ChatAgent)
assert isinstance(agent, Agent)
assert agent.id == created.id
finally:
await provider._agents_client.delete_agent(created.id) # type: ignore
python/packages/azure-ai/tests/test_azure_ai_agent_client.py
+260 -260
View File
File diff suppressed because it is too large Load Diff
python/packages/azure-ai/tests/test_azure_ai_client.py
+25 -25
View File
@@ -11,10 +11,8 @@ from uuid import uuid4
import pytest
from agent_framework import (
Agent,
AgentResponse,
ChatAgent,
ChatClientProtocol,
ChatMessage,
ChatOptions,
ChatResponse,
Content,
@@ -22,6 +20,8 @@ from agent_framework import (
HostedFileSearchTool,
HostedMCPTool,
HostedWebSearchTool,
Message,
SupportsChatGetResponse,
tool,
)
from agent_framework.exceptions import ServiceInitializationError
@@ -88,19 +88,19 @@ async def temporary_chat_client(agent_name: str) -> AsyncIterator[AzureAIClient]
"""Async context manager that creates an Azure AI agent and yields an `AzureAIClient`.
The underlying agent version is cleaned up automatically after use.
Tests can construct their own `ChatAgent` instances from the yielded client.
Tests can construct their own `Agent` instances from the yielded client.
"""
endpoint = os.environ["AZURE_AI_PROJECT_ENDPOINT"]
async with (
AzureCliCredential() as credential,
AIProjectClient(endpoint=endpoint, credential=credential) as project_client,
):
chat_client = AzureAIClient(
client = AzureAIClient(
project_client=project_client,
agent_name=agent_name,
)
try:
yield chat_client
yield client
finally:
await project_client.agents.delete(agent_name=agent_name)
@@ -179,7 +179,7 @@ def test_init_with_project_client(mock_project_client: MagicMock) -> None:
assert client.agent_name == "test-agent"
assert client.agent_version == "1.0"
assert not client._should_close_client # type: ignore
assert isinstance(client, ChatClientProtocol)
assert isinstance(client, SupportsChatGetResponse)
def test_init_auto_create_client(
@@ -298,9 +298,9 @@ async def test_prepare_messages_for_azure_ai_with_system_messages(
client = create_test_azure_ai_client(mock_project_client)
messages = [
ChatMessage(role="system", contents=[Content.from_text(text="You are a helpful assistant.")]),
ChatMessage(role="user", contents=[Content.from_text(text="Hello")]),
ChatMessage(role="assistant", contents=[Content.from_text(text="System response")]),
Message(role="system", contents=[Content.from_text(text="You are a helpful assistant.")]),
Message(role="user", contents=[Content.from_text(text="Hello")]),
Message(role="assistant", contents=[Content.from_text(text="System response")]),
]
result_messages, instructions = client._prepare_messages_for_azure_ai(messages) # type: ignore
@@ -318,8 +318,8 @@ async def test_prepare_messages_for_azure_ai_no_system_messages(
client = create_test_azure_ai_client(mock_project_client)
messages = [
ChatMessage(role="user", contents=[Content.from_text(text="Hello")]),
ChatMessage(role="assistant", contents=[Content.from_text(text="Hi there!")]),
Message(role="user", contents=[Content.from_text(text="Hello")]),
Message(role="assistant", contents=[Content.from_text(text="Hi there!")]),
]
result_messages, instructions = client._prepare_messages_for_azure_ai(messages) # type: ignore
@@ -419,7 +419,7 @@ async def test_prepare_options_basic(mock_project_client: MagicMock) -> None:
"""Test prepare_options basic functionality."""
client = create_test_azure_ai_client(mock_project_client, agent_name="test-agent", agent_version="1.0")
messages = [ChatMessage(role="user", contents=[Content.from_text(text="Hello")])]
messages = [Message(role="user", contents=[Content.from_text(text="Hello")])]
with (
patch(
@@ -456,7 +456,7 @@ async def test_prepare_options_with_application_endpoint(
agent_version="1",
)
messages = [ChatMessage(role="user", contents=[Content.from_text(text="Hello")])]
messages = [Message(role="user", contents=[Content.from_text(text="Hello")])]
with (
patch(
@@ -498,7 +498,7 @@ async def test_prepare_options_with_application_project_client(
agent_version="1",
)
messages = [ChatMessage(role="user", contents=[Content.from_text(text="Hello")])]
messages = [Message(role="user", contents=[Content.from_text(text="Hello")])]
with (
patch(
@@ -977,7 +977,7 @@ async def test_prepare_options_excludes_response_format(
"""Test that prepare_options excludes response_format, text, and text_format from final run options."""
client = create_test_azure_ai_client(mock_project_client, agent_name="test-agent", agent_version="1.0")
messages = [ChatMessage(role="user", contents=[Content.from_text(text="Hello")])]
messages = [Message(role="user", contents=[Content.from_text(text="Hello")])]
chat_options: ChatOptions = {}
with (
@@ -1363,10 +1363,10 @@ async def test_integration_options(
# Prepare test message
if option_name.startswith("tool_choice"):
# Use weather-related prompt for tool tests
messages = [ChatMessage(role="user", text="What is the weather in Seattle?")]
messages = [Message(role="user", text="What is the weather in Seattle?")]
else:
# Generic prompt for simple options
messages = [ChatMessage(role="user", text="Say 'Hello World' briefly.")]
messages = [Message(role="user", text="Say 'Hello World' briefly.")]
# Build options dict
options: dict[str, Any] = {option_name: option_value, "tools": [get_weather]}
@@ -1480,11 +1480,11 @@ async def test_integration_agent_options(
# Prepare test message
if option_name.startswith("response_format"):
# Use prompt that works well with structured output
messages = [ChatMessage(role="user", text="The weather in Seattle is sunny")]
messages.append(ChatMessage(role="user", text="What is the weather in Seattle?"))
messages = [Message(role="user", text="The weather in Seattle is sunny")]
messages.append(Message(role="user", text="What is the weather in Seattle?"))
else:
# Generic prompt for simple options
messages = [ChatMessage(role="user", text="Say 'Hello World' briefly.")]
messages = [Message(role="user", text="Say 'Hello World' briefly.")]
# Build options dict
options = {option_name: option_value}
@@ -1621,8 +1621,8 @@ async def test_integration_agent_existing_thread():
async with (
temporary_chat_client(agent_name="af-int-test-existing-thread") as client,
ChatAgent(
chat_client=client,
Agent(
client=client,
instructions="You are a helpful assistant with good memory.",
) as first_agent,
):
@@ -1640,8 +1640,8 @@ async def test_integration_agent_existing_thread():
if preserved_thread:
async with (
temporary_chat_client(agent_name="af-int-test-existing-thread-2") as client,
ChatAgent(
chat_client=client,
Agent(
client=client,
instructions="You are a helpful assistant with good memory.",
) as second_agent,
):
python/packages/azure-ai/tests/test_provider.py
+7 -7
View File
@@ -4,7 +4,7 @@ import os
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from agent_framework import ChatAgent, FunctionTool
from agent_framework import Agent, FunctionTool
from agent_framework._mcp import MCPTool
from agent_framework.exceptions import ServiceInitializationError
from azure.ai.projects.aio import AIProjectClient
@@ -158,7 +158,7 @@ async def test_provider_create_agent(
description="Test Agent",
)
assert isinstance(agent, ChatAgent)
assert isinstance(agent, Agent)
assert agent.name == "test-agent"
mock_project_client.agents.create_version.assert_called_once()
@@ -192,7 +192,7 @@ async def test_provider_create_agent_with_env_model(
# Call without model parameter - should use env var
agent = await provider.create_agent(name="test-agent")
assert isinstance(agent, ChatAgent)
assert isinstance(agent, Agent)
# Verify the model from env var was used
call_args = mock_project_client.agents.create_version.call_args
assert call_args[1]["definition"].model == azure_ai_unit_test_env["AZURE_AI_MODEL_DEPLOYMENT_NAME"]
@@ -322,7 +322,7 @@ async def test_provider_get_agent_with_name(mock_project_client: MagicMock) -> N
agent = await provider.get_agent(name="test-agent")
assert isinstance(agent, ChatAgent)
assert isinstance(agent, Agent)
assert agent.name == "test-agent"
mock_project_client.agents.get.assert_called_with(agent_name="test-agent")
@@ -350,7 +350,7 @@ async def test_provider_get_agent_with_reference(mock_project_client: MagicMock)
agent_reference = AgentReference(name="test-agent", version="1.0")
agent = await provider.get_agent(reference=agent_reference)
assert isinstance(agent, ChatAgent)
assert isinstance(agent, Agent)
assert agent.name == "test-agent"
mock_project_client.agents.get_version.assert_called_with(agent_name="test-agent", agent_version="1.0")
@@ -410,7 +410,7 @@ def test_provider_as_agent(mock_project_client: MagicMock) -> None:
with patch("agent_framework_azure_ai._project_provider.AzureAIClient") as mock_azure_ai_client:
agent = provider.as_agent(mock_agent_version)
assert isinstance(agent, ChatAgent)
assert isinstance(agent, Agent)
assert agent.name == "test-agent"
assert agent.description == "Test Agent"
@@ -709,7 +709,7 @@ async def test_provider_create_and_get_agent_integration() -> None:
instructions="You are a helpful assistant. Always respond with 'Hello from provider!'",
)
assert isinstance(agent, ChatAgent)
assert isinstance(agent, Agent)
assert agent.name == "ProviderTestAgent"
# Run the agent
python/packages/azurefunctions/tests/test_app.py
+6 -14
View File
@@ -12,7 +12,7 @@ from unittest.mock import ANY, AsyncMock, Mock, patch
import azure.durable_functions as df
import azure.functions as func
import pytest
from agent_framework import AgentResponse, ChatMessage
from agent_framework import AgentResponse, Message
from agent_framework_durabletask import (
MIMETYPE_APPLICATION_JSON,
MIMETYPE_TEXT_PLAIN,
@@ -356,7 +356,7 @@ class TestAgentEntityOperations:
"""Test that entity can run agent operation."""
mock_agent = Mock()
mock_agent.run = AsyncMock(
return_value=AgentResponse(messages=[ChatMessage(role="assistant", text="Test response")])
return_value=AgentResponse(messages=[Message(role="assistant", text="Test response")])
)
entity = AgentEntity(mock_agent, state_provider=_InMemoryStateProvider(thread_id="test-conv-123"))
@@ -373,9 +373,7 @@ class TestAgentEntityOperations:
async def test_entity_stores_conversation_history(self) -> None:
"""Test that the entity stores conversation history."""
mock_agent = Mock()
mock_agent.run = AsyncMock(
return_value=AgentResponse(messages=[ChatMessage(role="assistant", text="Response 1")])
)
mock_agent.run = AsyncMock(return_value=AgentResponse(messages=[Message(role="assistant", text="Response 1")]))
entity = AgentEntity(mock_agent, state_provider=_InMemoryStateProvider(thread_id="conv-1"))
@@ -407,9 +405,7 @@ class TestAgentEntityOperations:
async def test_entity_increments_message_count(self) -> None:
"""Test that the entity increments the message count."""
mock_agent = Mock()
mock_agent.run = AsyncMock(
return_value=AgentResponse(messages=[ChatMessage(role="assistant", text="Response")])
)
mock_agent.run = AsyncMock(return_value=AgentResponse(messages=[Message(role="assistant", text="Response")]))
entity = AgentEntity(mock_agent, state_provider=_InMemoryStateProvider(thread_id="conv-1"))
@@ -448,9 +444,7 @@ class TestAgentEntityFactory:
def test_entity_function_handles_run_operation(self) -> None:
"""Test that the entity function handles the run operation."""
mock_agent = Mock()
mock_agent.run = AsyncMock(
return_value=AgentResponse(messages=[ChatMessage(role="assistant", text="Response")])
)
mock_agent.run = AsyncMock(return_value=AgentResponse(messages=[Message(role="assistant", text="Response")]))
entity_function = create_agent_entity(mock_agent)
@@ -475,9 +469,7 @@ class TestAgentEntityFactory:
def test_entity_function_handles_run_agent_operation(self) -> None:
"""Test that the entity function handles the deprecated run_agent operation for backward compatibility."""
mock_agent = Mock()
mock_agent.run = AsyncMock(
return_value=AgentResponse(messages=[ChatMessage(role="assistant", text="Response")])
)
mock_agent.run = AsyncMock(return_value=AgentResponse(messages=[Message(role="assistant", text="Response")]))
entity_function = create_agent_entity(mock_agent)
python/packages/azurefunctions/tests/test_entities.py
+2 -2
View File
@@ -10,7 +10,7 @@ from typing import Any, TypeVar
from unittest.mock import AsyncMock, Mock
import pytest
from agent_framework import AgentResponse, ChatMessage
from agent_framework import AgentResponse, Message
from agent_framework_azurefunctions._entities import create_agent_entity
@@ -19,7 +19,7 @@ FuncT = TypeVar("FuncT", bound=Callable[..., Any])
def _agent_response(text: str | None) -> AgentResponse:
"""Create an AgentResponse with a single assistant message."""
message = ChatMessage(role="assistant", text=text) if text is not None else ChatMessage(role="assistant", text="")
message = Message(role="assistant", text=text) if text is not None else Message(role="assistant", text="")
return AgentResponse(messages=[message])
python/packages/azurefunctions/tests/test_orchestration.py
+3 -3
View File
@@ -6,7 +6,7 @@ from typing import Any
from unittest.mock import Mock
import pytest
from agent_framework import AgentResponse, ChatMessage
from agent_framework import AgentResponse, Message
from agent_framework_durabletask import DurableAIAgent
from azure.durable_functions.models.Task import TaskBase, TaskState
@@ -136,7 +136,7 @@ class TestAgentResponseHelpers:
# Simulate successful entity task completion
entity_task.state = TaskState.SUCCEEDED
entity_task.result = AgentResponse(messages=[ChatMessage(role="assistant", text="Test response")]).to_dict()
entity_task.result = AgentResponse(messages=[Message(role="assistant", text="Test response")]).to_dict()
# Clear pending_tasks to simulate that parent has processed the child
task.pending_tasks.clear()
@@ -178,7 +178,7 @@ class TestAgentResponseHelpers:
# Simulate successful entity task with JSON response
entity_task.state = TaskState.SUCCEEDED
entity_task.result = AgentResponse(messages=[ChatMessage(role="assistant", text='{"answer": "42"}')]).to_dict()
entity_task.result = AgentResponse(messages=[Message(role="assistant", text='{"answer": "42"}')]).to_dict()
# Clear pending_tasks to simulate that parent has processed the child
task.pending_tasks.clear()
python/packages/bedrock/agent_framework_bedrock/_chat_client.py
+6 -6
View File
@@ -14,7 +14,6 @@ from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
BaseChatClient,
ChatAndFunctionMiddlewareTypes,
ChatMessage,
ChatMiddlewareLayer,
ChatOptions,
ChatResponse,
@@ -24,6 +23,7 @@ from agent_framework import (
FunctionInvocationConfiguration,
FunctionInvocationLayer,
FunctionTool,
Message,
ResponseStream,
ToolProtocol,
UsageDetails,
@@ -325,7 +325,7 @@ class BedrockChatClient(
def _inner_get_response(
self,
*,
messages: Sequence[ChatMessage],
messages: Sequence[Message],
options: Mapping[str, Any],
stream: bool = False,
**kwargs: Any,
@@ -359,7 +359,7 @@ class BedrockChatClient(
def _prepare_options(
self,
messages: Sequence[ChatMessage],
messages: Sequence[Message],
options: Mapping[str, Any],
**kwargs: Any,
) -> dict[str, Any]:
@@ -410,7 +410,7 @@ class BedrockChatClient(
return run_options
def _prepare_bedrock_messages(
self, messages: Sequence[ChatMessage]
self, messages: Sequence[Message]
) -> tuple[list[dict[str, str]], list[dict[str, Any]]]:
prompts: list[dict[str, str]] = []
conversation: list[dict[str, Any]] = []
@@ -482,7 +482,7 @@ class BedrockChatClient(
return aligned_blocks
def _convert_message_to_content_blocks(self, message: ChatMessage) -> list[dict[str, Any]]:
def _convert_message_to_content_blocks(self, message: Message) -> list[dict[str, Any]]:
blocks: list[dict[str, Any]] = []
for content in message.contents:
block = self._convert_content_to_bedrock_block(content)
@@ -593,7 +593,7 @@ class BedrockChatClient(
message = output.get("message", {})
content_blocks = message.get("content", []) or []
contents = self._parse_message_contents(content_blocks)
chat_message = ChatMessage(role="assistant", contents=contents, raw_representation=message)
chat_message = Message(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")
python/packages/bedrock/samples/bedrock_sample.py
+3 -3
View File
@@ -3,7 +3,7 @@
import asyncio
import logging
from agent_framework import ChatAgent, tool
from agent_framework import Agent, tool
from agent_framework_bedrock import BedrockChatClient
@@ -17,8 +17,8 @@ def get_weather(city: str) -> dict[str, str]:
async def main() -> None:
"""Run the Bedrock sample agent, invoke the weather tool, and log the response."""
agent = ChatAgent(
chat_client=BedrockChatClient(),
agent = Agent(
client=BedrockChatClient(),
instructions="You are a concise travel assistant.",
name="BedrockWeatherAgent",
tool_choice="auto",
python/packages/bedrock/tests/test_bedrock_client.py
+4 -4
View File
@@ -5,7 +5,7 @@ from __future__ import annotations
from typing import Any
import pytest
from agent_framework import ChatMessage, Content
from agent_framework import Content, Message
from agent_framework.exceptions import ServiceInitializationError
from agent_framework_bedrock import BedrockChatClient
@@ -41,8 +41,8 @@ async def test_get_response_invokes_bedrock_runtime() -> None:
)
messages = [
ChatMessage(role="system", contents=[Content.from_text(text="You are concise.")]),
ChatMessage(role="user", contents=[Content.from_text(text="hello")]),
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})
@@ -62,7 +62,7 @@ def test_build_request_requires_non_system_messages() -> None:
client=_StubBedrockRuntime(),
)
messages = [ChatMessage(role="system", contents=[Content.from_text(text="Only system text")])]
messages = [Message(role="system", contents=[Content.from_text(text="Only system text")])]
with pytest.raises(ServiceInitializationError):
client._prepare_options(messages, {})
python/packages/bedrock/tests/test_bedrock_settings.py
+5 -5
View File
@@ -6,10 +6,10 @@ from unittest.mock import MagicMock
import pytest
from agent_framework import (
ChatMessage,
ChatOptions,
Content,
FunctionTool,
Message,
)
from pydantic import BaseModel
@@ -46,7 +46,7 @@ def test_build_request_includes_tool_config() -> None:
"tools": [tool],
"tool_choice": {"mode": "required", "required_function_name": "get_weather"},
}
messages = [ChatMessage(role="user", contents=[Content.from_text(text="hi")])]
messages = [Message(role="user", contents=[Content.from_text(text="hi")])]
request = client._prepare_options(messages, options)
@@ -58,14 +58,14 @@ def test_build_request_serializes_tool_history() -> None:
client = _build_client()
options: ChatOptions = {}
messages = [
ChatMessage(role="user", contents=[Content.from_text(text="how's weather?")]),
ChatMessage(
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"}')
],
),
ChatMessage(
Message(
role="tool",
contents=[Content.from_function_result(call_id="call-1", result={"answer": "72F"})],
),
python/packages/chatkit/README.md
+5 -5
View File
@@ -7,9 +7,9 @@ Specifically, it mirrors the [Agent SDK integration](https://github.com/openai/c
- `stream_agent_response`: A helper to convert a streamed `AgentResponseUpdate`
from a Microsoft Agent Framework agent that implements `SupportsAgentRun` to ChatKit events.
- `ThreadItemConverter`: A extendable helper class to convert ChatKit thread items to
`ChatMessage` objects that can be consumed by an Agent Framework agent.
`Message` objects that can be consumed by an Agent Framework agent.
- `simple_to_agent_input`: A helper function that uses the default implementation
of `ThreadItemConverter` to convert a ChatKit thread to a list of `ChatMessage`,
of `ThreadItemConverter` to convert a ChatKit thread to a list of `Message`,
useful for getting started quickly.
## Installation
@@ -63,7 +63,7 @@ from azure.identity import AzureCliCredential
from fastapi import FastAPI, Request
from fastapi.responses import Response, StreamingResponse
from agent_framework import ChatAgent
from agent_framework import Agent
from agent_framework.azure import AzureOpenAIChatClient
from agent_framework.chatkit import simple_to_agent_input, stream_agent_response
@@ -74,8 +74,8 @@ from chatkit.types import ThreadMetadata, UserMessageItem, ThreadStreamEvent
from your_store import YourStore # type: ignore[import-not-found] # Replace with your Store implementation
# Define your agent with tools
agent = ChatAgent(
chat_client=AzureOpenAIChatClient(credential=AzureCliCredential()),
agent = Agent(
client=AzureOpenAIChatClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant.",
tools=[], # Add your tools here
)
python/packages/chatkit/agent_framework_chatkit/_converter.py
+45 -47
View File
@@ -9,8 +9,8 @@ import sys
from collections.abc import Awaitable, Callable, Sequence
from agent_framework import (
ChatMessage,
Content,
Message,
)
from chatkit.types import (
AssistantMessageItem,
@@ -39,7 +39,7 @@ logger = logging.getLogger(__name__)
class ThreadItemConverter:
"""Helper class to convert ChatKit thread items to Agent Framework ChatMessage objects.
"""Helper class to convert ChatKit thread items to Agent Framework Message objects.
This class provides a base implementation for converting ChatKit thread items
to Agent Framework messages. It can be extended to handle attachments,
@@ -64,8 +64,8 @@ class ThreadItemConverter:
async def user_message_to_input(
self, item: UserMessageItem, is_last_message: bool = True
) -> ChatMessage | list[ChatMessage] | None:
"""Convert a ChatKit UserMessageItem to Agent Framework ChatMessage(s).
) -> Message | list[Message] | None:
"""Convert a ChatKit UserMessageItem to Agent Framework Message(s).
This method is called internally by `to_agent_input()`. Override this method
to customize how user messages are converted.
@@ -75,7 +75,7 @@ class ThreadItemConverter:
is_last_message: Whether this is the last message in the thread (used for quoted_text handling).
Returns:
A ChatMessage, list of messages, or None to skip.
A Message, list of messages, or None to skip.
Note:
Instead of calling this method directly, use `to_agent_input()` which handles
@@ -102,19 +102,19 @@ class ThreadItemConverter:
# If only text and no attachments, use text parameter for simplicity
if text_content.strip() and not data_contents:
user_message = ChatMessage(role="user", text=text_content.strip())
user_message = Message(role="user", text=text_content.strip())
else:
# Build contents list with both text and attachments
contents: list[Content] = []
if text_content.strip():
contents.append(Content.from_text(text=text_content.strip()))
contents.extend(data_contents)
user_message = ChatMessage(role="user", contents=contents)
user_message = Message(role="user", contents=contents)
# Handle quoted text if this is the last message
messages = [user_message]
if item.quoted_text and is_last_message:
quoted_context = ChatMessage(
quoted_context = Message(
role="user",
text=f"The user is referring to this in particular:\n{item.quoted_text}",
)
@@ -179,10 +179,8 @@ class ThreadItemConverter:
# Subclasses can override this method to provide custom handling
return None
def hidden_context_to_input(
self, item: HiddenContextItem | SDKHiddenContextItem
) -> ChatMessage | list[ChatMessage] | None:
"""Convert a ChatKit HiddenContextItem or SDKHiddenContextItem to Agent Framework ChatMessage(s).
def hidden_context_to_input(self, item: HiddenContextItem | SDKHiddenContextItem) -> Message | list[Message] | None:
"""Convert a ChatKit HiddenContextItem or SDKHiddenContextItem to Agent Framework Message(s).
This method is called internally by `to_agent_input()`. Override this method
to customize how hidden context is converted.
@@ -195,7 +193,7 @@ class ThreadItemConverter:
item: The ChatKit hidden context item to convert.
Returns:
A ChatMessage with system role, a list of messages, or None to skip.
A Message with system role, a list of messages, or None to skip.
Note:
Instead of calling this method directly, use `to_agent_input()` which handles
@@ -213,9 +211,9 @@ class ThreadItemConverter:
content="User's email: user@example.com",
)
message = converter.hidden_context_to_input(hidden_item)
# Returns: ChatMessage(role=SYSTEM, text="<HIDDEN_CONTEXT>User's email: ...</HIDDEN_CONTEXT>")
# Returns: Message(role=SYSTEM, text="<HIDDEN_CONTEXT>User's email: ...</HIDDEN_CONTEXT>")
"""
return ChatMessage(role="system", text=f"<HIDDEN_CONTEXT>{item.content}</HIDDEN_CONTEXT>")
return Message(role="system", text=f"<HIDDEN_CONTEXT>{item.content}</HIDDEN_CONTEXT>")
def tag_to_message_content(self, tag: UserMessageTagContent) -> Content:
"""Convert a ChatKit tag (@-mention) to Agent Framework content.
@@ -250,8 +248,8 @@ class ThreadItemConverter:
name = getattr(tag.data, "name", tag.text if hasattr(tag, "text") else "unknown")
return Content.from_text(text=f"<TAG>Name:{name}</TAG>")
def task_to_input(self, item: TaskItem) -> ChatMessage | list[ChatMessage] | None:
"""Convert a ChatKit TaskItem to Agent Framework ChatMessage(s).
def task_to_input(self, item: TaskItem) -> Message | list[Message] | None:
"""Convert a ChatKit TaskItem to Agent Framework Message(s).
This method is called internally by `to_agent_input()`. Override this method
to customize how tasks are converted.
@@ -263,7 +261,7 @@ class ThreadItemConverter:
item: The ChatKit task item to convert.
Returns:
A ChatMessage, a list of messages, or None to skip the task.
A Message, a list of messages, or None to skip the task.
Note:
Instead of calling this method directly, use `to_agent_input()` which handles
@@ -294,10 +292,10 @@ class ThreadItemConverter:
f"A message was displayed to the user that the following task was performed:\n<Task>\n{task_text}\n</Task>"
)
return ChatMessage(role="user", text=text)
return Message(role="user", text=text)
def workflow_to_input(self, item: WorkflowItem) -> ChatMessage | list[ChatMessage] | None:
"""Convert a ChatKit WorkflowItem to Agent Framework ChatMessage(s).
def workflow_to_input(self, item: WorkflowItem) -> Message | list[Message] | None:
"""Convert a ChatKit WorkflowItem to Agent Framework Message(s).
This method is called internally by `to_agent_input()`. Override this method
to customize how workflows are converted.
@@ -336,7 +334,7 @@ class ThreadItemConverter:
messages = converter.workflow_to_input(workflow_item)
# Returns list of messages for each task
"""
messages: list[ChatMessage] = []
messages: list[Message] = []
for task in item.workflow.tasks:
if task.type != "custom" or (not task.title and not task.content):
continue
@@ -349,12 +347,12 @@ class ThreadItemConverter:
f"<Task>\n{task_text}\n</Task>"
)
messages.append(ChatMessage(role="user", text=text))
messages.append(Message(role="user", text=text))
return messages if messages else None
def widget_to_input(self, item: WidgetItem) -> ChatMessage | list[ChatMessage] | None:
"""Convert a ChatKit WidgetItem to Agent Framework ChatMessage(s).
def widget_to_input(self, item: WidgetItem) -> Message | list[Message] | None:
"""Convert a ChatKit WidgetItem to Agent Framework Message(s).
This method is called internally by `to_agent_input()`. Override this method
to customize how widgets are converted.
@@ -367,7 +365,7 @@ class ThreadItemConverter:
item: The ChatKit widget item to convert.
Returns:
A ChatMessage describing the widget, or None to skip.
A Message describing the widget, or None to skip.
Note:
Instead of calling this method directly, use `to_agent_input()` which handles
@@ -391,13 +389,13 @@ class ThreadItemConverter:
try:
widget_json = item.widget.model_dump_json(exclude_unset=True, exclude_none=True)
text = f"The following graphical UI widget (id: {item.id}) was displayed to the user:{widget_json}"
return ChatMessage(role="user", text=text)
return Message(role="user", text=text)
except Exception:
# If JSON serialization fails, skip the widget
return None
async def assistant_message_to_input(self, item: AssistantMessageItem) -> ChatMessage | list[ChatMessage] | None:
"""Convert a ChatKit AssistantMessageItem to Agent Framework ChatMessage(s).
async def assistant_message_to_input(self, item: AssistantMessageItem) -> Message | list[Message] | None:
"""Convert a ChatKit AssistantMessageItem to Agent Framework Message(s).
The default implementation extracts text from all content parts and creates
an assistant message.
@@ -406,7 +404,7 @@ class ThreadItemConverter:
item: The ChatKit assistant message item to convert.
Returns:
A ChatMessage with assistant role, or None to skip.
A Message with assistant role, or None to skip.
Note:
Instead of calling this method directly, use `to_agent_input()` which handles
@@ -417,10 +415,10 @@ class ThreadItemConverter:
if not text_parts:
return None
return ChatMessage(role="assistant", text="".join(text_parts))
return Message(role="assistant", text="".join(text_parts))
async def client_tool_call_to_input(self, item: ClientToolCallItem) -> ChatMessage | list[ChatMessage] | None:
"""Convert a ChatKit ClientToolCallItem to Agent Framework ChatMessage(s).
async def client_tool_call_to_input(self, item: ClientToolCallItem) -> Message | list[Message] | None:
"""Convert a ChatKit ClientToolCallItem to Agent Framework Message(s).
The default implementation converts completed tool calls into function call
and result content.
@@ -442,7 +440,7 @@ class ThreadItemConverter:
import json
# Create function call message
function_call_msg = ChatMessage(
function_call_msg = Message(
role="assistant",
contents=[
Content.from_function_call(
@@ -454,7 +452,7 @@ class ThreadItemConverter:
)
# Create function result message
function_result_msg = ChatMessage(
function_result_msg = Message(
role="tool",
contents=[
Content.from_function_result(
@@ -466,8 +464,8 @@ class ThreadItemConverter:
return [function_call_msg, function_result_msg]
async def end_of_turn_to_input(self, item: EndOfTurnItem) -> ChatMessage | list[ChatMessage] | None:
"""Convert a ChatKit EndOfTurnItem to Agent Framework ChatMessage(s).
async def end_of_turn_to_input(self, item: EndOfTurnItem) -> Message | list[Message] | None:
"""Convert a ChatKit EndOfTurnItem to Agent Framework Message(s).
The default implementation skips end-of-turn markers as they are only UI hints.
@@ -488,15 +486,15 @@ class ThreadItemConverter:
self,
item: ThreadItem,
is_last_message: bool = True,
) -> list[ChatMessage]:
"""Internal method to convert a single ThreadItem to ChatMessage(s).
) -> list[Message]:
"""Internal method to convert a single ThreadItem to Message(s).
Args:
item: The thread item to convert.
is_last_message: Whether this is the last item in the thread.
Returns:
A list of ChatMessage objects (may be empty).
A list of Message objects (may be empty).
"""
match item:
case UserMessageItem():
@@ -535,7 +533,7 @@ class ThreadItemConverter:
async def to_agent_input(
self,
thread_items: Sequence[ThreadItem] | ThreadItem,
) -> list[ChatMessage]:
) -> list[Message]:
"""Convert ChatKit thread items to Agent Framework ChatMessages.
This is the main entry point for converting ChatKit thread items. It handles
@@ -546,7 +544,7 @@ class ThreadItemConverter:
thread_items: A single ThreadItem or a sequence of ThreadItems to convert.
Returns:
A list of ChatMessage objects that can be sent to an Agent Framework agent.
A list of Message objects that can be sent to an Agent Framework agent.
Examples:
.. code-block:: python
@@ -562,14 +560,14 @@ class ThreadItemConverter:
messages = await converter.to_agent_input([user_message_item, assistant_message_item, task_item])
# Use with agent
from agent_framework import ChatAgent
from agent_framework import Agent
agent = ChatAgent(...)
agent = Agent(...)
response = await agent.run(messages)
"""
thread_items = list(thread_items) if isinstance(thread_items, Sequence) else [thread_items]
output: list[ChatMessage] = []
output: list[Message] = []
for item in thread_items:
output.extend(
await self._thread_item_to_input_item(
@@ -584,7 +582,7 @@ class ThreadItemConverter:
_DEFAULT_CONVERTER = ThreadItemConverter()
async def simple_to_agent_input(thread_items: Sequence[ThreadItem] | ThreadItem) -> list[ChatMessage]:
async def simple_to_agent_input(thread_items: Sequence[ThreadItem] | ThreadItem) -> list[Message]:
"""Helper function that uses the default ThreadItemConverter.
This function provides a quick way to get started with ChatKit integration
@@ -594,7 +592,7 @@ async def simple_to_agent_input(thread_items: Sequence[ThreadItem] | ThreadItem)
thread_items: A single ThreadItem or a sequence of ThreadItems to convert.
Returns:
A list of ChatMessage objects that can be sent to an Agent Framework agent.
A list of Message objects that can be sent to an Agent Framework agent.
Examples:
.. code-block:: python
python/packages/chatkit/tests/test_converter.py
+10 -10
View File
@@ -5,7 +5,7 @@
from unittest.mock import Mock
import pytest
from agent_framework import ChatMessage
from agent_framework import Message
from chatkit.types import UserMessageTextContent
from agent_framework_chatkit import ThreadItemConverter, simple_to_agent_input
@@ -43,7 +43,7 @@ class TestThreadItemConverter:
result = await converter.to_agent_input(input_item)
assert len(result) == 1
assert isinstance(result[0], ChatMessage)
assert isinstance(result[0], Message)
assert result[0].role == "user"
assert result[0].text == "Hello, how can you help me?"
@@ -110,13 +110,13 @@ class TestThreadItemConverter:
assert result[0].text == "Hello world!"
def test_hidden_context_to_input(self, converter):
"""Test converting hidden context item to ChatMessage."""
"""Test converting hidden context item to Message."""
hidden_item = Mock()
hidden_item.content = "This is hidden context information"
result = converter.hidden_context_to_input(hidden_item)
assert isinstance(result, ChatMessage)
assert isinstance(result, Message)
assert result.role == "system"
assert result.text == "<HIDDEN_CONTEXT>This is hidden context information</HIDDEN_CONTEXT>"
@@ -288,7 +288,7 @@ class TestThreadItemConverter:
assert message.contents[1].media_type == "application/pdf"
def test_task_to_input(self, converter):
"""Test converting TaskItem to ChatMessage."""
"""Test converting TaskItem to Message."""
from datetime import datetime
from chatkit.types import CustomTask, TaskItem
@@ -302,7 +302,7 @@ class TestThreadItemConverter:
)
result = converter.task_to_input(task_item)
assert isinstance(result, ChatMessage)
assert isinstance(result, Message)
assert result.role == "user"
assert "Analysis: Analyzed the data" in result.text
assert "<Task>" in result.text
@@ -347,7 +347,7 @@ class TestThreadItemConverter:
result = converter.workflow_to_input(workflow_item)
assert isinstance(result, list)
assert len(result) == 2
assert all(isinstance(msg, ChatMessage) for msg in result)
assert all(isinstance(msg, Message) for msg in result)
assert "Step 1: First step" in result[0].text
assert "Step 2: Second step" in result[1].text
@@ -369,7 +369,7 @@ class TestThreadItemConverter:
assert result is None
def test_widget_to_input(self, converter):
"""Test converting WidgetItem to ChatMessage."""
"""Test converting WidgetItem to Message."""
from datetime import datetime
from chatkit.types import WidgetItem
@@ -384,7 +384,7 @@ class TestThreadItemConverter:
)
result = converter.widget_to_input(widget_item)
assert isinstance(result, ChatMessage)
assert isinstance(result, Message)
assert result.role == "user"
assert "widget_1" in result.text
assert "graphical UI widget" in result.text
@@ -417,6 +417,6 @@ class TestSimpleToAgentInput:
result = await simple_to_agent_input(input_item)
assert len(result) == 1
assert isinstance(result[0], ChatMessage)
assert isinstance(result[0], Message)
assert result[0].role == "user"
assert result[0].text == "Test message"
python/packages/claude/agent_framework_claude/_agent.py
+7 -7
View File
@@ -14,10 +14,10 @@ from agent_framework import (
AgentResponseUpdate,
AgentThread,
BaseAgent,
ChatMessage,
Content,
ContextProvider,
FunctionTool,
Message,
ToolProtocol,
get_logger,
normalize_messages,
@@ -541,7 +541,7 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]):
if "permission_mode" in options:
await self._client.set_permission_mode(options["permission_mode"])
def _format_prompt(self, messages: list[ChatMessage] | None) -> str:
def _format_prompt(self, messages: list[Message] | None) -> str:
"""Format messages into a prompt string.
Args:
@@ -557,7 +557,7 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]):
@overload
def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[True],
thread: AgentThread | None = None,
@@ -568,7 +568,7 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]):
@overload
async def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[False] = ...,
thread: AgentThread | None = None,
@@ -578,7 +578,7 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]):
def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
@@ -608,7 +608,7 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]):
async def _run_non_streaming(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
messages: str | Message | Sequence[str | Message] | None = None,
*,
thread: AgentThread | None = None,
options: OptionsT | MutableMapping[str, Any] | None = None,
@@ -622,7 +622,7 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]):
async def _run_streaming(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
messages: str | Message | Sequence[str | Message] | None = None,
*,
thread: AgentThread | None = None,
options: OptionsT | MutableMapping[str, Any] | None = None,
python/packages/claude/tests/test_claude_agent.py
+5 -5
View File
@@ -4,7 +4,7 @@ from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from agent_framework import AgentResponseUpdate, AgentThread, ChatMessage, Content, tool
from agent_framework import AgentResponseUpdate, AgentThread, Content, Message, tool
from agent_framework_claude import ClaudeAgent, ClaudeAgentOptions, ClaudeAgentSettings
from agent_framework_claude._agent import TOOLS_MCP_SERVER_NAME
@@ -686,7 +686,7 @@ class TestFormatPrompt:
def test_format_user_message(self) -> None:
"""Test formatting user message."""
agent = ClaudeAgent()
msg = ChatMessage(
msg = Message(
role="user",
contents=[Content.from_text(text="Hello")],
)
@@ -697,9 +697,9 @@ class TestFormatPrompt:
"""Test formatting multiple messages."""
agent = ClaudeAgent()
messages = [
ChatMessage(role="user", contents=[Content.from_text(text="Hi")]),
ChatMessage(role="assistant", contents=[Content.from_text(text="Hello!")]),
ChatMessage(role="user", contents=[Content.from_text(text="How are you?")]),
Message(role="user", contents=[Content.from_text(text="Hi")]),
Message(role="assistant", contents=[Content.from_text(text="Hello!")]),
Message(role="user", contents=[Content.from_text(text="How are you?")]),
]
result = agent._format_prompt(messages) # type: ignore[reportPrivateUsage]
assert "Hi" in result
python/packages/copilotstudio/agent_framework_copilotstudio/_agent.py
+10 -10
View File
@@ -11,9 +11,9 @@ from agent_framework import (
AgentResponseUpdate,
AgentThread,
BaseAgent,
ChatMessage,
Content,
ContextProvider,
Message,
ResponseStream,
normalize_messages,
)
@@ -210,7 +210,7 @@ class CopilotStudioAgent(BaseAgent):
@overload
def run(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
messages: str | Message | list[str] | list[Message] | None = None,
*,
stream: Literal[False] = False,
thread: AgentThread | None = None,
@@ -220,7 +220,7 @@ class CopilotStudioAgent(BaseAgent):
@overload
def run(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
messages: str | Message | list[str] | list[Message] | None = None,
*,
stream: Literal[True],
thread: AgentThread | None = None,
@@ -229,7 +229,7 @@ class CopilotStudioAgent(BaseAgent):
def run(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
messages: str | Message | list[str] | list[Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
@@ -259,7 +259,7 @@ class CopilotStudioAgent(BaseAgent):
async def _run_impl(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
messages: str | Message | list[str] | list[Message] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
@@ -274,7 +274,7 @@ class CopilotStudioAgent(BaseAgent):
question = "\n".join([message.text for message in input_messages])
activities = self.client.ask_question(question, thread.service_thread_id)
response_messages: list[ChatMessage] = []
response_messages: list[Message] = []
response_id: str | None = None
response_messages = [message async for message in self._process_activities(activities, streaming=False)]
@@ -284,7 +284,7 @@ class CopilotStudioAgent(BaseAgent):
def _run_stream_impl(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
messages: str | Message | list[str] | list[Message] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
@@ -338,7 +338,7 @@ class CopilotStudioAgent(BaseAgent):
return conversation_id
async def _process_activities(self, activities: AsyncIterable[Any], streaming: bool) -> AsyncIterable[ChatMessage]:
async def _process_activities(self, activities: AsyncIterable[Any], streaming: bool) -> AsyncIterable[Message]:
"""Process activities from the Copilot Studio agent.
Args:
@@ -347,13 +347,13 @@ class CopilotStudioAgent(BaseAgent):
or non-streaming (message activities) responses.
Yields:
ChatMessage objects created from the activities.
Message objects created from the activities.
"""
async for activity in activities:
if activity.text and (
(activity.type == "message" and not streaming) or (activity.type == "typing" and streaming)
):
yield ChatMessage(
yield Message(
role="assistant",
contents=[Content.from_text(activity.text)],
author_name=activity.from_property.name if activity.from_property else None,
python/packages/copilotstudio/tests/test_copilot_agent.py
+3 -3
View File
@@ -4,7 +4,7 @@ from typing import Any
from unittest.mock import MagicMock, patch
import pytest
from agent_framework import AgentResponse, AgentResponseUpdate, AgentThread, ChatMessage, Content
from agent_framework import AgentResponse, AgentResponseUpdate, AgentThread, Content, Message
from agent_framework.exceptions import ServiceException, ServiceInitializationError
from microsoft_agents.copilotstudio.client import CopilotClient
@@ -134,7 +134,7 @@ class TestCopilotStudioAgent:
assert response.messages[0].role == "assistant"
async def test_run_with_chat_message(self, mock_copilot_client: MagicMock, mock_activity: MagicMock) -> None:
"""Test run method with ChatMessage."""
"""Test run method with Message."""
agent = CopilotStudioAgent(client=mock_copilot_client)
conversation_activity = MagicMock()
@@ -143,7 +143,7 @@ class TestCopilotStudioAgent:
mock_copilot_client.start_conversation.return_value = create_async_generator([conversation_activity])
mock_copilot_client.ask_question.return_value = create_async_generator([mock_activity])
chat_message = ChatMessage(role="user", contents=[Content.from_text("test message")])
chat_message = Message(role="user", contents=[Content.from_text("test message")])
response = await agent.run(chat_message)
assert isinstance(response, AgentResponse)
python/packages/core/AGENTS.md
+11 -11
View File
@@ -9,7 +9,7 @@ agent_framework/
├── __init__.py # Public API exports
├── _agents.py # Agent implementations
├── _clients.py # Chat client base classes and protocols
├── _types.py # Core types (ChatMessage, ChatResponse, Content, etc.)
├── _types.py # Core types (Message, ChatResponse, Content, etc.)
├── _tools.py # Tool definitions and function invocation
├── _middleware.py # Middleware system for request/response interception
├── _threads.py # AgentThread and message store abstractions
@@ -27,16 +27,16 @@ agent_framework/
- **`SupportsAgentRun`** - Protocol defining the agent interface
- **`BaseAgent`** - Abstract base class for agents
- **`ChatAgent`** - Main agent class wrapping a chat client with tools, instructions, and middleware
- **`Agent`** - Main agent class wrapping a chat client with tools, instructions, and middleware
### Chat Clients (`_clients.py`)
- **`ChatClientProtocol`** - Protocol for chat client implementations
- **`SupportsChatGetResponse`** - Protocol for chat client implementations
- **`BaseChatClient`** - Abstract base class with middleware support; subclasses implement `_inner_get_response()` and `_inner_get_streaming_response()`
### Types (`_types.py`)
- **`ChatMessage`** - Represents a chat message with role, content, and metadata
- **`Message`** - Represents a chat message with role, content, and metadata
- **`ChatResponse`** - Response from a chat client containing messages and usage
- **`ChatResponseUpdate`** - Streaming response update
- **`AgentResponse`** / **`AgentResponseUpdate`** - Agent-level response wrappers
@@ -91,11 +91,11 @@ agent_framework/
### Creating an Agent
```python
from agent_framework import ChatAgent
from agent_framework import Agent
from agent_framework.openai import OpenAIChatClient
agent = ChatAgent(
chat_client=OpenAIChatClient(),
agent = Agent(
client=OpenAIChatClient(),
instructions="You are helpful.",
tools=[my_function],
)
@@ -114,7 +114,7 @@ agent = OpenAIChatClient().as_agent(
### Middleware Pipeline
```python
from agent_framework import ChatAgent, AgentMiddleware, AgentContext
from agent_framework import Agent, AgentMiddleware, AgentContext
class LoggingMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, call_next) -> None:
@@ -122,18 +122,18 @@ class LoggingMiddleware(AgentMiddleware):
await call_next(context)
print(f"Output: {context.result}")
agent = ChatAgent(..., middleware=[LoggingMiddleware()])
agent = Agent(..., middleware=[LoggingMiddleware()])
```
### Custom Chat Client
```python
from agent_framework import BaseChatClient, ChatResponse, ChatMessage
from agent_framework import BaseChatClient, ChatResponse, Message
class MyClient(BaseChatClient):
async def _inner_get_response(self, *, messages, options, **kwargs) -> ChatResponse:
# Call your LLM here
return ChatResponse(messages=[ChatMessage(role="assistant", text="Hi!")])
return ChatResponse(messages=[Message(role="assistant", text="Hi!")])
async def _inner_get_streaming_response(self, *, messages, options, **kwargs):
yield ChatResponseUpdate(...)
python/packages/core/README.md
+15 -15
View File
@@ -45,7 +45,7 @@ You can also override environment variables by explicitly passing configuration
```python
from agent_framework.azure import AzureOpenAIChatClient
chat_client = AzureOpenAIChatClient(
client = AzureOpenAIChatClient(
api_key="",
endpoint="",
deployment_name="",
@@ -61,12 +61,12 @@ Create agents and invoke them directly:
```python
import asyncio
from agent_framework import ChatAgent
from agent_framework import Agent
from agent_framework.openai import OpenAIChatClient
async def main():
agent = ChatAgent(
chat_client=OpenAIChatClient(),
agent = Agent(
client=OpenAIChatClient(),
instructions="""
1) A robot may not injure a human being...
2) A robot must obey orders given it by human beings...
@@ -90,14 +90,14 @@ You can use the chat client classes directly for advanced workflows:
```python
import asyncio
from agent_framework.openai import OpenAIChatClient
from agent_framework import ChatMessage, Role
from agent_framework import Message, Role
async def main():
client = OpenAIChatClient()
messages = [
ChatMessage("system", ["You are a helpful assistant."]),
ChatMessage("user", ["Write a haiku about Agent Framework."])
Message("system", ["You are a helpful assistant."]),
Message("user", ["Write a haiku about Agent Framework."])
]
response = await client.get_response(messages)
@@ -123,7 +123,7 @@ import asyncio
from typing import Annotated
from random import randint
from pydantic import Field
from agent_framework import ChatAgent
from agent_framework import Agent
from agent_framework.openai import OpenAIChatClient
@@ -145,8 +145,8 @@ def get_menu_specials() -> str:
async def main():
agent = ChatAgent(
chat_client=OpenAIChatClient(),
agent = Agent(
client=OpenAIChatClient(),
instructions="You are a helpful assistant that can provide weather and restaurant information.",
tools=[get_weather, get_menu_specials]
)
@@ -169,20 +169,20 @@ Coordinate multiple agents to collaborate on complex tasks using orchestration p
```python
import asyncio
from agent_framework import ChatAgent
from agent_framework import Agent
from agent_framework.openai import OpenAIChatClient
async def main():
# Create specialized agents
writer = ChatAgent(
chat_client=OpenAIChatClient(),
writer = Agent(
client=OpenAIChatClient(),
name="Writer",
instructions="You are a creative content writer. Generate and refine slogans based on feedback."
)
reviewer = ChatAgent(
chat_client=OpenAIChatClient(),
reviewer = Agent(
client=OpenAIChatClient(),
name="Reviewer",
instructions="You are a critical reviewer. Provide detailed feedback on proposed slogans."
)
python/packages/core/agent_framework/_agents.py
+58 -58
View File
@@ -28,7 +28,7 @@ from mcp.server.lowlevel import Server
from mcp.shared.exceptions import McpError
from pydantic import BaseModel, Field, create_model
from ._clients import BaseChatClient, ChatClientProtocol
from ._clients import BaseChatClient, SupportsChatGetResponse
from ._logging import get_logger
from ._mcp import LOG_LEVEL_MAPPING, MCPTool
from ._memory import Context, ContextProvider
@@ -43,9 +43,9 @@ from ._tools import (
from ._types import (
AgentResponse,
AgentResponseUpdate,
ChatMessage,
ChatResponse,
ChatResponseUpdate,
Message,
ResponseStream,
map_chat_to_agent_update,
normalize_messages,
@@ -157,15 +157,15 @@ def _sanitize_agent_name(agent_name: str | None) -> str | None:
class _RunContext(TypedDict):
thread: AgentThread
input_messages: list[ChatMessage]
thread_messages: list[ChatMessage]
input_messages: list[Message]
thread_messages: list[Message]
agent_name: str
chat_options: dict[str, Any]
filtered_kwargs: dict[str, Any]
finalize_kwargs: dict[str, Any]
__all__ = ["BaseAgent", "ChatAgent", "RawChatAgent", "SupportsAgentRun"]
__all__ = ["Agent", "BaseAgent", "RawAgent", "SupportsAgentRun"]
# region Agent Protocol
@@ -230,7 +230,7 @@ class SupportsAgentRun(Protocol):
@overload
def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[False] = ...,
thread: AgentThread | None = None,
@@ -242,7 +242,7 @@ class SupportsAgentRun(Protocol):
@overload
def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[True],
thread: AgentThread | None = None,
@@ -253,7 +253,7 @@ class SupportsAgentRun(Protocol):
def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
@@ -292,7 +292,7 @@ class BaseAgent(SerializationMixin):
"""Base class for all Agent Framework agents.
This is the minimal base class without middleware or telemetry layers.
For most use cases, prefer :class:`ChatAgent` which includes all standard layers.
For most use cases, prefer :class:`Agent` which includes all standard layers.
This class provides core functionality for agent implementations, including
context providers, middleware support, and thread management.
@@ -300,7 +300,7 @@ class BaseAgent(SerializationMixin):
Note:
BaseAgent cannot be instantiated directly as it doesn't implement the
``run()`` and other methods required by SupportsAgentRun.
Use a concrete implementation like ChatAgent or create a subclass.
Use a concrete implementation like Agent or create a subclass.
Examples:
.. code-block:: python
@@ -380,8 +380,8 @@ class BaseAgent(SerializationMixin):
async def _notify_thread_of_new_messages(
self,
thread: AgentThread,
input_messages: ChatMessage | Sequence[ChatMessage],
response_messages: ChatMessage | Sequence[ChatMessage],
input_messages: Message | Sequence[Message],
response_messages: Message | Sequence[Message],
**kwargs: Any,
) -> None:
"""Notify the thread of new messages.
@@ -394,9 +394,9 @@ class BaseAgent(SerializationMixin):
response_messages: The response messages to notify about.
**kwargs: Any extra arguments to pass from the agent run.
"""
if isinstance(input_messages, ChatMessage) or len(input_messages) > 0:
if isinstance(input_messages, Message) or len(input_messages) > 0:
await thread.on_new_messages(input_messages)
if isinstance(response_messages, ChatMessage) or len(response_messages) > 0:
if isinstance(response_messages, Message) or len(response_messages) > 0:
await thread.on_new_messages(response_messages)
if thread.context_provider:
await thread.context_provider.invoked(input_messages, response_messages, **kwargs)
@@ -459,16 +459,16 @@ class BaseAgent(SerializationMixin):
Examples:
.. code-block:: python
from agent_framework import ChatAgent
from agent_framework import Agent
# Create an agent
agent = ChatAgent(chat_client=client, name="research-agent", description="Performs research tasks")
agent = Agent(client=client, name="research-agent", description="Performs research tasks")
# Convert the agent to a tool
research_tool = agent.as_tool()
# Use the tool with another agent
coordinator = ChatAgent(chat_client=client, name="coordinator", tools=research_tool)
coordinator = Agent(client=client, name="coordinator", tools=research_tool)
"""
# Verify that self implements SupportsAgentRun
if not isinstance(self, SupportsAgentRun):
@@ -523,14 +523,14 @@ class BaseAgent(SerializationMixin):
return agent_tool
# region ChatAgent
# region Agent
class RawChatAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
"""A Chat Client Agent without middleware or telemetry layers.
This is the core chat agent implementation. For most use cases,
prefer :class:`ChatAgent` which includes all standard layers.
prefer :class:`Agent` which includes all standard layers.
This is the primary agent implementation that uses a chat client to interact
with language models. It supports tools, context providers, middleware, and
@@ -544,12 +544,12 @@ class RawChatAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
.. code-block:: python
from agent_framework import ChatAgent
from agent_framework import Agent
from agent_framework.openai import OpenAIChatClient
# Create a basic chat agent
client = OpenAIChatClient(model_id="gpt-4")
agent = ChatAgent(chat_client=client, name="assistant", description="A helpful assistant")
agent = Agent(client=client, name="assistant", description="A helpful assistant")
# Run the agent with a simple message
response = await agent.run("Hello, how are you?")
@@ -564,8 +564,8 @@ class RawChatAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
return f"The weather in {location} is sunny."
agent = ChatAgent(
chat_client=client,
agent = Agent(
client=client,
name="weather-agent",
instructions="You are a weather assistant.",
tools=get_weather,
@@ -583,12 +583,12 @@ class RawChatAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
.. code-block:: python
from agent_framework import ChatAgent
from agent_framework import Agent
from agent_framework.openai import OpenAIChatClient, OpenAIChatOptions
client = OpenAIChatClient(model_id="gpt-4o")
agent: ChatAgent[OpenAIChatOptions] = ChatAgent(
chat_client=client,
agent: Agent[OpenAIChatOptions] = Agent(
client=client,
name="reasoning-agent",
instructions="You are a reasoning assistant.",
options={
@@ -609,7 +609,7 @@ class RawChatAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
def __init__(
self,
chat_client: ChatClientProtocol[OptionsCoT],
client: SupportsChatGetResponse[OptionsCoT],
instructions: str | None = None,
*,
id: str | None = None,
@@ -625,10 +625,10 @@ class RawChatAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
context_provider: ContextProvider | None = None,
**kwargs: Any,
) -> None:
"""Initialize a ChatAgent instance.
"""Initialize a Agent instance.
Args:
chat_client: The chat client to use for the agent.
client: The chat client to use for the agent.
instructions: Optional instructions for the agent.
These will be put into the messages sent to the chat client service as a system message.
@@ -641,7 +641,7 @@ class RawChatAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
context_provider: The context providers to include during agent invocation.
middleware: List of middleware to intercept agent and function invocations.
default_options: A TypedDict containing chat options. When using a typed agent like
``ChatAgent[OpenAIChatOptions]``, this enables IDE autocomplete for
``Agent[OpenAIChatOptions]``, this enables IDE autocomplete for
provider-specific options including temperature, max_tokens, model_id,
tool_choice, and provider-specific options like reasoning_effort.
You can also create your own TypedDict for custom chat clients.
@@ -663,7 +663,7 @@ class RawChatAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
"Use conversation_id for service-managed threads or chat_message_store_factory for local storage."
)
if not isinstance(chat_client, FunctionInvocationLayer) and isinstance(chat_client, BaseChatClient):
if not isinstance(client, FunctionInvocationLayer) and isinstance(client, BaseChatClient):
logger.warning(
"The provided chat client does not support function invoking, this might limit agent capabilities."
)
@@ -675,7 +675,7 @@ class RawChatAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
context_provider=context_provider,
**kwargs,
)
self.chat_client = chat_client
self.client = client
self.chat_message_store_factory = chat_message_store_factory
# Get tools from options or named parameter (named param takes precedence)
@@ -702,7 +702,7 @@ class RawChatAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
# Build chat options dict
self.default_options: dict[str, Any] = {
"model_id": opts.pop("model_id", None) or (getattr(self.chat_client, "model_id", None)),
"model_id": opts.pop("model_id", None) or (getattr(self.client, "model_id", None)),
"allow_multiple_tool_calls": opts.pop("allow_multiple_tool_calls", None),
"conversation_id": conversation_id,
"frequency_penalty": opts.pop("frequency_penalty", None),
@@ -730,16 +730,16 @@ class RawChatAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
async def __aenter__(self) -> Self:
"""Enter the async context manager.
If any of the chat_client or local_mcp_tools are context managers,
If any of the client or local_mcp_tools are context managers,
they will be entered into the async exit stack to ensure proper cleanup.
Note:
This list might be extended in the future.
Returns:
The ChatAgent instance.
The Agent instance.
"""
for context_manager in chain([self.chat_client], self.mcp_tools):
for context_manager in chain([self.client], self.mcp_tools):
if isinstance(context_manager, AbstractAsyncContextManager):
await self._async_exit_stack.enter_async_context(context_manager)
return self
@@ -768,15 +768,15 @@ class RawChatAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
should check if there is already an agent name defined, and if not
set it to this value.
"""
if hasattr(self.chat_client, "_update_agent_name_and_description") and callable(
self.chat_client._update_agent_name_and_description
if hasattr(self.client, "_update_agent_name_and_description") and callable(
self.client._update_agent_name_and_description
): # type: ignore[reportAttributeAccessIssue, attr-defined]
self.chat_client._update_agent_name_and_description(self.name, self.description) # type: ignore[reportAttributeAccessIssue, attr-defined]
self.client._update_agent_name_and_description(self.name, self.description) # type: ignore[reportAttributeAccessIssue, attr-defined]
@overload
def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[False] = ...,
thread: AgentThread | None = None,
@@ -792,7 +792,7 @@ class RawChatAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
@overload
def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[False] = ...,
thread: AgentThread | None = None,
@@ -808,7 +808,7 @@ class RawChatAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
@overload
def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[True],
thread: AgentThread | None = None,
@@ -823,7 +823,7 @@ class RawChatAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
@@ -851,7 +851,7 @@ class RawChatAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
thread: The thread to use for the agent.
tools: The tools to use for this specific run (merged with default tools).
options: A TypedDict containing chat options. When using a typed agent like
``ChatAgent[OpenAIChatOptions]``, this enables IDE autocomplete for
``Agent[OpenAIChatOptions]``, this enables IDE autocomplete for
provider-specific options including temperature, max_tokens, model_id,
tool_choice, and provider-specific options like reasoning_effort.
kwargs: Additional keyword arguments for the agent.
@@ -872,7 +872,7 @@ class RawChatAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
options=options,
kwargs=kwargs,
)
response = await self.chat_client.get_response( # type: ignore[call-overload]
response = await self.client.get_response( # type: ignore[call-overload]
messages=ctx["thread_messages"],
stream=False,
options=ctx["chat_options"],
@@ -944,7 +944,7 @@ class RawChatAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
kwargs=kwargs,
)
ctx: _RunContext = ctx_holder["ctx"] # type: ignore[assignment] # Safe: we just assigned it
return self.chat_client.get_response( # type: ignore[call-overload, no-any-return]
return self.client.get_response( # type: ignore[call-overload, no-any-return]
messages=ctx["thread_messages"],
stream=True,
options=ctx["chat_options"],
@@ -979,7 +979,7 @@ class RawChatAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
async def _prepare_run_context(
self,
*,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None,
messages: str | Message | Sequence[str | Message] | None,
thread: AgentThread | None,
tools: ToolProtocol
| Callable[..., Any]
@@ -1067,7 +1067,7 @@ class RawChatAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
response: ChatResponse,
agent_name: str,
thread: AgentThread,
input_messages: list[ChatMessage],
input_messages: list[Message],
kwargs: dict[str, Any],
) -> None:
"""Finalize response by updating thread and setting author names.
@@ -1287,9 +1287,9 @@ class RawChatAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
self,
*,
thread: AgentThread | None,
input_messages: list[ChatMessage] | None = None,
input_messages: list[Message] | None = None,
**kwargs: Any,
) -> tuple[AgentThread, dict[str, Any], list[ChatMessage]]:
) -> tuple[AgentThread, dict[str, Any], list[Message]]:
"""Prepare the thread and messages for agent execution.
This method prepares the conversation thread, merges context provider data,
@@ -1325,7 +1325,7 @@ class RawChatAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
thread = thread or self.get_new_thread()
if thread.service_thread_id and thread.context_provider:
await thread.context_provider.thread_created(thread.service_thread_id)
thread_messages: list[ChatMessage] = []
thread_messages: list[Message] = []
if thread.message_store:
thread_messages.extend(await thread.message_store.list_messages() or [])
context: Context | None = None
@@ -1369,10 +1369,10 @@ class RawChatAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
return self.name or "UnnamedAgent"
class ChatAgent(
class Agent(
AgentTelemetryLayer,
AgentMiddlewareLayer,
RawChatAgent[OptionsCoT],
RawAgent[OptionsCoT],
Generic[OptionsCoT],
):
"""A Chat Client Agent with middleware, telemetry, and full layer support.
@@ -1381,12 +1381,12 @@ class ChatAgent(
- Agent middleware support for request/response interception
- OpenTelemetry-based telemetry for observability
For a minimal implementation without these features, use :class:`RawChatAgent`.
For a minimal implementation without these features, use :class:`RawAgent`.
"""
def __init__(
self,
chat_client: ChatClientProtocol[OptionsCoT],
client: SupportsChatGetResponse[OptionsCoT],
instructions: str | None = None,
*,
id: str | None = None,
@@ -1403,9 +1403,9 @@ class ChatAgent(
middleware: Sequence[MiddlewareTypes] | None = None,
**kwargs: Any,
) -> None:
"""Initialize a ChatAgent instance."""
"""Initialize a Agent instance."""
super().__init__(
chat_client=chat_client,
client=client,
instructions=instructions,
id=id,
name=name,
python/packages/core/agent_framework/_clients.py
+25 -25
View File
@@ -36,9 +36,9 @@ from ._tools import (
ToolProtocol,
)
from ._types import (
ChatMessage,
ChatResponse,
ChatResponseUpdate,
Message,
ResponseStream,
prepare_messages,
validate_chat_options,
@@ -51,7 +51,7 @@ else:
if TYPE_CHECKING:
from ._agents import ChatAgent
from ._agents import Agent
from ._middleware import (
MiddlewareTypes,
)
@@ -67,11 +67,11 @@ logger = get_logger()
__all__ = [
"BaseChatClient",
"ChatClientProtocol",
"SupportsChatGetResponse",
]
# region ChatClientProtocol Protocol
# region SupportsChatGetResponse Protocol
# Contravariant for the Protocol
OptionsContraT = TypeVar(
@@ -86,7 +86,7 @@ ResponseModelBoundT = TypeVar("ResponseModelBoundT", bound=BaseModel)
@runtime_checkable
class ChatClientProtocol(Protocol[OptionsContraT]):
class SupportsChatGetResponse(Protocol[OptionsContraT]):
"""A protocol for a chat client that can generate responses.
This protocol defines the interface that all chat clients must implement,
@@ -103,7 +103,7 @@ class ChatClientProtocol(Protocol[OptionsContraT]):
Examples:
.. code-block:: python
from agent_framework import ChatClientProtocol, ChatResponse, ChatMessage
from agent_framework import SupportsChatGetResponse, ChatResponse, Message
# Any class implementing the required methods is compatible
@@ -128,7 +128,7 @@ class ChatClientProtocol(Protocol[OptionsContraT]):
# Verify the instance satisfies the protocol
client = CustomChatClient()
assert isinstance(client, ChatClientProtocol)
assert isinstance(client, SupportsChatGetResponse)
"""
additional_properties: dict[str, Any]
@@ -136,7 +136,7 @@ class ChatClientProtocol(Protocol[OptionsContraT]):
@overload
def get_response(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage],
messages: str | Message | Sequence[str | Message],
*,
stream: Literal[False] = ...,
options: ChatOptions[ResponseModelBoundT],
@@ -146,7 +146,7 @@ class ChatClientProtocol(Protocol[OptionsContraT]):
@overload
def get_response(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage],
messages: str | Message | Sequence[str | Message],
*,
stream: Literal[False] = ...,
options: OptionsContraT | ChatOptions[None] | None = None,
@@ -156,7 +156,7 @@ class ChatClientProtocol(Protocol[OptionsContraT]):
@overload
def get_response(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage],
messages: str | Message | Sequence[str | Message],
*,
stream: Literal[True],
options: OptionsContraT | ChatOptions[Any] | None = None,
@@ -165,7 +165,7 @@ class ChatClientProtocol(Protocol[OptionsContraT]):
def get_response(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage],
messages: str | Message | Sequence[str | Message],
*,
stream: bool = False,
options: OptionsContraT | ChatOptions[Any] | None = None,
@@ -226,7 +226,7 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]):
Examples:
.. code-block:: python
from agent_framework import BaseChatClient, ChatResponse, ChatMessage
from agent_framework import BaseChatClient, ChatResponse, Message
from collections.abc import AsyncIterable
@@ -243,7 +243,7 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]):
else:
# Non-streaming implementation
return ChatResponse(
messages=[ChatMessage(role="assistant", text="Hello!")], response_id="custom-response"
messages=[Message(role="assistant", text="Hello!")], response_id="custom-response"
)
@@ -338,7 +338,7 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]):
def _inner_get_response(
self,
*,
messages: Sequence[ChatMessage],
messages: Sequence[Message],
stream: bool,
options: Mapping[str, Any],
**kwargs: Any,
@@ -365,7 +365,7 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]):
@overload
def get_response(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage],
messages: str | Message | Sequence[str | Message],
*,
stream: Literal[False] = ...,
options: ChatOptions[ResponseModelBoundT],
@@ -375,7 +375,7 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]):
@overload
def get_response(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage],
messages: str | Message | Sequence[str | Message],
*,
stream: Literal[False] = ...,
options: OptionsCoT | ChatOptions[None] | None = None,
@@ -385,7 +385,7 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]):
@overload
def get_response(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage],
messages: str | Message | Sequence[str | Message],
*,
stream: Literal[True],
options: OptionsCoT | ChatOptions[Any] | None = None,
@@ -394,7 +394,7 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]):
def get_response(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage],
messages: str | Message | Sequence[str | Message],
*,
stream: bool = False,
options: OptionsCoT | ChatOptions[Any] | None = None,
@@ -448,10 +448,10 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]):
middleware: Sequence[MiddlewareTypes] | None = None,
function_invocation_configuration: FunctionInvocationConfiguration | None = None,
**kwargs: Any,
) -> ChatAgent[OptionsCoT]:
"""Create a ChatAgent with this client.
) -> Agent[OptionsCoT]:
"""Create a Agent with this client.
This is a convenience method that creates a ChatAgent instance with this
This is a convenience method that creates a Agent instance with this
chat client already configured.
Keyword Args:
@@ -474,7 +474,7 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]):
kwargs: Any additional keyword arguments. Will be stored as ``additional_properties``.
Returns:
A ChatAgent instance configured with this chat client.
A Agent instance configured with this chat client.
Examples:
.. code-block:: python
@@ -494,10 +494,10 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]):
# Run the agent
response = await agent.run("Hello!")
"""
from ._agents import ChatAgent
from ._agents import Agent
return ChatAgent(
chat_client=self,
return Agent(
client=self,
id=id,
name=name,
description=description,
python/packages/core/agent_framework/_mcp.py
+28 -28
View File
@@ -32,8 +32,8 @@ from ._tools import (
_build_pydantic_model_from_json_schema,
)
from ._types import (
ChatMessage,
Content,
Message,
)
from .exceptions import ToolException, ToolExecutionException
@@ -43,7 +43,7 @@ else:
from typing_extensions import Self # pragma: no cover
if TYPE_CHECKING:
from ._clients import ChatClientProtocol
from ._clients import SupportsChatGetResponse
logger = logging.getLogger(__name__)
@@ -69,9 +69,9 @@ __all__ = [
def _parse_message_from_mcp(
mcp_type: types.PromptMessage | types.SamplingMessage,
) -> ChatMessage:
) -> Message:
"""Parse an MCP container type into an Agent Framework type."""
return ChatMessage(
return Message(
role=mcp_type.role,
contents=_parse_content_from_mcp(mcp_type.content),
raw_representation=mcp_type,
@@ -256,9 +256,9 @@ def _prepare_content_for_mcp(
def _prepare_message_for_mcp(
content: ChatMessage,
content: Message,
) -> list[types.TextContent | types.ImageContent | types.AudioContent | types.EmbeddedResource | types.ResourceLink]:
"""Prepare a ChatMessage for MCP format."""
"""Prepare a Message for MCP format."""
messages: list[
types.TextContent | types.ImageContent | types.AudioContent | types.EmbeddedResource | types.ResourceLink
] = []
@@ -335,7 +335,7 @@ class MCPTool:
parse_prompt_results: Literal[True] | Callable[[types.GetPromptResult], Any] | None = True,
session: ClientSession | None = None,
request_timeout: int | None = None,
chat_client: ChatClientProtocol | None = None,
client: SupportsChatGetResponse | None = None,
additional_properties: dict[str, Any] | None = None,
) -> None:
"""Initialize the MCP Tool base.
@@ -356,7 +356,7 @@ class MCPTool:
self._exit_stack = AsyncExitStack()
self.session = session
self.request_timeout = request_timeout
self.chat_client = chat_client
self.client = client
self._functions: list[FunctionTool[Any, Any]] = []
self.is_connected: bool = False
self._tools_loaded: bool = False
@@ -507,17 +507,17 @@ class MCPTool:
Returns:
Either a CreateMessageResult with the generated message or ErrorData if generation fails.
"""
if not self.chat_client:
if not self.client:
return types.ErrorData(
code=types.INTERNAL_ERROR,
message="No chat client available. Please set a chat client.",
)
logger.debug("Sampling callback called with params: %s", params)
messages: list[ChatMessage] = []
messages: list[Message] = []
for msg in params.messages:
messages.append(_parse_message_from_mcp(msg))
try:
response = await self.chat_client.get_response(
response = await self.client.get_response(
messages,
temperature=params.temperature,
max_tokens=params.maxTokens,
@@ -634,7 +634,7 @@ class MCPTool:
input_model = _get_input_model_from_mcp_prompt(prompt)
approval_mode = self._determine_approval_mode(local_name)
func: FunctionTool[BaseModel, list[ChatMessage] | Any | types.GetPromptResult] = FunctionTool(
func: FunctionTool[BaseModel, list[Message] | Any | types.GetPromptResult] = FunctionTool(
func=partial(self.get_prompt, prompt.name),
name=local_name,
description=prompt.description or "",
@@ -801,7 +801,7 @@ class MCPTool:
raise ToolExecutionException(f"Failed to call tool '{tool_name}'.", inner_exception=ex) from ex
raise ToolExecutionException(f"Failed to call tool '{tool_name}' after retries.")
async def get_prompt(self, prompt_name: str, **kwargs: Any) -> list[ChatMessage] | Any | types.GetPromptResult:
async def get_prompt(self, prompt_name: str, **kwargs: Any) -> list[Message] | Any | types.GetPromptResult:
"""Call a prompt with the given arguments.
Args:
@@ -909,7 +909,7 @@ class MCPStdioTool(MCPTool):
Examples:
.. code-block:: python
from agent_framework import MCPStdioTool, ChatAgent
from agent_framework import MCPStdioTool, Agent
# Create an MCP stdio tool
mcp_tool = MCPStdioTool(
@@ -921,7 +921,7 @@ class MCPStdioTool(MCPTool):
# Use with a chat agent
async with mcp_tool:
agent = ChatAgent(chat_client=client, name="assistant", tools=mcp_tool)
agent = Agent(client=client, name="assistant", tools=mcp_tool)
response = await agent.run("List files in the directory")
"""
@@ -942,7 +942,7 @@ class MCPStdioTool(MCPTool):
args: list[str] | None = None,
env: dict[str, str] | None = None,
encoding: str | None = None,
chat_client: ChatClientProtocol | None = None,
client: SupportsChatGetResponse | None = None,
additional_properties: dict[str, Any] | None = None,
**kwargs: Any,
) -> None:
@@ -982,7 +982,7 @@ class MCPStdioTool(MCPTool):
args: The arguments to pass to the command.
env: The environment variables to set for the command.
encoding: The encoding to use for the command output.
chat_client: The chat client to use for sampling.
client: The chat client to use for sampling.
kwargs: Any extra arguments to pass to the stdio client.
"""
super().__init__(
@@ -992,7 +992,7 @@ class MCPStdioTool(MCPTool):
allowed_tools=allowed_tools,
additional_properties=additional_properties,
session=session,
chat_client=chat_client,
client=client,
load_tools=load_tools,
parse_tool_results=parse_tool_results,
load_prompts=load_prompts,
@@ -1031,7 +1031,7 @@ class MCPStreamableHTTPTool(MCPTool):
Examples:
.. code-block:: python
from agent_framework import MCPStreamableHTTPTool, ChatAgent
from agent_framework import MCPStreamableHTTPTool, Agent
# Create an MCP HTTP tool
mcp_tool = MCPStreamableHTTPTool(
@@ -1042,7 +1042,7 @@ class MCPStreamableHTTPTool(MCPTool):
# Use with a chat agent
async with mcp_tool:
agent = ChatAgent(chat_client=client, name="assistant", tools=mcp_tool)
agent = Agent(client=client, name="assistant", tools=mcp_tool)
response = await agent.run("Fetch data from the API")
"""
@@ -1061,7 +1061,7 @@ class MCPStreamableHTTPTool(MCPTool):
approval_mode: (Literal["always_require", "never_require"] | HostedMCPSpecificApproval | None) = None,
allowed_tools: Collection[str] | None = None,
terminate_on_close: bool | None = None,
chat_client: ChatClientProtocol | None = None,
client: SupportsChatGetResponse | None = None,
additional_properties: dict[str, Any] | None = None,
http_client: httpx.AsyncClient | None = None,
**kwargs: Any,
@@ -1101,7 +1101,7 @@ class MCPStreamableHTTPTool(MCPTool):
allowed_tools: A list of tools that are allowed to use this tool.
additional_properties: Additional properties.
terminate_on_close: Close the transport when the MCP client is terminated.
chat_client: The chat client to use for sampling.
client: The chat client to use for sampling.
http_client: Optional httpx.AsyncClient to use. If not provided, the
``streamable_http_client`` API will create and manage a default client.
To configure headers, timeouts, or other HTTP client settings, create
@@ -1115,7 +1115,7 @@ class MCPStreamableHTTPTool(MCPTool):
allowed_tools=allowed_tools,
additional_properties=additional_properties,
session=session,
chat_client=chat_client,
client=client,
load_tools=load_tools,
parse_tool_results=parse_tool_results,
load_prompts=load_prompts,
@@ -1148,7 +1148,7 @@ class MCPWebsocketTool(MCPTool):
Examples:
.. code-block:: python
from agent_framework import MCPWebsocketTool, ChatAgent
from agent_framework import MCPWebsocketTool, Agent
# Create an MCP WebSocket tool
mcp_tool = MCPWebsocketTool(
@@ -1157,7 +1157,7 @@ class MCPWebsocketTool(MCPTool):
# Use with a chat agent
async with mcp_tool:
agent = ChatAgent(chat_client=client, name="assistant", tools=mcp_tool)
agent = Agent(client=client, name="assistant", tools=mcp_tool)
response = await agent.run("Connect to the real-time service")
"""
@@ -1175,7 +1175,7 @@ class MCPWebsocketTool(MCPTool):
description: str | None = None,
approval_mode: (Literal["always_require", "never_require"] | HostedMCPSpecificApproval | None) = None,
allowed_tools: Collection[str] | None = None,
chat_client: ChatClientProtocol | None = None,
client: SupportsChatGetResponse | None = None,
additional_properties: dict[str, Any] | None = None,
**kwargs: Any,
) -> None:
@@ -1213,7 +1213,7 @@ class MCPWebsocketTool(MCPTool):
A tool should not be listed in both, if so, it will require approval.
allowed_tools: A list of tools that are allowed to use this tool.
additional_properties: Additional properties.
chat_client: The chat client to use for sampling.
client: The chat client to use for sampling.
kwargs: Any extra arguments to pass to the WebSocket client.
"""
super().__init__(
@@ -1223,7 +1223,7 @@ class MCPWebsocketTool(MCPTool):
allowed_tools=allowed_tools,
additional_properties=additional_properties,
session=session,
chat_client=chat_client,
client=client,
load_tools=load_tools,
parse_tool_results=parse_tool_results,
load_prompts=load_prompts,
python/packages/core/agent_framework/_memory.py
+10 -10
View File
@@ -8,7 +8,7 @@ from collections.abc import MutableSequence, Sequence
from types import TracebackType
from typing import TYPE_CHECKING, Any, Final
from ._types import ChatMessage
from ._types import Message
if TYPE_CHECKING:
from ._tools import ToolProtocol
@@ -34,12 +34,12 @@ class Context:
Examples:
.. code-block:: python
from agent_framework import Context, ChatMessage
from agent_framework import Context, Message
# Create context with instructions
context = Context(
instructions="Use a professional tone when responding.",
messages=[ChatMessage(content="Previous context", role="user")],
messages=[Message(content="Previous context", role="user")],
tools=[my_tool],
)
@@ -51,7 +51,7 @@ class Context:
def __init__(
self,
instructions: str | None = None,
messages: Sequence[ChatMessage] | None = None,
messages: Sequence[Message] | None = None,
tools: Sequence[ToolProtocol] | None = None,
):
"""Create a new Context object.
@@ -62,7 +62,7 @@ class Context:
tools: The list of tools to provide to this run.
"""
self.instructions = instructions
self.messages: Sequence[ChatMessage] = messages or []
self.messages: Sequence[Message] = messages or []
self.tools: Sequence[ToolProtocol] = tools or []
@@ -85,7 +85,7 @@ class ContextProvider(ABC):
Examples:
.. code-block:: python
from agent_framework import ContextProvider, Context, ChatMessage
from agent_framework import ContextProvider, Context, Message
class CustomContextProvider(ContextProvider):
@@ -96,7 +96,7 @@ class ContextProvider(ABC):
# Use with a chat agent
async with CustomContextProvider() as provider:
agent = ChatAgent(chat_client=client, name="assistant", context_provider=provider)
agent = Agent(client=client, name="assistant", context_provider=provider)
"""
# Default prompt to be used by all context providers when assembling memories/instructions
@@ -116,8 +116,8 @@ class ContextProvider(ABC):
async def invoked(
self,
request_messages: ChatMessage | Sequence[ChatMessage],
response_messages: ChatMessage | Sequence[ChatMessage] | None = None,
request_messages: Message | Sequence[Message],
response_messages: Message | Sequence[Message] | None = None,
invoke_exception: Exception | None = None,
**kwargs: Any,
) -> None:
@@ -136,7 +136,7 @@ class ContextProvider(ABC):
pass
@abstractmethod
async def invoking(self, messages: ChatMessage | MutableSequence[ChatMessage], **kwargs: Any) -> Context:
async def invoking(self, messages: Message | MutableSequence[Message], **kwargs: Any) -> Context:
"""Called just before the model/agent is invoked.
Implementers can load any additional context required at this time,
python/packages/core/agent_framework/_middleware.py
+34 -34
View File
@@ -10,13 +10,13 @@ from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, Sequenc
from enum import Enum
from typing import TYPE_CHECKING, Any, Generic, Literal, TypeAlias, overload
from ._clients import ChatClientProtocol
from ._clients import SupportsChatGetResponse
from ._types import (
AgentResponse,
AgentResponseUpdate,
ChatMessage,
ChatResponse,
ChatResponseUpdate,
Message,
ResponseStream,
prepare_messages,
)
@@ -35,7 +35,7 @@ if TYPE_CHECKING:
from pydantic import BaseModel
from ._agents import SupportsAgentRun
from ._clients import ChatClientProtocol
from ._clients import SupportsChatGetResponse
from ._threads import AgentThread
from ._tools import FunctionTool
from ._types import ChatOptions, ChatResponse, ChatResponseUpdate
@@ -155,7 +155,7 @@ class AgentContext:
self,
*,
agent: SupportsAgentRun,
messages: list[ChatMessage],
messages: list[Message],
thread: AgentThread | None = None,
options: Mapping[str, Any] | None = None,
stream: bool = False,
@@ -263,7 +263,7 @@ class ChatContext:
about the chat request.
Attributes:
chat_client: The chat client being invoked.
client: The chat client being invoked.
messages: The messages being sent to the chat client.
options: The options for the chat request as a dict.
stream: Whether this is a streaming invocation.
@@ -302,8 +302,8 @@ class ChatContext:
def __init__(
self,
chat_client: ChatClientProtocol,
messages: Sequence[ChatMessage],
client: SupportsChatGetResponse,
messages: Sequence[Message],
options: Mapping[str, Any] | None,
stream: bool = False,
metadata: Mapping[str, Any] | None = None,
@@ -319,7 +319,7 @@ class ChatContext:
"""Initialize the ChatContext.
Args:
chat_client: The chat client being invoked.
client: The chat client being invoked.
messages: The messages being sent to the chat client.
options: The options for the chat request as a dict.
stream: Whether this is a streaming invocation.
@@ -330,7 +330,7 @@ class ChatContext:
stream_result_hooks: Result hooks to apply to the finalized streaming response.
stream_cleanup_hooks: Cleanup hooks to run after streaming completes.
"""
self.chat_client = chat_client
self.client = client
self.messages = messages
self.options = options
self.stream = stream
@@ -356,7 +356,7 @@ class AgentMiddleware(ABC):
Examples:
.. code-block:: python
from agent_framework import AgentMiddleware, AgentContext, ChatAgent
from agent_framework import AgentMiddleware, AgentContext, Agent
class RetryMiddleware(AgentMiddleware):
@@ -372,7 +372,7 @@ class AgentMiddleware(ABC):
# Use with an agent
agent = ChatAgent(chat_client=client, name="assistant", middleware=[RetryMiddleware()])
agent = Agent(client=client, name="assistant", middleware=[RetryMiddleware()])
"""
@abstractmethod
@@ -415,7 +415,7 @@ class FunctionMiddleware(ABC):
Examples:
.. code-block:: python
from agent_framework import FunctionMiddleware, FunctionInvocationContext, ChatAgent
from agent_framework import FunctionMiddleware, FunctionInvocationContext, Agent
class CachingMiddleware(FunctionMiddleware):
@@ -439,7 +439,7 @@ class FunctionMiddleware(ABC):
# Use with an agent
agent = ChatAgent(chat_client=client, name="assistant", middleware=[CachingMiddleware()])
agent = Agent(client=client, name="assistant", middleware=[CachingMiddleware()])
"""
@abstractmethod
@@ -479,7 +479,7 @@ class ChatMiddleware(ABC):
Examples:
.. code-block:: python
from agent_framework import ChatMiddleware, ChatContext, ChatAgent
from agent_framework import ChatMiddleware, ChatContext, Agent
class SystemPromptMiddleware(ChatMiddleware):
@@ -488,17 +488,17 @@ class ChatMiddleware(ABC):
async def process(self, context: ChatContext, call_next):
# Add system prompt to messages
from agent_framework import ChatMessage
from agent_framework import Message
context.messages.insert(0, ChatMessage(role="system", text=self.system_prompt))
context.messages.insert(0, Message(role="system", text=self.system_prompt))
# Continue execution
await call_next(context)
# Use with an agent
agent = ChatAgent(
chat_client=client,
agent = Agent(
client=client,
name="assistant",
middleware=[SystemPromptMiddleware("You are a helpful assistant.")],
)
@@ -572,7 +572,7 @@ def agent_middleware(func: AgentMiddlewareCallable) -> AgentMiddlewareCallable:
Examples:
.. code-block:: python
from agent_framework import agent_middleware, AgentContext, ChatAgent
from agent_framework import agent_middleware, AgentContext, Agent
@agent_middleware
@@ -583,7 +583,7 @@ def agent_middleware(func: AgentMiddlewareCallable) -> AgentMiddlewareCallable:
# Use with an agent
agent = ChatAgent(chat_client=client, name="assistant", middleware=[logging_middleware])
agent = Agent(client=client, name="assistant", middleware=[logging_middleware])
"""
# Add marker attribute to identify this as agent middleware
func._middleware_type: MiddlewareType = MiddlewareType.AGENT # type: ignore
@@ -605,7 +605,7 @@ def function_middleware(func: FunctionMiddlewareCallable) -> FunctionMiddlewareC
Examples:
.. code-block:: python
from agent_framework import function_middleware, FunctionInvocationContext, ChatAgent
from agent_framework import function_middleware, FunctionInvocationContext, Agent
@function_middleware
@@ -616,7 +616,7 @@ def function_middleware(func: FunctionMiddlewareCallable) -> FunctionMiddlewareC
# Use with an agent
agent = ChatAgent(chat_client=client, name="assistant", middleware=[logging_middleware])
agent = Agent(client=client, name="assistant", middleware=[logging_middleware])
"""
# Add marker attribute to identify this as function middleware
func._middleware_type: MiddlewareType = MiddlewareType.FUNCTION # type: ignore
@@ -638,7 +638,7 @@ def chat_middleware(func: ChatMiddlewareCallable) -> ChatMiddlewareCallable:
Examples:
.. code-block:: python
from agent_framework import chat_middleware, ChatContext, ChatAgent
from agent_framework import chat_middleware, ChatContext, Agent
@chat_middleware
@@ -649,7 +649,7 @@ def chat_middleware(func: ChatMiddlewareCallable) -> ChatMiddlewareCallable:
# Use with an agent
agent = ChatAgent(chat_client=client, name="assistant", middleware=[logging_middleware])
agent = Agent(client=client, name="assistant", middleware=[logging_middleware])
"""
# Add marker attribute to identify this as chat middleware
func._middleware_type: MiddlewareType = MiddlewareType.CHAT # type: ignore
@@ -980,7 +980,7 @@ class ChatMiddlewareLayer(Generic[OptionsCoT]):
@overload
def get_response(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage],
messages: str | Message | Sequence[str | Message],
*,
stream: Literal[False] = ...,
options: ChatOptions[ResponseModelBoundT],
@@ -990,7 +990,7 @@ class ChatMiddlewareLayer(Generic[OptionsCoT]):
@overload
def get_response(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage],
messages: str | Message | Sequence[str | Message],
*,
stream: Literal[False] = ...,
options: OptionsCoT | ChatOptions[None] | None = None,
@@ -1000,7 +1000,7 @@ class ChatMiddlewareLayer(Generic[OptionsCoT]):
@overload
def get_response(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage],
messages: str | Message | Sequence[str | Message],
*,
stream: Literal[True],
options: OptionsCoT | ChatOptions[Any] | None = None,
@@ -1009,7 +1009,7 @@ class ChatMiddlewareLayer(Generic[OptionsCoT]):
def get_response(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage],
messages: str | Message | Sequence[str | Message],
*,
stream: bool = False,
options: OptionsCoT | ChatOptions[Any] | None = None,
@@ -1035,7 +1035,7 @@ class ChatMiddlewareLayer(Generic[OptionsCoT]):
)
context = ChatContext(
chat_client=self, # type: ignore[arg-type]
client=self, # type: ignore[arg-type]
messages=prepare_messages(messages),
options=options,
stream=stream,
@@ -1090,14 +1090,14 @@ class AgentMiddlewareLayer:
self.agent_middleware = middleware_list["agent"]
# Pass middleware to super so BaseAgent can store it for dynamic rebuild
super().__init__(*args, middleware=middleware, **kwargs) # type: ignore[call-arg]
# Note: We intentionally don't extend chat_client's middleware lists here.
# Note: We intentionally don't extend client's middleware lists here.
# Chat and function middleware is passed to the chat client at runtime via kwargs
# in AgentMiddlewareLayer.run(), where it's properly combined with run-level middleware.
@overload
def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[False] = ...,
thread: AgentThread | None = None,
@@ -1109,7 +1109,7 @@ class AgentMiddlewareLayer:
@overload
def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[False] = ...,
thread: AgentThread | None = None,
@@ -1121,7 +1121,7 @@ class AgentMiddlewareLayer:
@overload
def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[True],
thread: AgentThread | None = None,
@@ -1132,7 +1132,7 @@ class AgentMiddlewareLayer:
def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
python/packages/core/agent_framework/_serialization.py
+11 -11
View File
@@ -31,16 +31,16 @@ class SerializationProtocol(Protocol):
ensuring consistent behavior across the framework.
Examples:
The framework's ``ChatMessage`` class demonstrates the protocol in action:
The framework's ``Message`` class demonstrates the protocol in action:
.. code-block:: python
from agent_framework import ChatMessage
from agent_framework import Message
from agent_framework._serialization import SerializationProtocol
# ChatMessage implements SerializationProtocol via SerializationMixin
user_msg = ChatMessage(role="user", text="What's the weather like today?")
# Message implements SerializationProtocol via SerializationMixin
user_msg = Message(role="user", text="What's the weather like today?")
# Serialize to dictionary - automatic type identification and nested serialization
msg_dict = user_msg.to_dict()
@@ -52,8 +52,8 @@ class SerializationProtocol(Protocol):
# "additional_properties": {}
# }
# Deserialize back to ChatMessage instance - automatic type reconstruction
restored_msg = ChatMessage.from_dict(msg_dict)
# Deserialize back to Message instance - automatic type reconstruction
restored_msg = Message.from_dict(msg_dict)
print(restored_msg.text) # "What's the weather like today?"
print(restored_msg.role) # "user"
@@ -170,15 +170,15 @@ class SerializationMixin:
.. code-block:: python
from agent_framework import ChatMessage
from agent_framework import Message
from agent_framework._threads import AgentThreadState, ChatMessageStoreState
# ChatMessageStoreState handles nested ChatMessage serialization
# ChatMessageStoreState handles nested Message serialization
store_state = ChatMessageStoreState(
messages=[
ChatMessage(role="user", text="Hello agent"),
ChatMessage(role="assistant", text="Hi! How can I help?"),
Message(role="user", text="Hello agent"),
Message(role="assistant", text="Hi! How can I help?"),
]
)
@@ -443,7 +443,7 @@ class SerializationMixin:
dependencies = {"open_ai_chat_client": {"client": openai_client}}
# The chat client is reconstructed with the OpenAI client injected
chat_client = OpenAIChatClient.from_dict(client_data, dependencies=dependencies)
client = OpenAIChatClient.from_dict(client_data, dependencies=dependencies)
# Now ready to make API calls with the injected client
**Function Injection for Tools** - FunctionTool runtime dependency:
python/packages/core/agent_framework/_sessions.py
+16 -16
View File
@@ -19,7 +19,7 @@ from collections.abc import Sequence
from typing import TYPE_CHECKING, Any
from ._tools import ToolProtocol
from ._types import AgentResponse, ChatMessage
from ._types import AgentResponse, Message
if TYPE_CHECKING:
from ._agents import SupportsAgentRun
@@ -80,7 +80,7 @@ def _deserialize_state(state: dict[str, Any]) -> dict[str, Any]:
# Register known types
_register_state_type(ChatMessage)
_register_state_type(Message)
class SessionContext:
@@ -107,8 +107,8 @@ class SessionContext:
*,
session_id: str | None = None,
service_session_id: str | None = None,
input_messages: list[ChatMessage],
context_messages: dict[str, list[ChatMessage]] | None = None,
input_messages: list[Message],
context_messages: dict[str, list[Message]] | None = None,
instructions: list[str] | None = None,
tools: list[ToolProtocol] | None = None,
options: dict[str, Any] | None = None,
@@ -129,7 +129,7 @@ class SessionContext:
self.session_id = session_id
self.service_session_id = service_session_id
self.input_messages = input_messages
self.context_messages: dict[str, list[ChatMessage]] = context_messages or {}
self.context_messages: dict[str, list[Message]] = context_messages or {}
self.instructions: list[str] = instructions or []
self.tools: list[ToolProtocol] = tools or []
self._response: AgentResponse | None = None
@@ -141,7 +141,7 @@ class SessionContext:
"""The agent's response. Set by the framework after invocation, read-only for providers."""
return self._response
def extend_messages(self, source: str | object, messages: Sequence[ChatMessage]) -> None:
def extend_messages(self, source: str | object, messages: Sequence[Message]) -> None:
"""Add context messages from a specific source.
Messages are copied before attribution is added, so the caller's
@@ -164,7 +164,7 @@ class SessionContext:
source_id = source.source_id # type: ignore[attr-defined]
attribution = {"source_id": source_id, "source_type": type(source).__name__}
copied: list[ChatMessage] = []
copied: list[Message] = []
for message in messages:
msg_copy = copy.copy(message)
msg_copy.additional_properties = dict(message.additional_properties)
@@ -206,7 +206,7 @@ class SessionContext:
exclude_sources: set[str] | None = None,
include_input: bool = False,
include_response: bool = False,
) -> list[ChatMessage]:
) -> list[Message]:
"""Get context messages, optionally filtered and including input/response.
Returns messages in provider execution order (dict insertion order),
@@ -221,7 +221,7 @@ class SessionContext:
Returns:
Flattened list of messages in conversation order.
"""
result: list[ChatMessage] = []
result: list[Message] = []
for source_id, messages in self.context_messages.items():
if sources is not None and source_id not in sources:
continue
@@ -353,7 +353,7 @@ class BaseHistoryProvider(BaseContextProvider):
self.store_outputs = store_outputs
@abstractmethod
async def get_messages(self, session_id: str | None, **kwargs: Any) -> list[ChatMessage]:
async def get_messages(self, session_id: str | None, **kwargs: Any) -> list[Message]:
"""Retrieve stored messages for this session.
Args:
@@ -366,7 +366,7 @@ class BaseHistoryProvider(BaseContextProvider):
...
@abstractmethod
async def save_messages(self, session_id: str | None, messages: Sequence[ChatMessage], **kwargs: Any) -> None:
async def save_messages(self, session_id: str | None, messages: Sequence[Message], **kwargs: Any) -> None:
"""Persist messages for this session.
Args:
@@ -376,7 +376,7 @@ class BaseHistoryProvider(BaseContextProvider):
"""
...
def _get_context_messages_to_store(self, context: SessionContext) -> list[ChatMessage]:
def _get_context_messages_to_store(self, context: SessionContext) -> list[Message]:
"""Get context messages that should be stored based on configuration."""
if not self.store_context_messages:
return []
@@ -405,7 +405,7 @@ class BaseHistoryProvider(BaseContextProvider):
state: dict[str, Any],
) -> None:
"""Store messages based on configuration."""
messages_to_store: list[ChatMessage] = []
messages_to_store: list[Message] = []
messages_to_store.extend(self._get_context_messages_to_store(context))
if self.store_inputs:
messages_to_store.extend(context.input_messages)
@@ -487,7 +487,7 @@ class InMemoryHistoryProvider(BaseHistoryProvider):
"""Built-in history provider that stores messages in session.state.
Messages are stored in ``state[source_id]["messages"]`` as a list of
``ChatMessage`` objects. Serialization to/from dicts is handled by
``Message`` objects. Serialization to/from dicts is handled by
``AgentSession.to_dict()``/``from_dict()`` using ``SerializationProtocol``.
This provider holds no instance state — all data lives in the session's
@@ -499,7 +499,7 @@ class InMemoryHistoryProvider(BaseHistoryProvider):
async def get_messages(
self, session_id: str | None, *, state: dict[str, Any] | None = None, **kwargs: Any
) -> list[ChatMessage]:
) -> list[Message]:
"""Retrieve messages from session state."""
if state is None:
return []
@@ -509,7 +509,7 @@ class InMemoryHistoryProvider(BaseHistoryProvider):
async def save_messages(
self,
session_id: str | None,
messages: Sequence[ChatMessage],
messages: Sequence[Message],
*,
state: dict[str, Any] | None = None,
**kwargs: Any,
python/packages/core/agent_framework/_threads.py
+26 -26
View File
@@ -7,7 +7,7 @@ from typing import Any, Protocol, TypeVar
from ._memory import ContextProvider
from ._serialization import SerializationMixin
from ._types import ChatMessage
from ._types import Message
from .exceptions import AgentThreadException
__all__ = ["AgentThread", "ChatMessageStore", "ChatMessageStoreProtocol"]
@@ -22,17 +22,17 @@ class ChatMessageStoreProtocol(Protocol):
Examples:
.. code-block:: python
from agent_framework import ChatMessage
from agent_framework import Message
class MyMessageStore:
def __init__(self):
self._messages = []
async def list_messages(self) -> list[ChatMessage]:
async def list_messages(self) -> list[Message]:
return self._messages
async def add_messages(self, messages: Sequence[ChatMessage]) -> None:
async def add_messages(self, messages: Sequence[Message]) -> None:
self._messages.extend(messages)
@classmethod
@@ -52,7 +52,7 @@ class ChatMessageStoreProtocol(Protocol):
store = MyMessageStore()
"""
async def list_messages(self) -> list[ChatMessage]:
async def list_messages(self) -> list[Message]:
"""Gets all the messages from the store that should be used for the next agent invocation.
Messages are returned in ascending chronological order, with the oldest message first.
@@ -65,11 +65,11 @@ class ChatMessageStoreProtocol(Protocol):
"""
...
async def add_messages(self, messages: Sequence[ChatMessage]) -> None:
async def add_messages(self, messages: Sequence[Message]) -> None:
"""Adds messages to the store.
Args:
messages: The sequence of ChatMessage objects to add to the store.
messages: The sequence of Message objects to add to the store.
"""
...
@@ -128,7 +128,7 @@ class ChatMessageStoreState(SerializationMixin):
def __init__(
self,
messages: Sequence[ChatMessage] | Sequence[MutableMapping[str, Any]] | None = None,
messages: Sequence[Message] | Sequence[MutableMapping[str, Any]] | None = None,
**kwargs: Any,
) -> None:
"""Create the store state.
@@ -141,16 +141,16 @@ class ChatMessageStoreState(SerializationMixin):
"""
if not messages:
self.messages: list[ChatMessage] = []
self.messages: list[Message] = []
return
if not isinstance(messages, list):
raise TypeError("Messages should be a list")
new_messages: list[ChatMessage] = []
new_messages: list[Message] = []
for msg in messages:
if isinstance(msg, ChatMessage):
if isinstance(msg, Message):
new_messages.append(msg)
else:
new_messages.append(ChatMessage.from_dict(msg))
new_messages.append(Message.from_dict(msg))
self.messages = new_messages
@@ -198,13 +198,13 @@ class ChatMessageStore:
Examples:
.. code-block:: python
from agent_framework import ChatMessageStore, ChatMessage
from agent_framework import ChatMessageStore, Message
# Create an empty store
store = ChatMessageStore()
# Add messages
message = ChatMessage(role="user", text="Hello")
message = Message(role="user", text="Hello")
await store.add_messages([message])
# Retrieve messages
@@ -217,7 +217,7 @@ class ChatMessageStore:
restored_store = await ChatMessageStore.deserialize(state)
"""
def __init__(self, messages: Sequence[ChatMessage] | None = None):
def __init__(self, messages: Sequence[Message] | None = None):
"""Create a ChatMessageStore for use in a thread.
Args:
@@ -225,19 +225,19 @@ class ChatMessageStore:
"""
self.messages = list(messages) if messages else []
async def add_messages(self, messages: Sequence[ChatMessage]) -> None:
async def add_messages(self, messages: Sequence[Message]) -> None:
"""Add messages to the store.
Args:
messages: Sequence of ChatMessage objects to add to the store.
messages: Sequence of Message objects to add to the store.
"""
self.messages.extend(messages)
async def list_messages(self) -> list[ChatMessage]:
async def list_messages(self) -> list[Message]:
"""Get all messages from the store in chronological order.
Returns:
List of ChatMessage objects, ordered from oldest to newest.
List of Message objects, ordered from oldest to newest.
"""
return self.messages
@@ -302,21 +302,21 @@ class AgentThread:
Examples:
.. code-block:: python
from agent_framework import ChatAgent, ChatMessageStore
from agent_framework import Agent, ChatMessageStore
from agent_framework.openai import OpenAIChatClient
client = OpenAIChatClient(model="gpt-4o")
# Create agent with service-managed threads using a service_thread_id
service_agent = ChatAgent(name="assistant", client=client)
service_agent = Agent(name="assistant", client=client)
service_thread = await service_agent.get_new_thread(service_thread_id="thread_abc123")
# Create agent with service-managed threads using conversation_id
conversation_agent = ChatAgent(name="assistant", client=client, conversation_id="thread_abc123")
conversation_agent = Agent(name="assistant", client=client, conversation_id="thread_abc123")
conversation_thread = await conversation_agent.get_new_thread()
# Create agent with custom message store factory
local_agent = ChatAgent(name="assistant", client=client, chat_message_store_factory=ChatMessageStore)
local_agent = Agent(name="assistant", client=client, chat_message_store_factory=ChatMessageStore)
local_thread = await local_agent.get_new_thread()
# Serialize and restore thread state
@@ -401,11 +401,11 @@ class AgentThread:
self._message_store = message_store
async def on_new_messages(self, new_messages: ChatMessage | Sequence[ChatMessage]) -> None:
async def on_new_messages(self, new_messages: Message | Sequence[Message]) -> None:
"""Invoked when a new message has been contributed to the chat by any participant.
Args:
new_messages: The new ChatMessage or sequence of ChatMessage objects to add to the thread.
new_messages: The new Message or sequence of Message objects to add to the thread.
"""
if self._service_thread_id is not None:
# If the thread messages are stored in the service there is nothing to do here,
@@ -416,7 +416,7 @@ class AgentThread:
# create a default in memory store.
self._message_store = ChatMessageStore()
# If a store has been provided, we need to add the messages to the store.
if isinstance(new_messages, ChatMessage):
if isinstance(new_messages, Message):
new_messages = [new_messages]
await self._message_store.add_messages(new_messages)
python/packages/core/agent_framework/_tools.py
+21 -21
View File
@@ -65,14 +65,14 @@ else:
if TYPE_CHECKING:
from ._clients import ChatClientProtocol
from ._clients import SupportsChatGetResponse
from ._middleware import FunctionMiddlewarePipeline, FunctionMiddlewareTypes
from ._types import (
ChatMessage,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
Content,
Message,
ResponseStream,
)
@@ -100,7 +100,7 @@ __all__ = [
logger = get_logger()
DEFAULT_MAX_ITERATIONS: Final[int] = 40
DEFAULT_MAX_CONSECUTIVE_ERRORS_PER_REQUEST: Final[int] = 3
ChatClientT = TypeVar("ChatClientT", bound="ChatClientProtocol[Any]")
ChatClientT = TypeVar("ChatClientT", bound="SupportsChatGetResponse[Any]")
# region Helpers
ArgsT = TypeVar("ArgsT", bound=BaseModel, default=BaseModel)
@@ -1857,14 +1857,14 @@ def _extract_tools(options: dict[str, Any] | None) -> Any:
def _collect_approval_responses(
messages: list[ChatMessage],
messages: list[Message],
) -> dict[str, Content]:
"""Collect approval responses (both approved and rejected) from messages."""
from ._types import ChatMessage
from ._types import Message
fcc_todo: dict[str, Content] = {}
for msg in messages:
for content in msg.contents if isinstance(msg, ChatMessage) else []:
for content in msg.contents if isinstance(msg, Message) else []:
# Collect BOTH approved and rejected responses
if content.type == "function_approval_response":
fcc_todo[content.id] = content # type: ignore[attr-defined, index]
@@ -1872,7 +1872,7 @@ def _collect_approval_responses(
def _replace_approval_contents_with_results(
messages: list[ChatMessage],
messages: list[Message],
fcc_todo: dict[str, Content],
approved_function_results: list[Content],
) -> None:
@@ -1941,7 +1941,7 @@ def _extract_function_calls(response: ChatResponse) -> list[Content]:
]
def _prepend_fcc_messages(response: ChatResponse, fcc_messages: list[ChatMessage]) -> None:
def _prepend_fcc_messages(response: ChatResponse, fcc_messages: list[Message]) -> None:
if not fcc_messages:
return
for msg in reversed(fcc_messages):
@@ -1961,7 +1961,7 @@ class FunctionRequestResult(TypedDict, total=False):
action: Literal["return", "continue", "stop"]
errors_in_a_row: int
result_message: ChatMessage | None
result_message: Message | None
update_role: Literal["assistant", "tool"] | None
function_call_results: list[Content] | None
@@ -1970,12 +1970,12 @@ def _handle_function_call_results(
*,
response: ChatResponse,
function_call_results: list[Content],
fcc_messages: list[ChatMessage],
fcc_messages: list[Message],
errors_in_a_row: int,
had_errors: bool,
max_errors: int,
) -> FunctionRequestResult:
from ._types import ChatMessage
from ._types import Message
if any(fccr.type in {"function_approval_request", "function_call"} for fccr in function_call_results):
# Only add items that aren't already in the message (e.g. function_approval_request wrappers).
@@ -1985,7 +1985,7 @@ def _handle_function_call_results(
if response.messages and response.messages[0].role == "assistant":
response.messages[0].contents.extend(new_items)
else:
response.messages.append(ChatMessage(role="assistant", contents=new_items))
response.messages.append(Message(role="assistant", contents=new_items))
return {
"action": "return",
"errors_in_a_row": errors_in_a_row,
@@ -2012,7 +2012,7 @@ def _handle_function_call_results(
else:
errors_in_a_row = 0
result_message = ChatMessage(role="tool", contents=function_call_results)
result_message = Message(role="tool", contents=function_call_results)
response.messages.append(result_message)
fcc_messages.extend(response.messages)
return {
@@ -2027,10 +2027,10 @@ def _handle_function_call_results(
async def _process_function_requests(
*,
response: ChatResponse | None,
prepped_messages: list[ChatMessage] | None,
prepped_messages: list[Message] | None,
tool_options: dict[str, Any] | None,
attempt_idx: int,
fcc_messages: list[ChatMessage] | None,
fcc_messages: list[Message] | None,
errors_in_a_row: int,
max_errors: int,
execute_function_calls: Callable[..., Awaitable[tuple[list[Content], bool, bool]]],
@@ -2139,7 +2139,7 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
@overload
def get_response(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage],
messages: str | Message | Sequence[str | Message],
*,
stream: Literal[False] = ...,
options: ChatOptions[ResponseModelBoundT],
@@ -2149,7 +2149,7 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
@overload
def get_response(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage],
messages: str | Message | Sequence[str | Message],
*,
stream: Literal[False] = ...,
options: OptionsCoT | ChatOptions[None] | None = None,
@@ -2159,7 +2159,7 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
@overload
def get_response(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage],
messages: str | Message | Sequence[str | Message],
*,
stream: Literal[True],
options: OptionsCoT | ChatOptions[Any] | None = None,
@@ -2168,7 +2168,7 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
def get_response(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage],
messages: str | Message | Sequence[str | Message],
*,
stream: bool = False,
options: OptionsCoT | ChatOptions[Any] | None = None,
@@ -2213,7 +2213,7 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
nonlocal filtered_kwargs
errors_in_a_row: int = 0
prepped_messages = prepare_messages(messages)
fcc_messages: list[ChatMessage] = []
fcc_messages: list[Message] = []
response: ChatResponse | None = None
for attempt_idx in range(
@@ -2307,7 +2307,7 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
nonlocal stream_result_hooks
errors_in_a_row: int = 0
prepped_messages = prepare_messages(messages)
fcc_messages: list[ChatMessage] = []
fcc_messages: list[Message] = []
response: ChatResponse | None = None
for attempt_idx in range(
python/packages/core/agent_framework/_types.py
+69 -69
View File
@@ -31,7 +31,6 @@ __all__ = [
"AgentResponse",
"AgentResponseUpdate",
"Annotation",
"ChatMessage",
"ChatOptions",
"ChatResponse",
"ChatResponseUpdate",
@@ -40,6 +39,7 @@ __all__ = [
"FinalT",
"FinishReason",
"FinishReasonLiteral",
"Message",
"OuterFinalT",
"OuterUpdateT",
"ResponseStream",
@@ -1420,14 +1420,14 @@ Known values: "system", "user", "assistant", "tool"
Examples:
.. code-block:: python
from agent_framework import ChatMessage
from agent_framework import Message
# Use string values directly
user_msg = ChatMessage("user", ["Hello"])
assistant_msg = ChatMessage("assistant", ["Hi there!"])
user_msg = Message("user", ["Hello"])
assistant_msg = Message("assistant", ["Hi there!"])
# Custom roles are also supported
custom_msg = ChatMessage("custom", ["Custom role message"])
custom_msg = Message("custom", ["Custom role message"])
# Compare roles directly as strings
if user_msg.role == "user":
@@ -1461,10 +1461,10 @@ Examples:
"""
# region ChatMessage
# region Message
class ChatMessage(SerializationMixin):
class Message(SerializationMixin):
"""Represents a chat message.
Attributes:
@@ -1479,17 +1479,17 @@ class ChatMessage(SerializationMixin):
Examples:
.. code-block:: python
from agent_framework import ChatMessage, Content
from agent_framework import Message, Content
# Create a message with text content
user_msg = ChatMessage("user", ["What's the weather?"])
user_msg = Message("user", ["What's the weather?"])
print(user_msg.text) # "What's the weather?"
# Create a system message
system_msg = ChatMessage("system", ["You are a helpful assistant."])
system_msg = Message("system", ["You are a helpful assistant."])
# Create a message with mixed content types
assistant_msg = ChatMessage(
assistant_msg = Message(
"assistant",
["The weather is sunny!", Content.from_image_uri("https://...")],
)
@@ -1499,13 +1499,13 @@ class ChatMessage(SerializationMixin):
msg_dict = user_msg.to_dict()
# {'type': 'chat_message', 'role': 'user',
# 'contents': [{'type': 'text', 'text': "What's the weather?"}], 'additional_properties': {}}
restored_msg = ChatMessage.from_dict(msg_dict)
restored_msg = Message.from_dict(msg_dict)
print(restored_msg.text) # "What's the weather?"
# Serialization - to_json and from_json
msg_json = user_msg.to_json()
# '{"type": "chat_message", "role": "user", "contents": [...], ...}'
restored_from_json = ChatMessage.from_json(msg_json)
restored_from_json = Message.from_json(msg_json)
print(restored_from_json.role) # "user"
"""
@@ -1523,7 +1523,7 @@ class ChatMessage(SerializationMixin):
additional_properties: MutableMapping[str, Any] | None = None,
raw_representation: Any | None = None,
) -> None:
"""Initialize ChatMessage.
"""Initialize Message.
Args:
role: The role of the author of the message (e.g., "user", "assistant", "system", "tool").
@@ -1568,86 +1568,86 @@ class ChatMessage(SerializationMixin):
def prepare_messages(
messages: str | Content | ChatMessage | Sequence[str | Content | ChatMessage],
messages: str | Content | Message | Sequence[str | Content | Message],
system_instructions: str | Sequence[str] | None = None,
) -> list[ChatMessage]:
"""Convert various message input formats into a list of ChatMessage objects.
) -> list[Message]:
"""Convert various message input formats into a list of Message objects.
Args:
messages: The input messages in various supported formats. Can be:
- A string (converted to a user message)
- A Content object (wrapped in a user ChatMessage)
- A ChatMessage object
- A Content object (wrapped in a user Message)
- A Message object
- A sequence containing any mix of the above
system_instructions: The system instructions. They will be inserted to the start of the messages list.
Returns:
A list of ChatMessage objects.
A list of Message objects.
"""
if system_instructions is not None:
if isinstance(system_instructions, str):
system_instructions = [system_instructions]
system_instruction_messages = [ChatMessage("system", [instr]) for instr in system_instructions]
system_instruction_messages = [Message("system", [instr]) for instr in system_instructions]
else:
system_instruction_messages = []
if isinstance(messages, str):
return [*system_instruction_messages, ChatMessage("user", [messages])]
return [*system_instruction_messages, Message("user", [messages])]
if isinstance(messages, Content):
return [*system_instruction_messages, ChatMessage("user", [messages])]
if isinstance(messages, ChatMessage):
return [*system_instruction_messages, Message("user", [messages])]
if isinstance(messages, Message):
return [*system_instruction_messages, messages]
return_messages: list[ChatMessage] = system_instruction_messages
return_messages: list[Message] = system_instruction_messages
for msg in messages:
if isinstance(msg, (str, Content)):
msg = ChatMessage("user", [msg])
msg = Message("user", [msg])
return_messages.append(msg)
return return_messages
def normalize_messages(
messages: str | Content | ChatMessage | Sequence[str | Content | ChatMessage] | None = None,
) -> list[ChatMessage]:
"""Normalize message inputs to a list of ChatMessage objects.
messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
) -> list[Message]:
"""Normalize message inputs to a list of Message objects.
Args:
messages: The input messages in various supported formats. Can be:
- None (returns empty list)
- A string (converted to a user message)
- A Content object (wrapped in a user ChatMessage)
- A ChatMessage object
- A Content object (wrapped in a user Message)
- A Message object
- A sequence containing any mix of the above
Returns:
A list of ChatMessage objects.
A list of Message objects.
"""
if messages is None:
return []
if isinstance(messages, str):
return [ChatMessage("user", [messages])]
return [Message("user", [messages])]
if isinstance(messages, Content):
return [ChatMessage("user", [messages])]
return [Message("user", [messages])]
if isinstance(messages, ChatMessage):
if isinstance(messages, Message):
return [messages]
result: list[ChatMessage] = []
result: list[Message] = []
for msg in messages:
if isinstance(msg, (str, Content)):
result.append(ChatMessage("user", [msg]))
result.append(Message("user", [msg]))
else:
result.append(msg)
return result
def prepend_instructions_to_messages(
messages: list[ChatMessage],
messages: list[Message],
instructions: str | Sequence[str] | None,
role: RoleLiteral | str = "system",
) -> list[ChatMessage]:
) -> list[Message]:
"""Prepend instructions to a list of messages with a specified role.
This is a helper method for chat clients that need to add instructions
@@ -1655,7 +1655,7 @@ def prepend_instructions_to_messages(
instructions (e.g., OpenAI uses "system", some providers might use "user").
Args:
messages: The existing list of ChatMessage objects.
messages: The existing list of Message objects.
instructions: The instructions to prepend. Can be a single string or a sequence of strings.
role: The role to use for the instruction messages. Defaults to "system".
@@ -1665,9 +1665,9 @@ def prepend_instructions_to_messages(
Examples:
.. code-block:: python
from agent_framework import prepend_instructions_to_messages, ChatMessage
from agent_framework import prepend_instructions_to_messages, Message
messages = [ChatMessage("user", ["Hello"])]
messages = [Message("user", ["Hello"])]
instructions = "You are a helpful assistant"
# Prepend as system message (default)
@@ -1682,7 +1682,7 @@ def prepend_instructions_to_messages(
if isinstance(instructions, str):
instructions = [instructions]
instruction_messages = [ChatMessage(role, [instr]) for instr in instructions]
instruction_messages = [Message(role, [instr]) for instr in instructions]
return [*instruction_messages, *messages]
@@ -1704,7 +1704,7 @@ def _process_update(response: ChatResponse | AgentResponse, update: ChatResponse
is_new_message = True
if is_new_message:
message = ChatMessage("assistant", [])
message = Message("assistant", [])
response.messages.append(message)
else:
message = response.messages[-1]
@@ -1847,17 +1847,17 @@ class ChatResponse(SerializationMixin, Generic[ResponseModelT]):
raw_representation: The raw representation of the chat response from an underlying implementation.
Note:
The `author_name` attribute is available on the `ChatMessage` objects inside `messages`,
The `author_name` attribute is available on the `Message` objects inside `messages`,
not on the `ChatResponse` itself. Use `response.messages[0].author_name` to access
the author name of individual messages.
Examples:
.. code-block:: python
from agent_framework import ChatResponse, ChatMessage
from agent_framework import ChatResponse, Message
# Create a response with messages
msg = ChatMessage("assistant", ["The weather is sunny."])
msg = Message("assistant", ["The weather is sunny."])
response = ChatResponse(
messages=[msg],
finish_reason="stop",
@@ -1887,7 +1887,7 @@ class ChatResponse(SerializationMixin, Generic[ResponseModelT]):
def __init__(
self,
*,
messages: ChatMessage | Sequence[ChatMessage] | None = None,
messages: Message | Sequence[Message] | None = None,
response_id: str | None = None,
conversation_id: str | None = None,
model_id: str | None = None,
@@ -1903,7 +1903,7 @@ class ChatResponse(SerializationMixin, Generic[ResponseModelT]):
"""Initializes a ChatResponse with the provided parameters.
Keyword Args:
messages: A single ChatMessage or sequence of ChatMessage objects to include in the response.
messages: A single Message or sequence of Message objects to include in the response.
response_id: Optional ID of the chat response.
conversation_id: Optional identifier for the state of the conversation.
model_id: Optional model ID used in the creation of the chat response.
@@ -1918,17 +1918,17 @@ class ChatResponse(SerializationMixin, Generic[ResponseModelT]):
raw_representation: Optional raw representation of the chat response from an underlying implementation.
"""
if messages is None:
self.messages: list[ChatMessage] = []
elif isinstance(messages, ChatMessage):
self.messages: list[Message] = []
elif isinstance(messages, Message):
self.messages = [messages]
else:
# Handle both ChatMessage objects and dicts (for from_dict support)
processed_messages: list[ChatMessage] = []
# Handle both Message objects and dicts (for from_dict support)
processed_messages: list[Message] = []
for msg in messages:
if isinstance(msg, ChatMessage):
if isinstance(msg, Message):
processed_messages.append(msg)
elif isinstance(msg, dict):
processed_messages.append(ChatMessage.from_dict(msg))
processed_messages.append(Message.from_dict(msg))
else:
processed_messages.append(msg)
self.messages = processed_messages
@@ -2057,7 +2057,7 @@ class ChatResponse(SerializationMixin, Generic[ResponseModelT]):
@property
def text(self) -> str:
"""Returns the concatenated text of all messages in the response."""
return ("\n".join(message.text for message in self.messages if isinstance(message, ChatMessage))).strip()
return ("\n".join(message.text for message in self.messages if isinstance(message, Message))).strip()
@property
def value(self) -> ResponseModelT | None:
@@ -2096,7 +2096,7 @@ class ChatResponseUpdate(SerializationMixin):
author_name: The name of the author of the response update. This is primarily used in
multi-agent scenarios to identify which agent or participant generated the response.
When updates are combined into a `ChatResponse`, the `author_name` is propagated
to the resulting `ChatMessage` objects.
to the resulting `Message` objects.
response_id: The ID of the response of which this update is a part.
message_id: The ID of the message of which this update is a part.
conversation_id: An identifier for the state of the conversation of which this update is a part.
@@ -2217,17 +2217,17 @@ class AgentResponse(SerializationMixin, Generic[ResponseModelT]):
messages in scenarios involving function calls, RAG retrievals, or complex logic.
Note:
The `author_name` attribute is available on the `ChatMessage` objects inside `messages`,
The `author_name` attribute is available on the `Message` objects inside `messages`,
not on the `AgentResponse` itself. Use `response.messages[0].author_name` to access
the author name of individual messages.
Examples:
.. code-block:: python
from agent_framework import AgentResponse, ChatMessage
from agent_framework import AgentResponse, Message
# Create agent response
msg = ChatMessage("assistant", ["Task completed successfully."])
msg = Message("assistant", ["Task completed successfully."])
response = AgentResponse(messages=[msg], response_id="run_123")
print(response.text) # "Task completed successfully."
@@ -2258,7 +2258,7 @@ class AgentResponse(SerializationMixin, Generic[ResponseModelT]):
def __init__(
self,
*,
messages: ChatMessage | Sequence[ChatMessage] | None = None,
messages: Message | Sequence[Message] | None = None,
response_id: str | None = None,
agent_id: str | None = None,
created_at: CreatedAtT | None = None,
@@ -2272,7 +2272,7 @@ class AgentResponse(SerializationMixin, Generic[ResponseModelT]):
"""Initialize an AgentResponse.
Keyword Args:
messages: A single ChatMessage or sequence of ChatMessage objects to include in the response.
messages: A single Message or sequence of Message objects to include in the response.
response_id: The ID of the chat response.
agent_id: The identifier of the agent that produced this response. Useful in multi-agent
scenarios to track which agent generated the response.
@@ -2286,17 +2286,17 @@ class AgentResponse(SerializationMixin, Generic[ResponseModelT]):
raw_representation: The raw representation of the chat response from an underlying implementation.
"""
if messages is None:
self.messages: list[ChatMessage] = []
elif isinstance(messages, ChatMessage):
self.messages: list[Message] = []
elif isinstance(messages, Message):
self.messages = [messages]
else:
# Handle both ChatMessage objects and dicts (for from_dict support)
processed_messages: list[ChatMessage] = []
# Handle both Message objects and dicts (for from_dict support)
processed_messages: list[Message] = []
for msg in messages:
if isinstance(msg, ChatMessage):
if isinstance(msg, Message):
processed_messages.append(msg)
elif isinstance(msg, dict):
processed_messages.append(ChatMessage.from_dict(msg))
processed_messages.append(Message.from_dict(msg))
else:
processed_messages.append(msg)
self.messages = processed_messages
@@ -2440,7 +2440,7 @@ class AgentResponseUpdate(SerializationMixin):
role: The role of the author of the response update.
author_name: The name of the author of the response update. In multi-agent scenarios,
this identifies which agent generated this update. When updates are combined into
an `AgentResponse`, the `author_name` is propagated to the resulting `ChatMessage` objects.
an `AgentResponse`, the `author_name` is propagated to the resulting `Message` objects.
agent_id: The identifier of the agent that produced this update. Useful in multi-agent
scenarios to track which agent generated specific parts of the response.
response_id: The ID of the response of which this update is a part.
python/packages/core/agent_framework/_workflows/__init__.py
+2 -2
View File
@@ -52,8 +52,8 @@ from ._request_info_mixin import response_handler
from ._runner import Runner
from ._runner_context import (
InProcRunnerContext,
Message,
RunnerContext,
WorkflowMessage,
)
from ._validation import (
EdgeDuplicationError,
@@ -92,7 +92,6 @@ __all__ = [
"GraphConnectivityError",
"InMemoryCheckpointStorage",
"InProcRunnerContext",
"Message",
"Runner",
"RunnerContext",
"SingleEdgeGroup",
@@ -117,6 +116,7 @@ __all__ = [
"WorkflowEventType",
"WorkflowException",
"WorkflowExecutor",
"WorkflowMessage",
"WorkflowRunResult",
"WorkflowRunState",
"WorkflowRunnerException",
python/packages/core/agent_framework/_workflows/_agent.py
+28 -28
View File
@@ -16,8 +16,8 @@ from agent_framework import (
AgentResponseUpdate,
AgentThread,
BaseAgent,
ChatMessage,
Content,
Message,
UsageDetails,
)
@@ -107,8 +107,8 @@ class WorkflowAgent(BaseAgent):
except KeyError as exc: # Defensive: workflow lacks a configured entry point
raise ValueError("Workflow's start executor is not defined.") from exc
if not any(is_type_compatible(list[ChatMessage], input_type) for input_type in start_executor.input_types):
raise ValueError("Workflow's start executor cannot handle list[ChatMessage]")
if not any(is_type_compatible(list[Message], input_type) for input_type in start_executor.input_types):
raise ValueError("Workflow's start executor cannot handle list[Message]")
super().__init__(id=id, name=name, description=description, **kwargs)
self._workflow: Workflow = workflow
@@ -127,7 +127,7 @@ class WorkflowAgent(BaseAgent):
@overload
def run(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
messages: str | Message | list[str] | list[Message] | None = None,
*,
stream: Literal[True],
thread: AgentThread | None = None,
@@ -139,7 +139,7 @@ class WorkflowAgent(BaseAgent):
@overload
async def run(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
messages: str | Message | list[str] | list[Message] | None = None,
*,
stream: Literal[False] = ...,
thread: AgentThread | None = None,
@@ -150,7 +150,7 @@ class WorkflowAgent(BaseAgent):
def run(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
messages: str | Message | list[str] | list[Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
@@ -202,7 +202,7 @@ class WorkflowAgent(BaseAgent):
async def _run_non_streaming(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
messages: str | Message | list[str] | list[Message] | None = None,
*,
thread: AgentThread | None = None,
checkpoint_id: str | None = None,
@@ -225,7 +225,7 @@ class WorkflowAgent(BaseAgent):
async def _run_streaming(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
messages: str | Message | list[str] | list[Message] | None = None,
*,
thread: AgentThread | None = None,
checkpoint_id: str | None = None,
@@ -257,7 +257,7 @@ class WorkflowAgent(BaseAgent):
async def _run_impl(
self,
input_messages: list[ChatMessage],
input_messages: list[Message],
response_id: str,
thread: AgentThread,
checkpoint_id: str | None = None,
@@ -289,7 +289,7 @@ class WorkflowAgent(BaseAgent):
async def _run_stream_impl(
self,
input_messages: list[ChatMessage],
input_messages: list[Message],
response_id: str,
thread: AgentThread,
checkpoint_id: str | None = None,
@@ -319,7 +319,7 @@ class WorkflowAgent(BaseAgent):
async def _run_core(
self,
input_messages: list[ChatMessage],
input_messages: list[Message],
thread: AgentThread,
checkpoint_id: str | None,
checkpoint_storage: CheckpointStorage | None,
@@ -393,8 +393,8 @@ class WorkflowAgent(BaseAgent):
async def _build_conversation_messages(
self,
thread: AgentThread,
input_messages: list[ChatMessage],
) -> list[ChatMessage]:
input_messages: list[Message],
) -> list[Message]:
"""Build the complete conversation by prepending thread history to input messages.
Args:
@@ -402,9 +402,9 @@ class WorkflowAgent(BaseAgent):
input_messages: The new input messages to append.
Returns:
A list of ChatMessage objects representing the full conversation.
A list of Message objects representing the full conversation.
"""
conversation_messages: list[ChatMessage] = []
conversation_messages: list[Message] = []
if thread.message_store:
history = await thread.message_store.list_messages()
if history:
@@ -412,7 +412,7 @@ class WorkflowAgent(BaseAgent):
conversation_messages.extend(input_messages)
return conversation_messages
def _process_pending_requests(self, input_messages: list[ChatMessage]) -> dict[str, Any]:
def _process_pending_requests(self, input_messages: list[Message]) -> dict[str, Any]:
"""Process pending requests by extracting function responses and updating state.
Args:
@@ -444,7 +444,7 @@ class WorkflowAgent(BaseAgent):
output_events: list[WorkflowEvent[Any]],
) -> AgentResponse:
"""Convert a list of workflow output events to an AgentResponse."""
messages: list[ChatMessage] = []
messages: list[Message] = []
raw_representations: list[object] = []
merged_usage: UsageDetails | None = None
latest_created_at: str | None = None
@@ -453,7 +453,7 @@ class WorkflowAgent(BaseAgent):
if output_event.type == "request_info":
function_call, approval_request = self._process_request_info_event(output_event)
messages.append(
ChatMessage(
Message(
contents=[function_call, approval_request],
role="assistant",
author_name=output_event.source_executor_id,
@@ -484,11 +484,11 @@ class WorkflowAgent(BaseAgent):
if data.created_at
else latest_created_at
)
elif isinstance(data, ChatMessage):
elif isinstance(data, Message):
messages.append(data)
raw_representations.append(data.raw_representation)
elif is_instance_of(data, list[ChatMessage]):
chat_messages = cast(list[ChatMessage], data)
elif is_instance_of(data, list[Message]):
chat_messages = cast(list[Message], data)
messages.extend(chat_messages)
raw_representations.append(data)
else:
@@ -497,7 +497,7 @@ class WorkflowAgent(BaseAgent):
continue
messages.append(
ChatMessage(
Message(
contents=contents,
role="assistant",
author_name=output_event.executor_id,
@@ -591,7 +591,7 @@ class WorkflowAgent(BaseAgent):
)
)
return updates
if isinstance(data, ChatMessage):
if isinstance(data, Message):
return [
AgentResponseUpdate(
contents=list(data.contents),
@@ -603,9 +603,9 @@ class WorkflowAgent(BaseAgent):
raw_representation=data,
)
]
if is_instance_of(data, list[ChatMessage]):
# Convert each ChatMessage to an AgentResponseUpdate
chat_messages = cast(list[ChatMessage], data)
if is_instance_of(data, list[Message]):
# Convert each Message to an AgentResponseUpdate
chat_messages = cast(list[Message], data)
updates = []
for msg in chat_messages:
updates.append(
@@ -669,7 +669,7 @@ class WorkflowAgent(BaseAgent):
# Ignore workflow-internal events
return []
def _extract_function_responses(self, input_messages: list[ChatMessage]) -> dict[str, Any]:
def _extract_function_responses(self, input_messages: list[Message]) -> dict[str, Any]:
"""Extract function responses from input messages."""
function_responses: dict[str, Any] = {}
for message in input_messages:
@@ -820,7 +820,7 @@ class WorkflowAgent(BaseAgent):
)
# PHASE 2: CONVERT GROUPED UPDATES TO RESPONSES AND MERGE
final_messages: list[ChatMessage] = []
final_messages: list[Message] = []
merged_usage: UsageDetails | None = None
latest_created_at: str | None = None
merged_additional_properties: dict[str, Any] | None = None
python/packages/core/agent_framework/_workflows/_agent_executor.py
+12 -12
View File
@@ -11,7 +11,7 @@ from agent_framework import Content
from .._agents import SupportsAgentRun
from .._threads import AgentThread
from .._types import AgentResponse, AgentResponseUpdate, ChatMessage
from .._types import AgentResponse, AgentResponseUpdate, Message
from ._agent_utils import resolve_agent_id
from ._checkpoint_encoding import decode_checkpoint_value, encode_checkpoint_value
from ._const import WORKFLOW_RUN_KWARGS_KEY
@@ -40,7 +40,7 @@ class AgentExecutorRequest:
If False, the messages will be saved to the executor's cache but not sent to the agent.
"""
messages: list[ChatMessage]
messages: list[Message]
should_respond: bool = True
@@ -58,7 +58,7 @@ class AgentExecutorResponse:
executor_id: str
agent_response: AgentResponse
full_conversation: list[ChatMessage] | None = None
full_conversation: list[Message] | None = None
class AgentExecutor(Executor):
@@ -104,9 +104,9 @@ class AgentExecutor(Executor):
self._pending_responses_to_agent: list[Content] = []
# AgentExecutor maintains an internal cache of messages in between runs
self._cache: list[ChatMessage] = []
self._cache: list[Message] = []
# This tracks the full conversation after each run
self._full_conversation: list[ChatMessage] = []
self._full_conversation: list[Message] = []
@property
def description(self) -> str | None:
@@ -157,20 +157,20 @@ class AgentExecutor(Executor):
@handler
async def from_message(
self,
message: ChatMessage,
message: Message,
ctx: WorkflowContext[AgentExecutorResponse, AgentResponse | AgentResponseUpdate],
) -> None:
"""Accept a single ChatMessage as input."""
"""Accept a single Message as input."""
self._cache = normalize_messages_input(message)
await self._run_agent_and_emit(ctx)
@handler
async def from_messages(
self,
messages: list[str | ChatMessage],
messages: list[str | Message],
ctx: WorkflowContext[AgentExecutorResponse, AgentResponse | AgentResponseUpdate],
) -> None:
"""Accept a list of chat inputs (strings or ChatMessage) as conversation context."""
"""Accept a list of chat inputs (strings or Message) as conversation context."""
self._cache = normalize_messages_input(messages)
await self._run_agent_and_emit(ctx)
@@ -198,7 +198,7 @@ class AgentExecutor(Executor):
# Use role="tool" for function_result responses (from declaration-only tools)
# so the LLM receives proper tool results instead of orphaned tool_calls.
role = "tool" if all(r.type == "function_result" for r in self._pending_responses_to_agent) else "user"
self._cache = normalize_messages_input(ChatMessage(role=role, contents=self._pending_responses_to_agent))
self._cache = normalize_messages_input(Message(role=role, contents=self._pending_responses_to_agent))
self._pending_responses_to_agent.clear()
await self._run_agent_and_emit(ctx)
@@ -216,8 +216,8 @@ class AgentExecutor(Executor):
"""
# Check if using AzureAIAgentClient with server-side thread and warn about checkpointing limitations
if is_chat_agent(self._agent) and self._agent_thread.service_thread_id is not None:
client_class_name = self._agent.chat_client.__class__.__name__
client_module = self._agent.chat_client.__class__.__module__
client_class_name = self._agent.client.__class__.__name__
client_module = self._agent.client.__class__.__module__
if client_class_name == "AzureAIAgentClient" and "azure_ai" in client_module:
logger.warning(
python/packages/core/agent_framework/_workflows/_conversation_history.py
+4 -4
View File
@@ -2,16 +2,16 @@
"""Helpers for managing chat conversation history.
These utilities operate on standard `list[ChatMessage]` collections and simple
These utilities operate on standard `list[Message]` collections and simple
dictionary snapshots so orchestrators can share logic without new mixins.
"""
from collections.abc import Sequence
from .._types import ChatMessage
from .._types import Message
def latest_user_message(conversation: Sequence[ChatMessage]) -> ChatMessage:
def latest_user_message(conversation: Sequence[Message]) -> Message:
"""Return the most recent user-authored message from `conversation`."""
for message in reversed(conversation):
role_value = getattr(message.role, "value", message.role)
@@ -20,7 +20,7 @@ def latest_user_message(conversation: Sequence[ChatMessage]) -> ChatMessage:
raise ValueError("No user message in conversation")
def ensure_author(message: ChatMessage, fallback: str) -> ChatMessage:
def ensure_author(message: Message, fallback: str) -> Message:
"""Attach `fallback` author if message is missing `author_name`."""
message.author_name = message.author_name or fallback
return message
python/packages/core/agent_framework/_workflows/_conversation_state.py
+6 -6
View File
@@ -3,20 +3,20 @@
from collections.abc import Iterable
from typing import Any, cast
from agent_framework import ChatMessage
from agent_framework import Message
from ._checkpoint_encoding import decode_checkpoint_value, encode_checkpoint_value
"""Utilities for serializing and deserializing chat conversations for persistence.
These helpers convert rich `ChatMessage` instances to checkpoint-friendly payloads
These helpers convert rich `Message` instances to checkpoint-friendly payloads
using the same encoding primitives as the workflow runner. This preserves
`additional_properties` and other metadata without relying on unsafe mechanisms
such as pickling.
"""
def encode_chat_messages(messages: Iterable[ChatMessage]) -> list[dict[str, Any]]:
def encode_chat_messages(messages: Iterable[Message]) -> list[dict[str, Any]]:
"""Serialize chat messages into checkpoint-safe payloads."""
encoded: list[dict[str, Any]] = []
for message in messages:
@@ -32,9 +32,9 @@ def encode_chat_messages(messages: Iterable[ChatMessage]) -> list[dict[str, Any]
return encoded
def decode_chat_messages(payload: Iterable[dict[str, Any]]) -> list[ChatMessage]:
def decode_chat_messages(payload: Iterable[dict[str, Any]]) -> list[Message]:
"""Restore chat messages from checkpoint-safe payloads."""
restored: list[ChatMessage] = []
restored: list[Message] = []
for item in payload:
if not isinstance(item, dict):
continue
@@ -64,7 +64,7 @@ def decode_chat_messages(payload: Iterable[dict[str, Any]]) -> list[ChatMessage]
additional[key] = decode_checkpoint_value(value)
restored.append(
ChatMessage( # type: ignore[call-overload]
Message( # type: ignore[call-overload]
role=role,
contents=contents,
author_name=item.get("author_name"),
python/packages/core/agent_framework/_workflows/_edge_runner.py
+10 -10
View File
@@ -18,7 +18,7 @@ from ._edge import (
SwitchCaseEdgeGroup,
)
from ._executor import Executor
from ._runner_context import Message, RunnerContext
from ._runner_context import RunnerContext, WorkflowMessage
from ._state import State
logger = logging.getLogger(__name__)
@@ -38,7 +38,7 @@ class EdgeRunner(ABC):
self._executors = executors
@abstractmethod
async def send_message(self, message: Message, state: State, ctx: RunnerContext) -> bool:
async def send_message(self, message: WorkflowMessage, state: State, ctx: RunnerContext) -> bool:
"""Send a message through the edge group.
Args:
@@ -52,7 +52,7 @@ class EdgeRunner(ABC):
"""
raise NotImplementedError
def _can_handle(self, executor_id: str, message: Message) -> bool:
def _can_handle(self, executor_id: str, message: WorkflowMessage) -> bool:
"""Check if an executor can handle the given message data."""
if executor_id not in self._executors:
return False
@@ -62,7 +62,7 @@ class EdgeRunner(ABC):
self,
target_id: str,
source_ids: list[str],
message: Message,
message: WorkflowMessage,
state: State,
ctx: RunnerContext,
) -> None:
@@ -90,7 +90,7 @@ class SingleEdgeRunner(EdgeRunner):
super().__init__(edge_group, executors)
self._edge = edge_group.edges[0]
async def send_message(self, message: Message, state: State, ctx: RunnerContext) -> bool:
async def send_message(self, message: WorkflowMessage, state: State, ctx: RunnerContext) -> bool:
"""Send a message through the single edge."""
should_execute = False
target_id: str | None = None
@@ -162,7 +162,7 @@ class FanOutEdgeRunner(EdgeRunner):
Callable[[Any, list[str]], list[str]] | None, getattr(edge_group, "selection_func", None)
)
async def send_message(self, message: Message, state: State, ctx: RunnerContext) -> bool:
async def send_message(self, message: WorkflowMessage, state: State, ctx: RunnerContext) -> bool:
"""Send a message through all edges in the fan-out edge group."""
deliverable_edges: list[Edge] = []
single_target_edge: Edge | None = None
@@ -283,9 +283,9 @@ class FanInEdgeRunner(EdgeRunner):
self._edges = edge_group.edges
# Buffer to hold messages before sending them to the target executor
# Key is the source executor ID, value is a list of messages
self._buffer: dict[str, list[Message]] = defaultdict(list)
self._buffer: dict[str, list[WorkflowMessage]] = defaultdict(list)
async def send_message(self, message: Message, state: State, ctx: RunnerContext) -> bool:
async def send_message(self, message: WorkflowMessage, state: State, ctx: RunnerContext) -> bool:
"""Send a message through all edges in the fan-in edge group."""
execution_data: dict[str, Any] | None = None
with create_edge_group_processing_span(
@@ -306,7 +306,7 @@ class FanInEdgeRunner(EdgeRunner):
# Check if target can handle list of message data (fan-in aggregates multiple messages)
if self._can_handle(
self._edges[0].target_id, Message(data=[message.data], source_id=message.source_id)
self._edges[0].target_id, WorkflowMessage(data=[message.data], source_id=message.source_id)
):
# If the edge can handle the data, buffer the message
self._buffer[message.source_id].append(message)
@@ -334,7 +334,7 @@ class FanInEdgeRunner(EdgeRunner):
source_span_ids = [msg.source_span_id for msg in messages_to_send if msg.source_span_id]
# Create a new Message object for the aggregated data
aggregated_message = Message(
aggregated_message = WorkflowMessage(
data=aggregated_data,
source_id=self._edge_group.__class__.__name__, # This won't be used in self._execute_on_target.
trace_contexts=trace_contexts,
python/packages/core/agent_framework/_workflows/_executor.py
+6 -6
View File
@@ -17,7 +17,7 @@ from ._events import (
)
from ._model_utils import DictConvertible
from ._request_info_mixin import RequestInfoMixin
from ._runner_context import Message, MessageType, RunnerContext
from ._runner_context import MessageType, RunnerContext, WorkflowMessage
from ._state import State
from ._typing_utils import is_instance_of, normalize_type_to_list, resolve_type_annotation
from ._workflow_context import WorkflowContext, validate_workflow_context_annotation
@@ -244,7 +244,7 @@ class Executor(RequestInfoMixin, DictConvertible):
with create_processing_span(
self.id,
self.__class__.__name__,
str(MessageType.STANDARD if not isinstance(message, Message) else message.type),
str(MessageType.STANDARD if not isinstance(message, WorkflowMessage) else message.type),
type(message).__name__,
source_trace_contexts=trace_contexts,
source_span_ids=source_span_ids,
@@ -253,7 +253,7 @@ class Executor(RequestInfoMixin, DictConvertible):
handler = self._find_handler(message)
original_message = message
if isinstance(message, Message):
if isinstance(message, WorkflowMessage):
# Unwrap raw data for handler call
message = message.data
@@ -265,7 +265,7 @@ class Executor(RequestInfoMixin, DictConvertible):
trace_contexts=trace_contexts,
source_span_ids=source_span_ids,
request_id=original_message.original_request_info_event.request_id
if isinstance(original_message, Message) and original_message.original_request_info_event
if isinstance(original_message, WorkflowMessage) and original_message.original_request_info_event
else None,
)
@@ -351,7 +351,7 @@ class Executor(RequestInfoMixin, DictConvertible):
# Add to unified handler specs list
self._handler_specs.append({**handler_spec})
def can_handle(self, message: Message) -> bool:
def can_handle(self, message: WorkflowMessage) -> bool:
"""Check if the executor can handle a given message type.
Args:
@@ -460,7 +460,7 @@ class Executor(RequestInfoMixin, DictConvertible):
Returns:
The handler function if found, None otherwise
"""
if isinstance(message, Message):
if isinstance(message, WorkflowMessage):
# Case where Message wrapper is passed instead of raw data
# Handler can be a standard handler or a response handler
if message.type == MessageType.STANDARD:
python/packages/core/agent_framework/_workflows/_message_utils.py
+12 -12
View File
@@ -4,38 +4,38 @@
from collections.abc import Sequence
from agent_framework import ChatMessage
from agent_framework import Message
def normalize_messages_input(
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
) -> list[ChatMessage]:
"""Normalize heterogeneous message inputs to a list of ChatMessage objects.
messages: str | Message | Sequence[str | Message] | None = None,
) -> list[Message]:
"""Normalize heterogeneous message inputs to a list of Message objects.
Args:
messages: String, ChatMessage, or sequence of either. None yields empty list.
messages: String, Message, or sequence of either. None yields empty list.
Returns:
List of ChatMessage instances suitable for workflow consumption.
List of Message instances suitable for workflow consumption.
"""
if messages is None:
return []
if isinstance(messages, str):
return [ChatMessage(role="user", text=messages)]
return [Message(role="user", text=messages)]
if isinstance(messages, ChatMessage):
if isinstance(messages, Message):
return [messages]
normalized: list[ChatMessage] = []
normalized: list[Message] = []
for item in messages:
if isinstance(item, str):
normalized.append(ChatMessage(role="user", text=item))
elif isinstance(item, ChatMessage):
normalized.append(Message(role="user", text=item))
elif isinstance(item, Message):
normalized.append(item)
else:
raise TypeError(
f"Messages sequence must contain only str or ChatMessage instances; found {type(item).__name__}."
f"Messages sequence must contain only str or Message instances; found {type(item).__name__}."
)
return normalized
python/packages/core/agent_framework/_workflows/_runner.py
+4 -4
View File
@@ -24,8 +24,8 @@ from ._exceptions import (
)
from ._executor import Executor
from ._runner_context import (
Message,
RunnerContext,
WorkflowMessage,
)
from ._state import State
@@ -162,14 +162,14 @@ class Runner:
self._running = False
async def _run_iteration(self) -> None:
async def _deliver_messages(source_executor_id: str, messages: list[Message]) -> None:
async def _deliver_messages(source_executor_id: str, messages: list[WorkflowMessage]) -> None:
"""Outer loop to concurrently deliver messages from all sources to their targets."""
async def _deliver_message_inner(edge_runner: EdgeRunner, message: Message) -> bool:
async def _deliver_message_inner(edge_runner: EdgeRunner, message: WorkflowMessage) -> bool:
"""Inner loop to deliver a single message through an edge runner."""
return await edge_runner.send_message(message, self._state, self._ctx)
def _normalize_message_payload(message: Message) -> None:
def _normalize_message_payload(message: WorkflowMessage) -> None:
data = message.data
if not isinstance(data, dict):
return
python/packages/core/agent_framework/_workflows/_runner_context.py
+23 -23
View File
@@ -29,25 +29,25 @@ T = TypeVar("T")
class MessageType(Enum):
"""Enumeration of message types in the workflow."""
"""Enumeration of WorkflowMessage types in the workflow."""
STANDARD = "standard"
"""A standard message between executors."""
"""A standard WorkflowMessage between executors."""
RESPONSE = "response"
"""A response message to a pending request."""
"""A response WorkflowMessage to a pending request."""
@dataclass
class Message:
"""A class representing a message in the workflow."""
class WorkflowMessage:
"""A class representing a WorkflowMessage in the workflow."""
data: Any
source_id: str
target_id: str | None = None
type: MessageType = MessageType.STANDARD
# OpenTelemetry trace context fields for message propagation
# OpenTelemetry trace context fields for WorkflowMessage propagation
# These are plural to support fan-in scenarios where multiple messages are aggregated
trace_contexts: list[dict[str, str]] | None = None # W3C Trace Context headers from multiple sources
source_span_ids: list[str] | None = None # Publishing span IDs for linking from multiple sources
@@ -67,7 +67,7 @@ class Message:
return self.source_span_ids[0] if self.source_span_ids else None
def to_dict(self) -> dict[str, Any]:
"""Convert the Message to a dictionary for serialization."""
"""Convert the WorkflowMessage to a dictionary for serialization."""
return {
"data": encode_checkpoint_value(self.data),
"source_id": self.source_id,
@@ -79,16 +79,16 @@ class Message:
}
@staticmethod
def from_dict(data: dict[str, Any]) -> Message:
"""Create a Message from a dictionary."""
def from_dict(data: dict[str, Any]) -> WorkflowMessage:
"""Create a WorkflowMessage from a dictionary."""
# Validation
if "data" not in data:
raise KeyError("Missing 'data' field in Message dictionary.")
raise KeyError("Missing 'data' field in WorkflowMessage dictionary.")
if "source_id" not in data:
raise KeyError("Missing 'source_id' field in Message dictionary.")
raise KeyError("Missing 'source_id' field in WorkflowMessage dictionary.")
return Message(
return WorkflowMessage(
data=decode_checkpoint_value(data["data"]),
source_id=data["source_id"],
target_id=data.get("target_id"),
@@ -119,15 +119,15 @@ class RunnerContext(Protocol):
If checkpoint storage is not configured, checkpoint methods may raise.
"""
async def send_message(self, message: Message) -> None:
"""Send a message from the executor to the context.
async def send_message(self, WorkflowMessage: WorkflowMessage) -> None:
"""Send a WorkflowMessage from the executor to the context.
Args:
message: The message to be sent.
WorkflowMessage: The WorkflowMessage to be sent.
"""
...
async def drain_messages(self) -> dict[str, list[Message]]:
async def drain_messages(self) -> dict[str, list[WorkflowMessage]]:
"""Drain all messages from the context.
Returns:
@@ -291,7 +291,7 @@ class InProcRunnerContext:
Args:
checkpoint_storage: Optional storage to enable checkpointing.
"""
self._messages: dict[str, list[Message]] = {}
self._messages: dict[str, list[WorkflowMessage]] = {}
# Event queue for immediate streaming of events
self._event_queue: asyncio.Queue[WorkflowEvent] = asyncio.Queue()
@@ -307,11 +307,11 @@ class InProcRunnerContext:
self._streaming: bool = False
# region Messaging and Events
async def send_message(self, message: Message) -> None:
self._messages.setdefault(message.source_id, [])
self._messages[message.source_id].append(message)
async def send_message(self, WorkflowMessage: WorkflowMessage) -> None:
self._messages.setdefault(WorkflowMessage.source_id, [])
self._messages[WorkflowMessage.source_id].append(WorkflowMessage)
async def drain_messages(self) -> dict[str, list[Message]]:
async def drain_messages(self) -> dict[str, list[WorkflowMessage]]:
messages = copy(self._messages)
self._messages.clear()
return messages
@@ -422,7 +422,7 @@ class InProcRunnerContext:
self._messages.clear()
messages_data = checkpoint.messages
for source_id, message_list in messages_data.items():
self._messages[source_id] = [Message.from_dict(msg) for msg in message_list]
self._messages[source_id] = [WorkflowMessage.from_dict(msg) for msg in message_list]
# Restore pending request info events
self._pending_request_info_events.clear()
@@ -504,7 +504,7 @@ class InProcRunnerContext:
source_executor_id = event.source_executor_id
# Create ResponseMessage instance
response_msg = Message(
response_msg = WorkflowMessage(
data=response,
source_id=INTERNAL_SOURCE_ID(source_executor_id),
target_id=source_executor_id,
python/packages/core/agent_framework/_workflows/_typing_utils.py
+6 -6
View File
@@ -3,19 +3,19 @@
from types import UnionType
from typing import Any, TypeGuard, Union, cast, get_args, get_origin
from .._agents import ChatAgent
from .._agents import Agent
def is_chat_agent(agent: Any) -> TypeGuard[ChatAgent]:
"""Check if the given agent is a ChatAgent.
def is_chat_agent(agent: Any) -> TypeGuard[Agent]:
"""Check if the given agent is a Agent.
Args:
agent (Any): The agent to check.
Returns:
TypeGuard[ChatAgent]: True if the agent is a ChatAgent, False otherwise.
TypeGuard[Agent]: True if the agent is a Agent, False otherwise.
"""
return isinstance(agent, ChatAgent)
return isinstance(agent, Agent)
def resolve_type_annotation(
@@ -255,7 +255,7 @@ def is_type_compatible(source_type: type | UnionType | Any, target_type: type |
A type is compatible if values of source_type can be assigned to variables of target_type.
For example:
- list[ChatMessage] is compatible with list[str | ChatMessage]
- list[Message] is compatible with list[str | Message]
- str is compatible with str | int
- int is compatible with Any
python/packages/core/agent_framework/_workflows/_workflow.py
+7 -7
View File
@@ -841,14 +841,14 @@ class Workflow(DictConvertible):
def as_agent(self, name: str | None = None) -> WorkflowAgent:
"""Create a WorkflowAgent that wraps this workflow.
The returned agent converts standard agent inputs (strings, ChatMessage, or lists of these)
into a list[ChatMessage] that is passed to the workflow's start executor. This conversion
The returned agent converts standard agent inputs (strings, Message, or lists of these)
into a list[Message] that is passed to the workflow's start executor. This conversion
happens in WorkflowAgent._normalize_messages() which transforms:
- str -> [ChatMessage(USER, [str])]
- ChatMessage -> [ChatMessage]
- list[str | ChatMessage] -> list[ChatMessage] (with string elements converted)
- str -> [Message(USER, [str])]
- Message -> [Message]
- list[str | Message] -> list[Message] (with string elements converted)
The workflow's start executor must accept list[ChatMessage] as an input type, otherwise
The workflow's start executor must accept list[Message] as an input type, otherwise
initialization will fail with a ValueError.
Args:
@@ -858,7 +858,7 @@ class Workflow(DictConvertible):
A WorkflowAgent instance that wraps this workflow.
Raises:
ValueError: If the workflow's start executor cannot handle list[ChatMessage] input.
ValueError: If the workflow's start executor cannot handle list[Message] input.
"""
# Import here to avoid circular imports
from ._agent import WorkflowAgent
python/packages/core/agent_framework/_workflows/_workflow_context.py
+2 -2
View File
@@ -19,7 +19,7 @@ from ._events import (
WorkflowEventSource,
_framework_event_origin, # type: ignore
)
from ._runner_context import Message, RunnerContext
from ._runner_context import RunnerContext, WorkflowMessage
from ._state import State
if TYPE_CHECKING:
@@ -321,7 +321,7 @@ class WorkflowContext(Generic[OutT, W_OutT]):
attributes[OtelAttr.MESSAGE_DESTINATION_EXECUTOR_ID] = target_id
with create_workflow_span(OtelAttr.MESSAGE_SEND_SPAN, attributes, kind=SpanKind.PRODUCER) as span:
# Create Message wrapper
msg = Message(data=message, source_id=self._executor_id, target_id=target_id)
msg = WorkflowMessage(data=message, source_id=self._executor_id, target_id=target_id)
# Track sent message for executor_completed event (type='executor_completed')
self._sent_messages.append(message)
python/packages/core/agent_framework/_workflows/_workflow_executor.py
+2 -2
View File
@@ -19,7 +19,7 @@ from ._events import (
)
from ._executor import Executor, handler
from ._request_info_mixin import response_handler
from ._runner_context import Message
from ._runner_context import WorkflowMessage
from ._typing_utils import is_instance_of
from ._workflow import WorkflowRunResult
from ._workflow_context import WorkflowContext
@@ -340,7 +340,7 @@ class WorkflowExecutor(Executor):
data["workflow"] = self.workflow.to_dict()
return data
def can_handle(self, message: Message) -> bool:
def can_handle(self, message: WorkflowMessage) -> bool:
"""Override can_handle to only accept messages that the wrapped workflow can handle.
This prevents the WorkflowExecutor from accepting messages that should go to other
python/packages/core/agent_framework/devui/__init__.py
+1
View File
@@ -14,6 +14,7 @@ _IMPORTS = [
"OpenAIResponse",
"ResponseStreamEvent",
"main",
"register_cleanup",
"serve",
"__version__",
]
python/packages/core/agent_framework/devui/__init__.pyi
+2
View File
@@ -10,6 +10,7 @@ from agent_framework_devui import (
ResponseStreamEvent,
__version__,
main,
register_cleanup,
serve,
)
@@ -23,5 +24,6 @@ __all__ = [
"ResponseStreamEvent",
"__version__",
"main",
"register_cleanup",
"serve",
]
python/packages/core/agent_framework/observability.py
+15 -15
View File
@@ -39,18 +39,18 @@ if TYPE_CHECKING: # pragma: no cover
from pydantic import BaseModel
from ._agents import SupportsAgentRun
from ._clients import ChatClientProtocol
from ._clients import SupportsChatGetResponse
from ._threads import AgentThread
from ._tools import FunctionTool
from ._types import (
AgentResponse,
AgentResponseUpdate,
ChatMessage,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
Content,
FinishReason,
Message,
ResponseStream,
)
@@ -71,7 +71,7 @@ __all__ = [
AgentT = TypeVar("AgentT", bound="SupportsAgentRun")
ChatClientT = TypeVar("ChatClientT", bound="ChatClientProtocol[Any]")
ChatClientT = TypeVar("ChatClientT", bound="SupportsChatGetResponse[Any]")
logger = get_logger()
@@ -122,7 +122,7 @@ OPERATION_DURATION_BUCKET_BOUNDARIES: Final[tuple[float, ...]] = (
#
# This is a workaround, we'll find a generic and better solution - see
# https://github.com/open-telemetry/semantic-conventions/issues/1701
class ChatMessageListTimestampFilter(logging.Filter):
class MessageListTimestampFilter(logging.Filter):
"""A filter to increment the timestamp of INFO logs by 1 microsecond."""
INDEX_KEY: ClassVar[str] = "chat_message_index"
@@ -135,7 +135,7 @@ class ChatMessageListTimestampFilter(logging.Filter):
return True
logger.addFilter(ChatMessageListTimestampFilter())
logger.addFilter(MessageListTimestampFilter())
class OtelAttr(str, Enum):
@@ -1070,7 +1070,7 @@ class ChatTelemetryLayer(Generic[OptionsCoT]):
@overload
def get_response(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage],
messages: str | Message | Sequence[str | Message],
*,
stream: Literal[False] = ...,
options: ChatOptions[ResponseModelBoundT],
@@ -1080,7 +1080,7 @@ class ChatTelemetryLayer(Generic[OptionsCoT]):
@overload
def get_response(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage],
messages: str | Message | Sequence[str | Message],
*,
stream: Literal[False] = ...,
options: OptionsCoT | ChatOptions[None] | None = None,
@@ -1090,7 +1090,7 @@ class ChatTelemetryLayer(Generic[OptionsCoT]):
@overload
def get_response(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage],
messages: str | Message | Sequence[str | Message],
*,
stream: Literal[True],
options: OptionsCoT | ChatOptions[Any] | None = None,
@@ -1099,7 +1099,7 @@ class ChatTelemetryLayer(Generic[OptionsCoT]):
def get_response(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage],
messages: str | Message | Sequence[str | Message],
*,
stream: bool = False,
options: OptionsCoT | ChatOptions[Any] | None = None,
@@ -1263,7 +1263,7 @@ class AgentTelemetryLayer:
@overload
def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[False] = ...,
thread: AgentThread | None = None,
@@ -1273,7 +1273,7 @@ class AgentTelemetryLayer:
@overload
def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: Literal[True],
thread: AgentThread | None = None,
@@ -1282,7 +1282,7 @@ class AgentTelemetryLayer:
def run(
self,
messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
messages: str | Message | Sequence[str | Message] | None = None,
*,
stream: bool = False,
thread: AgentThread | None = None,
@@ -1600,7 +1600,7 @@ def capture_exception(span: trace.Span, exception: Exception, timestamp: int | N
def _capture_messages(
span: trace.Span,
provider_name: str,
messages: str | ChatMessage | Sequence[str | ChatMessage],
messages: str | Message | Sequence[str | Message],
system_instructions: str | list[str] | None = None,
output: bool = False,
finish_reason: FinishReason | None = None,
@@ -1620,7 +1620,7 @@ def _capture_messages(
extra={
OtelAttr.EVENT_NAME: OtelAttr.CHOICE if output else ROLE_EVENT_MAP.get(message.role),
OtelAttr.PROVIDER_NAME: provider_name,
ChatMessageListTimestampFilter.INDEX_KEY: index,
MessageListTimestampFilter.INDEX_KEY: index,
},
)
if finish_reason:
@@ -1633,7 +1633,7 @@ def _capture_messages(
span.set_attribute(OtelAttr.SYSTEM_INSTRUCTIONS, json.dumps(otel_sys_instructions))
def _to_otel_message(message: ChatMessage) -> dict[str, Any]:
def _to_otel_message(message: Message) -> dict[str, Any]:
"""Create a otel representation of a message."""
return {"role": message.role, "parts": [_to_otel_part(content) for content in message.contents]}
python/packages/core/agent_framework/openai/_assistant_provider.py
+31 -31
View File
@@ -10,7 +10,7 @@ from openai import AsyncOpenAI
from openai.types.beta.assistant import Assistant
from pydantic import BaseModel, SecretStr, ValidationError
from .._agents import ChatAgent
from .._agents import Agent
from .._memory import ContextProvider
from .._middleware import MiddlewareTypes
from .._tools import FunctionTool, ToolProtocol
@@ -51,10 +51,10 @@ _ToolsType = (
class OpenAIAssistantProvider(Generic[OptionsCoT]):
"""Provider for creating ChatAgent instances from OpenAI Assistants API.
"""Provider for creating Agent instances from OpenAI Assistants API.
This provider allows you to create, retrieve, and wrap OpenAI Assistants
as ChatAgent instances for use in the agent framework.
as Agent instances for use in the agent framework.
Examples:
Basic usage with automatic client creation:
@@ -208,11 +208,11 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]):
default_options: OptionsCoT | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
) -> ChatAgent[OptionsCoT]:
"""Create a new assistant on OpenAI and return a ChatAgent.
) -> Agent[OptionsCoT]:
"""Create a new assistant on OpenAI and return a Agent.
This method creates a new assistant on the OpenAI service and wraps it
in a ChatAgent instance. The assistant will persist on OpenAI until deleted.
in a Agent instance. The assistant will persist on OpenAI until deleted.
Keyword Args:
name: The name of the assistant (required).
@@ -228,11 +228,11 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]):
default_options: A TypedDict containing default chat options for the agent.
These options are applied to every run unless overridden.
Include ``response_format`` here for structured output responses.
middleware: MiddlewareTypes for the ChatAgent.
context_provider: Context provider for the ChatAgent.
middleware: MiddlewareTypes for the Agent.
context_provider: Context provider for the Agent.
Returns:
A ChatAgent instance wrapping the created assistant.
A Agent instance wrapping the created assistant.
Raises:
ServiceInitializationError: If assistant creation fails.
@@ -297,7 +297,7 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]):
assistant = await self._client.beta.assistants.create(**create_params)
# Create ChatAgent - pass default_options which contains response_format
# Create Agent - pass default_options which contains response_format
return self._create_chat_agent_from_assistant(
assistant=assistant,
tools=normalized_tools,
@@ -316,11 +316,11 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]):
default_options: OptionsCoT | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
) -> ChatAgent[OptionsCoT]:
"""Retrieve an existing assistant by ID and return a ChatAgent.
) -> Agent[OptionsCoT]:
"""Retrieve an existing assistant by ID and return a Agent.
This method fetches an existing assistant from OpenAI by its ID
and wraps it in a ChatAgent instance.
and wraps it in a Agent instance.
Args:
assistant_id: The ID of the assistant to retrieve (e.g., "asst_123").
@@ -333,11 +333,11 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]):
instructions: Override the assistant's instructions (optional).
default_options: A TypedDict containing default chat options for the agent.
These options are applied to every run unless overridden.
middleware: MiddlewareTypes for the ChatAgent.
context_provider: Context provider for the ChatAgent.
middleware: MiddlewareTypes for the Agent.
context_provider: Context provider for the Agent.
Returns:
A ChatAgent instance wrapping the retrieved assistant.
A Agent instance wrapping the retrieved assistant.
Raises:
ServiceInitializationError: If the assistant cannot be retrieved.
@@ -382,11 +382,11 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]):
default_options: OptionsCoT | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
context_provider: ContextProvider | None = None,
) -> ChatAgent[OptionsCoT]:
"""Wrap an existing SDK Assistant object as a ChatAgent.
) -> Agent[OptionsCoT]:
"""Wrap an existing SDK Assistant object as a Agent.
This method does NOT make any HTTP calls. It simply wraps an already-
fetched Assistant object in a ChatAgent.
fetched Assistant object in a Agent.
Args:
assistant: The OpenAI Assistant SDK object to wrap.
@@ -398,11 +398,11 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]):
instructions: Override the assistant's instructions (optional).
default_options: A TypedDict containing default chat options for the agent.
These options are applied to every run unless overridden.
middleware: MiddlewareTypes for the ChatAgent.
context_provider: Context provider for the ChatAgent.
middleware: MiddlewareTypes for the Agent.
context_provider: Context provider for the Agent.
Returns:
A ChatAgent instance wrapping the assistant.
A Agent instance wrapping the assistant.
Raises:
ValueError: If required function tools are missing.
@@ -429,7 +429,7 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]):
# Merge hosted tools with user-provided function tools
merged_tools = self._merge_tools(assistant.tools or [], tools)
# Create ChatAgent
# Create Agent
return self._create_chat_agent_from_assistant(
assistant=assistant,
tools=merged_tools,
@@ -526,8 +526,8 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]):
context_provider: ContextProvider | None,
default_options: OptionsCoT | None = None,
**kwargs: Any,
) -> ChatAgent[OptionsCoT]:
"""Create a ChatAgent from an Assistant.
) -> Agent[OptionsCoT]:
"""Create a Agent from an Assistant.
Args:
assistant: The OpenAI Assistant object.
@@ -536,13 +536,13 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]):
middleware: MiddlewareTypes for the agent.
context_provider: Context provider for the agent.
default_options: Default chat options for the agent (may include response_format).
**kwargs: Additional arguments passed to ChatAgent.
**kwargs: Additional arguments passed to Agent.
Returns:
A configured ChatAgent instance.
A configured Agent instance.
"""
# Create the chat client with the assistant
chat_client = OpenAIAssistantsClient(
client = OpenAIAssistantsClient(
model_id=assistant.model,
assistant_id=assistant.id,
assistant_name=assistant.name,
@@ -553,9 +553,9 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]):
# Use instructions from assistant if not overridden
final_instructions = instructions if instructions is not None else assistant.instructions
# Create and return ChatAgent
return ChatAgent(
chat_client=chat_client,
# Create and return Agent
return Agent(
client=client,
id=assistant.id,
name=assistant.name,
description=assistant.description,
python/packages/core/agent_framework/openai/_assistants_client.py
+3 -3
View File
@@ -39,11 +39,11 @@ from .._tools import (
HostedFileSearchTool,
)
from .._types import (
ChatMessage,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
Content,
Message,
ResponseStream,
UsageDetails,
prepare_function_call_results,
@@ -352,7 +352,7 @@ class OpenAIAssistantsClient( # type: ignore[misc]
def _inner_get_response(
self,
*,
messages: Sequence[ChatMessage],
messages: Sequence[Message],
options: Mapping[str, Any],
stream: bool = False,
**kwargs: Any,
@@ -605,7 +605,7 @@ class OpenAIAssistantsClient( # type: ignore[misc]
def _prepare_options(
self,
messages: Sequence[ChatMessage],
messages: Sequence[Message],
options: Mapping[str, Any],
**kwargs: Any,
) -> tuple[dict[str, Any], list[Content] | None]:
python/packages/core/agent_framework/openai/_chat_client.py
+7 -7
View File
@@ -29,12 +29,12 @@ from .._tools import (
ToolProtocol,
)
from .._types import (
ChatMessage,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
Content,
FinishReason,
Message,
ResponseStream,
UsageDetails,
prepare_function_call_results,
@@ -158,7 +158,7 @@ class RawOpenAIChatClient( # type: ignore[misc]
def _inner_get_response(
self,
*,
messages: Sequence[ChatMessage],
messages: Sequence[Message],
options: Mapping[str, Any],
stream: bool = False,
**kwargs: Any,
@@ -252,7 +252,7 @@ class RawOpenAIChatClient( # type: ignore[misc]
ret_dict["web_search_options"] = web_search_options
return ret_dict
def _prepare_options(self, messages: Sequence[ChatMessage], options: Mapping[str, Any]) -> dict[str, Any]:
def _prepare_options(self, messages: Sequence[Message], options: Mapping[str, Any]) -> dict[str, Any]:
# Prepend instructions from options if they exist
from .._types import prepend_instructions_to_messages, validate_tool_mode
@@ -310,7 +310,7 @@ class RawOpenAIChatClient( # type: ignore[misc]
def _parse_response_from_openai(self, response: ChatCompletion, options: Mapping[str, Any]) -> ChatResponse:
"""Parse a response from OpenAI into a ChatResponse."""
response_metadata = self._get_metadata_from_chat_response(response)
messages: list[ChatMessage] = []
messages: list[Message] = []
finish_reason: FinishReason | None = None
for choice in response.choices:
response_metadata.update(self._get_metadata_from_chat_choice(choice))
@@ -323,7 +323,7 @@ class RawOpenAIChatClient( # type: ignore[misc]
contents.extend(parsed_tool_calls)
if reasoning_details := getattr(choice.message, "reasoning_details", None):
contents.append(Content.from_text_reasoning(protected_data=json.dumps(reasoning_details)))
messages.append(ChatMessage(role="assistant", contents=contents))
messages.append(Message(role="assistant", contents=contents))
return ChatResponse(
response_id=response.id,
created_at=datetime.fromtimestamp(response.created, tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
@@ -448,7 +448,7 @@ class RawOpenAIChatClient( # type: ignore[misc]
def _prepare_messages_for_openai(
self,
chat_messages: Sequence[ChatMessage],
chat_messages: Sequence[Message],
role_key: str = "role",
content_key: str = "content",
) -> list[dict[str, Any]]:
@@ -476,7 +476,7 @@ class RawOpenAIChatClient( # type: ignore[misc]
# region Parsers
def _prepare_message_for_openai(self, message: ChatMessage) -> list[dict[str, Any]]:
def _prepare_message_for_openai(self, message: Message) -> list[dict[str, Any]]:
"""Prepare a chat message for OpenAI."""
all_messages: list[dict[str, Any]] = []
for content in message.contents:
python/packages/core/agent_framework/openai/_responses_client.py
+7 -7
View File
@@ -51,12 +51,12 @@ from .._tools import (
)
from .._types import (
Annotation,
ChatMessage,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
Content,
ContinuationToken,
Message,
ResponseStream,
Role,
TextSpanRegion,
@@ -250,7 +250,7 @@ class RawOpenAIResponsesClient( # type: ignore[misc]
async def _prepare_request(
self,
messages: Sequence[ChatMessage],
messages: Sequence[Message],
options: Mapping[str, Any],
**kwargs: Any,
) -> tuple[AsyncOpenAI, dict[str, Any], dict[str, Any]]:
@@ -280,7 +280,7 @@ class RawOpenAIResponsesClient( # type: ignore[misc]
def _inner_get_response(
self,
*,
messages: Sequence[ChatMessage],
messages: Sequence[Message],
options: Mapping[str, Any],
stream: bool = False,
**kwargs: Any,
@@ -567,7 +567,7 @@ class RawOpenAIResponsesClient( # type: ignore[misc]
async def _prepare_options(
self,
messages: Sequence[ChatMessage],
messages: Sequence[Message],
options: Mapping[str, Any],
**kwargs: Any,
) -> dict[str, Any]:
@@ -673,7 +673,7 @@ class RawOpenAIResponsesClient( # type: ignore[misc]
"""
return kwargs.get("conversation_id") or options.get("conversation_id")
def _prepare_messages_for_openai(self, chat_messages: Sequence[ChatMessage]) -> list[dict[str, Any]]:
def _prepare_messages_for_openai(self, chat_messages: Sequence[Message]) -> list[dict[str, Any]]:
"""Prepare the chat messages for a request.
Allowing customization of the key names for role/author, and optionally overriding the role.
@@ -705,7 +705,7 @@ class RawOpenAIResponsesClient( # type: ignore[misc]
def _prepare_message_for_openai(
self,
message: ChatMessage,
message: Message,
call_id_to_id: dict[str, str],
) -> list[dict[str, Any]]:
"""Prepare a chat message for the OpenAI Responses API format."""
@@ -1095,7 +1095,7 @@ class RawOpenAIResponsesClient( # type: ignore[misc]
)
case _:
logger.debug("Unparsed output of type: %s: %s", item.type, item)
response_message = ChatMessage(role="assistant", contents=contents)
response_message = Message(role="assistant", contents=contents)
args: dict[str, Any] = {
"response_id": response.id,
"created_at": datetime.fromtimestamp(response.created_at, tz=timezone.utc).strftime(
python/packages/core/tests/azure/conftest.py
+2 -2
View File
@@ -3,7 +3,7 @@ from typing import Any
from pytest import fixture
from agent_framework import ChatMessage
from agent_framework import Message
# region: Connector Settings fixtures
@@ -58,5 +58,5 @@ def azure_openai_unit_test_env(monkeypatch, exclude_list, override_env_param_dic
@fixture(scope="function")
def chat_history() -> list[ChatMessage]:
def chat_history() -> list[Message]:
return []
python/packages/core/tests/azure/test_azure_assistants_client.py
+73 -73
View File
@@ -9,15 +9,15 @@ from azure.identity import AzureCliCredential
from pydantic import Field
from agent_framework import (
Agent,
AgentResponse,
AgentResponseUpdate,
AgentThread,
ChatAgent,
ChatClientProtocol,
ChatMessage,
ChatResponse,
ChatResponseUpdate,
HostedCodeInterpreterTool,
Message,
SupportsChatGetResponse,
tool,
)
from agent_framework.azure import AzureOpenAIAssistantsClient
@@ -83,19 +83,19 @@ def mock_async_azure_openai() -> MagicMock:
def test_azure_assistants_client_init_with_client(mock_async_azure_openai: MagicMock) -> None:
"""Test AzureOpenAIAssistantsClient initialization with existing client."""
chat_client = create_test_azure_assistants_client(
client = create_test_azure_assistants_client(
mock_async_azure_openai,
deployment_name="test_chat_deployment",
assistant_id="existing-assistant-id",
thread_id="test-thread-id",
)
assert chat_client.client is mock_async_azure_openai
assert chat_client.model_id == "test_chat_deployment"
assert chat_client.assistant_id == "existing-assistant-id"
assert chat_client.thread_id == "test-thread-id"
assert not chat_client._should_delete_assistant # type: ignore
assert isinstance(chat_client, ChatClientProtocol)
assert client.client is mock_async_azure_openai
assert client.model_id == "test_chat_deployment"
assert client.assistant_id == "existing-assistant-id"
assert client.thread_id == "test-thread-id"
assert not client._should_delete_assistant # type: ignore
assert isinstance(client, SupportsChatGetResponse)
def test_azure_assistants_client_init_auto_create_client(
@@ -103,7 +103,7 @@ def test_azure_assistants_client_init_auto_create_client(
mock_async_azure_openai: MagicMock,
) -> None:
"""Test AzureOpenAIAssistantsClient initialization with auto-created client."""
chat_client = AzureOpenAIAssistantsClient(
client = AzureOpenAIAssistantsClient(
deployment_name=azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
assistant_name="TestAssistant",
api_key=azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"],
@@ -111,11 +111,11 @@ def test_azure_assistants_client_init_auto_create_client(
async_client=mock_async_azure_openai,
)
assert chat_client.client is mock_async_azure_openai
assert chat_client.model_id == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]
assert chat_client.assistant_id is None
assert chat_client.assistant_name == "TestAssistant"
assert not chat_client._should_delete_assistant # type: ignore
assert client.client is mock_async_azure_openai
assert client.model_id == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]
assert client.assistant_id is None
assert client.assistant_name == "TestAssistant"
assert not client._should_delete_assistant # type: ignore
def test_azure_assistants_client_init_validation_fail() -> None:
@@ -138,32 +138,32 @@ def test_azure_assistants_client_init_with_default_headers(azure_openai_unit_tes
"""Test AzureOpenAIAssistantsClient initialization with default headers."""
default_headers = {"X-Unit-Test": "test-guid"}
chat_client = AzureOpenAIAssistantsClient(
client = AzureOpenAIAssistantsClient(
deployment_name="test_chat_deployment",
api_key=azure_openai_unit_test_env["AZURE_OPENAI_API_KEY"],
endpoint=azure_openai_unit_test_env["AZURE_OPENAI_ENDPOINT"],
default_headers=default_headers,
)
assert chat_client.model_id == "test_chat_deployment"
assert isinstance(chat_client, ChatClientProtocol)
assert client.model_id == "test_chat_deployment"
assert isinstance(client, SupportsChatGetResponse)
# Assert that the default header we added is present in the client's default headers
for key, value in default_headers.items():
assert key in chat_client.client.default_headers
assert chat_client.client.default_headers[key] == value
assert key in client.client.default_headers
assert client.client.default_headers[key] == value
async def test_azure_assistants_client_get_assistant_id_or_create_existing_assistant(
mock_async_azure_openai: MagicMock,
) -> None:
"""Test _get_assistant_id_or_create when assistant_id is already provided."""
chat_client = create_test_azure_assistants_client(mock_async_azure_openai, assistant_id="existing-assistant-id")
client = create_test_azure_assistants_client(mock_async_azure_openai, assistant_id="existing-assistant-id")
assistant_id = await chat_client._get_assistant_id_or_create() # type: ignore
assistant_id = await client._get_assistant_id_or_create() # type: ignore
assert assistant_id == "existing-assistant-id"
assert not chat_client._should_delete_assistant # type: ignore
assert not client._should_delete_assistant # type: ignore
mock_async_azure_openai.beta.assistants.create.assert_not_called()
@@ -171,14 +171,14 @@ async def test_azure_assistants_client_get_assistant_id_or_create_create_new(
mock_async_azure_openai: MagicMock,
) -> None:
"""Test _get_assistant_id_or_create when creating a new assistant."""
chat_client = create_test_azure_assistants_client(
client = create_test_azure_assistants_client(
mock_async_azure_openai, deployment_name="test_chat_deployment", assistant_name="TestAssistant"
)
assistant_id = await chat_client._get_assistant_id_or_create() # type: ignore
assistant_id = await client._get_assistant_id_or_create() # type: ignore
assert assistant_id == "test-assistant-id"
assert chat_client._should_delete_assistant # type: ignore
assert client._should_delete_assistant # type: ignore
mock_async_azure_openai.beta.assistants.create.assert_called_once()
@@ -186,38 +186,38 @@ async def test_azure_assistants_client_aclose_should_not_delete(
mock_async_azure_openai: MagicMock,
) -> None:
"""Test close when assistant should not be deleted."""
chat_client = create_test_azure_assistants_client(
client = create_test_azure_assistants_client(
mock_async_azure_openai, assistant_id="assistant-to-keep", should_delete_assistant=False
)
await chat_client.close() # type: ignore
await client.close() # type: ignore
# Verify assistant deletion was not called
mock_async_azure_openai.beta.assistants.delete.assert_not_called()
assert not chat_client._should_delete_assistant # type: ignore
assert not client._should_delete_assistant # type: ignore
async def test_azure_assistants_client_aclose_should_delete(mock_async_azure_openai: MagicMock) -> None:
"""Test close method calls cleanup."""
chat_client = create_test_azure_assistants_client(
client = create_test_azure_assistants_client(
mock_async_azure_openai, assistant_id="assistant-to-delete", should_delete_assistant=True
)
await chat_client.close()
await client.close()
# Verify assistant deletion was called
mock_async_azure_openai.beta.assistants.delete.assert_called_once_with("assistant-to-delete")
assert not chat_client._should_delete_assistant # type: ignore
assert not client._should_delete_assistant # type: ignore
async def test_azure_assistants_client_async_context_manager(mock_async_azure_openai: MagicMock) -> None:
"""Test async context manager functionality."""
chat_client = create_test_azure_assistants_client(
client = create_test_azure_assistants_client(
mock_async_azure_openai, assistant_id="assistant-to-delete", should_delete_assistant=True
)
# Test context manager
async with chat_client:
async with client:
pass # Just test that we can enter and exit
# Verify cleanup was called on exit
@@ -229,7 +229,7 @@ def test_azure_assistants_client_serialize(azure_openai_unit_test_env: dict[str,
default_headers = {"X-Unit-Test": "test-guid"}
# Test basic initialization and to_dict
chat_client = AzureOpenAIAssistantsClient(
client = AzureOpenAIAssistantsClient(
deployment_name="test_chat_deployment",
assistant_id="test-assistant-id",
assistant_name="TestAssistant",
@@ -239,7 +239,7 @@ def test_azure_assistants_client_serialize(azure_openai_unit_test_env: dict[str,
default_headers=default_headers,
)
dumped_settings = chat_client.to_dict()
dumped_settings = client.to_dict()
assert dumped_settings["model_id"] == "test_chat_deployment"
assert dumped_settings["assistant_id"] == "test-assistant-id"
@@ -267,17 +267,17 @@ def get_weather(
async def test_azure_assistants_client_get_response() -> None:
"""Test Azure Assistants Client response."""
async with AzureOpenAIAssistantsClient(credential=AzureCliCredential()) as azure_assistants_client:
assert isinstance(azure_assistants_client, ChatClientProtocol)
assert isinstance(azure_assistants_client, SupportsChatGetResponse)
messages: list[ChatMessage] = []
messages: list[Message] = []
messages.append(
ChatMessage(
Message(
role="user",
text="The weather in Seattle is currently sunny with a high of 25°C. "
"It's a beautiful day for outdoor activities.",
)
)
messages.append(ChatMessage(role="user", text="What's the weather like today?"))
messages.append(Message(role="user", text="What's the weather like today?"))
# Test that the client can be used to get a response
response = await azure_assistants_client.get_response(messages=messages)
@@ -292,10 +292,10 @@ async def test_azure_assistants_client_get_response() -> None:
async def test_azure_assistants_client_get_response_tools() -> None:
"""Test Azure Assistants Client response with tools."""
async with AzureOpenAIAssistantsClient(credential=AzureCliCredential()) as azure_assistants_client:
assert isinstance(azure_assistants_client, ChatClientProtocol)
assert isinstance(azure_assistants_client, SupportsChatGetResponse)
messages: list[ChatMessage] = []
messages.append(ChatMessage(role="user", text="What's the weather like in Seattle?"))
messages: list[Message] = []
messages.append(Message(role="user", text="What's the weather like in Seattle?"))
# Test that the client can be used to get a response
response = await azure_assistants_client.get_response(
@@ -313,17 +313,17 @@ async def test_azure_assistants_client_get_response_tools() -> None:
async def test_azure_assistants_client_streaming() -> None:
"""Test Azure Assistants Client streaming response."""
async with AzureOpenAIAssistantsClient(credential=AzureCliCredential()) as azure_assistants_client:
assert isinstance(azure_assistants_client, ChatClientProtocol)
assert isinstance(azure_assistants_client, SupportsChatGetResponse)
messages: list[ChatMessage] = []
messages: list[Message] = []
messages.append(
ChatMessage(
Message(
role="user",
text="The weather in Seattle is currently sunny with a high of 25°C. "
"It's a beautiful day for outdoor activities.",
)
)
messages.append(ChatMessage(role="user", text="What's the weather like today?"))
messages.append(Message(role="user", text="What's the weather like today?"))
# Test that the client can be used to get a response
response = azure_assistants_client.get_response(messages=messages, stream=True)
@@ -344,10 +344,10 @@ async def test_azure_assistants_client_streaming() -> None:
async def test_azure_assistants_client_streaming_tools() -> None:
"""Test Azure Assistants Client streaming response with tools."""
async with AzureOpenAIAssistantsClient(credential=AzureCliCredential()) as azure_assistants_client:
assert isinstance(azure_assistants_client, ChatClientProtocol)
assert isinstance(azure_assistants_client, SupportsChatGetResponse)
messages: list[ChatMessage] = []
messages.append(ChatMessage(role="user", text="What's the weather like in Seattle?"))
messages: list[Message] = []
messages.append(Message(role="user", text="What's the weather like in Seattle?"))
# Test that the client can be used to get a response
response = azure_assistants_client.get_response(
@@ -373,7 +373,7 @@ async def test_azure_assistants_client_with_existing_assistant() -> None:
# First create an assistant to use in the test
async with AzureOpenAIAssistantsClient(credential=AzureCliCredential()) as temp_client:
# Get the assistant ID by triggering assistant creation
messages = [ChatMessage(role="user", text="Hello")]
messages = [Message(role="user", text="Hello")]
await temp_client.get_response(messages=messages)
assistant_id = temp_client.assistant_id
@@ -381,10 +381,10 @@ async def test_azure_assistants_client_with_existing_assistant() -> None:
async with AzureOpenAIAssistantsClient(
assistant_id=assistant_id, credential=AzureCliCredential()
) as azure_assistants_client:
assert isinstance(azure_assistants_client, ChatClientProtocol)
assert isinstance(azure_assistants_client, SupportsChatGetResponse)
assert azure_assistants_client.assistant_id == assistant_id
messages = [ChatMessage(role="user", text="What can you do?")]
messages = [Message(role="user", text="What can you do?")]
# Test that the client can be used to get a response
response = await azure_assistants_client.get_response(messages=messages)
@@ -397,9 +397,9 @@ async def test_azure_assistants_client_with_existing_assistant() -> None:
@pytest.mark.flaky
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_agent_basic_run():
"""Test ChatAgent basic run functionality with AzureOpenAIAssistantsClient."""
async with ChatAgent(
chat_client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
"""Test Agent basic run functionality with AzureOpenAIAssistantsClient."""
async with Agent(
client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
) as agent:
# Run a simple query
response = await agent.run("Hello! Please respond with 'Hello World' exactly.")
@@ -414,9 +414,9 @@ async def test_azure_assistants_agent_basic_run():
@pytest.mark.flaky
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_agent_basic_run_streaming():
"""Test ChatAgent basic streaming functionality with AzureOpenAIAssistantsClient."""
async with ChatAgent(
chat_client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
"""Test Agent basic streaming functionality with AzureOpenAIAssistantsClient."""
async with Agent(
client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
) as agent:
# Run streaming query
full_message: str = ""
@@ -434,9 +434,9 @@ async def test_azure_assistants_agent_basic_run_streaming():
@pytest.mark.flaky
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_agent_thread_persistence():
"""Test ChatAgent thread persistence across runs with AzureOpenAIAssistantsClient."""
async with ChatAgent(
chat_client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
"""Test Agent thread persistence across runs with AzureOpenAIAssistantsClient."""
async with Agent(
client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant with good memory.",
) as agent:
# Create a new thread that will be reused
@@ -463,12 +463,12 @@ async def test_azure_assistants_agent_thread_persistence():
@pytest.mark.flaky
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_agent_existing_thread_id():
"""Test ChatAgent with existing thread ID to continue conversations across agent instances."""
"""Test Agent with existing thread ID to continue conversations across agent instances."""
# First, create a conversation and capture the thread ID
existing_thread_id = None
async with ChatAgent(
chat_client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
async with Agent(
client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
instructions="You are a helpful weather agent.",
tools=[get_weather],
) as agent:
@@ -487,8 +487,8 @@ async def test_azure_assistants_agent_existing_thread_id():
# Now continue with the same thread ID in a new agent instance
async with ChatAgent(
chat_client=AzureOpenAIAssistantsClient(thread_id=existing_thread_id, credential=AzureCliCredential()),
async with Agent(
client=AzureOpenAIAssistantsClient(thread_id=existing_thread_id, credential=AzureCliCredential()),
instructions="You are a helpful weather agent.",
tools=[get_weather],
) as agent:
@@ -508,10 +508,10 @@ async def test_azure_assistants_agent_existing_thread_id():
@pytest.mark.flaky
@skip_if_azure_integration_tests_disabled
async def test_azure_assistants_agent_code_interpreter():
"""Test ChatAgent with code interpreter through AzureOpenAIAssistantsClient."""
"""Test Agent with code interpreter through AzureOpenAIAssistantsClient."""
async with ChatAgent(
chat_client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
async with Agent(
client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant that can write and execute Python code.",
tools=[HostedCodeInterpreterTool()],
) as agent:
@@ -530,8 +530,8 @@ async def test_azure_assistants_agent_code_interpreter():
async def test_azure_assistants_client_agent_level_tool_persistence():
"""Test that agent-level tools persist across multiple runs with Azure Assistants Client."""
async with ChatAgent(
chat_client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
async with Agent(
client=AzureOpenAIAssistantsClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant that uses available tools.",
tools=[get_weather], # Agent-level tool
) as agent:
python/packages/core/tests/azure/test_azure_chat_client.py
+60 -60
View File
@@ -17,13 +17,13 @@ from openai.types.chat.chat_completion_chunk import ChoiceDelta as ChunkChoiceDe
from openai.types.chat.chat_completion_message import ChatCompletionMessage
from agent_framework import (
Agent,
AgentResponse,
AgentResponseUpdate,
ChatAgent,
ChatClientProtocol,
ChatMessage,
ChatResponse,
ChatResponseUpdate,
Message,
SupportsChatGetResponse,
tool,
)
from agent_framework._telemetry import USER_AGENT_KEY
@@ -52,7 +52,7 @@ def test_init(azure_openai_unit_test_env: dict[str, str]) -> None:
assert azure_chat_client.client is not None
assert isinstance(azure_chat_client.client, AsyncAzureOpenAI)
assert azure_chat_client.model_id == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]
assert isinstance(azure_chat_client, ChatClientProtocol)
assert isinstance(azure_chat_client, SupportsChatGetResponse)
def test_init_client(azure_openai_unit_test_env: dict[str, str]) -> None:
@@ -75,7 +75,7 @@ def test_init_base_url(azure_openai_unit_test_env: dict[str, str]) -> None:
assert azure_chat_client.client is not None
assert isinstance(azure_chat_client.client, AsyncAzureOpenAI)
assert azure_chat_client.model_id == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]
assert isinstance(azure_chat_client, ChatClientProtocol)
assert isinstance(azure_chat_client, SupportsChatGetResponse)
for key, value in default_headers.items():
assert key in azure_chat_client.client.default_headers
assert azure_chat_client.client.default_headers[key] == value
@@ -88,7 +88,7 @@ def test_init_endpoint(azure_openai_unit_test_env: dict[str, str]) -> None:
assert azure_chat_client.client is not None
assert isinstance(azure_chat_client.client, AsyncAzureOpenAI)
assert azure_chat_client.model_id == azure_openai_unit_test_env["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]
assert isinstance(azure_chat_client, ChatClientProtocol)
assert isinstance(azure_chat_client, SupportsChatGetResponse)
@pytest.mark.parametrize("exclude_list", [["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"]], indirect=True)
@@ -178,11 +178,11 @@ def mock_streaming_chat_completion_response() -> AsyncStream[ChatCompletionChunk
async def test_cmc(
mock_create: AsyncMock,
azure_openai_unit_test_env: dict[str, str],
chat_history: list[ChatMessage],
chat_history: list[Message],
mock_chat_completion_response: ChatCompletion,
) -> None:
mock_create.return_value = mock_chat_completion_response
chat_history.append(ChatMessage(text="hello world", role="user"))
chat_history.append(Message(text="hello world", role="user"))
azure_chat_client = AzureOpenAIChatClient()
await azure_chat_client.get_response(
@@ -199,12 +199,12 @@ async def test_cmc(
async def test_cmc_with_logit_bias(
mock_create: AsyncMock,
azure_openai_unit_test_env: dict[str, str],
chat_history: list[ChatMessage],
chat_history: list[Message],
mock_chat_completion_response: ChatCompletion,
) -> None:
mock_create.return_value = mock_chat_completion_response
prompt = "hello world"
chat_history.append(ChatMessage(text=prompt, role="user"))
chat_history.append(Message(text=prompt, role="user"))
token_bias: dict[str | int, float] = {"1": -100}
@@ -224,12 +224,12 @@ async def test_cmc_with_logit_bias(
async def test_cmc_with_stop(
mock_create: AsyncMock,
azure_openai_unit_test_env: dict[str, str],
chat_history: list[ChatMessage],
chat_history: list[Message],
mock_chat_completion_response: ChatCompletion,
) -> None:
mock_create.return_value = mock_chat_completion_response
prompt = "hello world"
chat_history.append(ChatMessage(text=prompt, role="user"))
chat_history.append(Message(text=prompt, role="user"))
stop = ["!"]
@@ -249,7 +249,7 @@ async def test_cmc_with_stop(
async def test_azure_on_your_data(
mock_create: AsyncMock,
azure_openai_unit_test_env: dict[str, str],
chat_history: list[ChatMessage],
chat_history: list[Message],
mock_chat_completion_response: ChatCompletion,
) -> None:
mock_chat_completion_response.choices = [
@@ -277,9 +277,9 @@ async def test_azure_on_your_data(
mock_create.return_value = mock_chat_completion_response
prompt = "hello world"
messages_in = chat_history
chat_history.append(ChatMessage(text=prompt, role="user"))
messages_out: list[ChatMessage] = []
messages_out.append(ChatMessage(text=prompt, role="user"))
chat_history.append(Message(text=prompt, role="user"))
messages_out: list[Message] = []
messages_out.append(Message(text=prompt, role="user"))
expected_data_settings = {
"data_sources": [
@@ -319,7 +319,7 @@ async def test_azure_on_your_data(
async def test_azure_on_your_data_string(
mock_create: AsyncMock,
azure_openai_unit_test_env: dict[str, str],
chat_history: list[ChatMessage],
chat_history: list[Message],
mock_chat_completion_response: ChatCompletion,
) -> None:
mock_chat_completion_response.choices = [
@@ -347,9 +347,9 @@ async def test_azure_on_your_data_string(
mock_create.return_value = mock_chat_completion_response
prompt = "hello world"
messages_in = chat_history
messages_in.append(ChatMessage(text=prompt, role="user"))
messages_out: list[ChatMessage] = []
messages_out.append(ChatMessage(text=prompt, role="user"))
messages_in.append(Message(text=prompt, role="user"))
messages_out: list[Message] = []
messages_out.append(Message(text=prompt, role="user"))
expected_data_settings = {
"data_sources": [
@@ -389,7 +389,7 @@ async def test_azure_on_your_data_string(
async def test_azure_on_your_data_fail(
mock_create: AsyncMock,
azure_openai_unit_test_env: dict[str, str],
chat_history: list[ChatMessage],
chat_history: list[Message],
mock_chat_completion_response: ChatCompletion,
) -> None:
mock_chat_completion_response.choices = [
@@ -406,9 +406,9 @@ async def test_azure_on_your_data_fail(
mock_create.return_value = mock_chat_completion_response
prompt = "hello world"
messages_in = chat_history
messages_in.append(ChatMessage(text=prompt, role="user"))
messages_out: list[ChatMessage] = []
messages_out.append(ChatMessage(text=prompt, role="user"))
messages_in.append(Message(text=prompt, role="user"))
messages_out: list[Message] = []
messages_out.append(Message(text=prompt, role="user"))
expected_data_settings = {
"data_sources": [
@@ -459,10 +459,10 @@ CONTENT_FILTERED_ERROR_FULL_MESSAGE = (
async def test_content_filtering_raises_correct_exception(
mock_create: AsyncMock,
azure_openai_unit_test_env: dict[str, str],
chat_history: list[ChatMessage],
chat_history: list[Message],
) -> None:
prompt = "some prompt that would trigger the content filtering"
chat_history.append(ChatMessage(text=prompt, role="user"))
chat_history.append(Message(text=prompt, role="user"))
test_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
assert test_endpoint is not None
@@ -504,10 +504,10 @@ async def test_content_filtering_raises_correct_exception(
async def test_content_filtering_without_response_code_raises_with_default_code(
mock_create: AsyncMock,
azure_openai_unit_test_env: dict[str, str],
chat_history: list[ChatMessage],
chat_history: list[Message],
) -> None:
prompt = "some prompt that would trigger the content filtering"
chat_history.append(ChatMessage(text=prompt, role="user"))
chat_history.append(Message(text=prompt, role="user"))
test_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
assert test_endpoint is not None
@@ -543,10 +543,10 @@ async def test_content_filtering_without_response_code_raises_with_default_code(
async def test_bad_request_non_content_filter(
mock_create: AsyncMock,
azure_openai_unit_test_env: dict[str, str],
chat_history: list[ChatMessage],
chat_history: list[Message],
) -> None:
prompt = "some prompt that would trigger the content filtering"
chat_history.append(ChatMessage(text=prompt, role="user"))
chat_history.append(Message(text=prompt, role="user"))
test_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
assert test_endpoint is not None
@@ -566,11 +566,11 @@ async def test_bad_request_non_content_filter(
async def test_get_streaming(
mock_create: AsyncMock,
azure_openai_unit_test_env: dict[str, str],
chat_history: list[ChatMessage],
chat_history: list[Message],
mock_streaming_chat_completion_response: AsyncStream[ChatCompletionChunk],
) -> None:
mock_create.return_value = mock_streaming_chat_completion_response
chat_history.append(ChatMessage(text="hello world", role="user"))
chat_history.append(Message(text="hello world", role="user"))
azure_chat_client = AzureOpenAIChatClient()
async for msg in azure_chat_client.get_response(
@@ -595,7 +595,7 @@ async def test_get_streaming(
async def test_streaming_with_none_delta(
mock_create: AsyncMock,
azure_openai_unit_test_env: dict[str, str],
chat_history: list[ChatMessage],
chat_history: list[Message],
) -> None:
"""Test streaming handles None delta from async content filtering."""
# First chunk has None delta (simulates async filtering)
@@ -619,7 +619,7 @@ async def test_streaming_with_none_delta(
stream.__aiter__.return_value = [chunk_with_none_delta, chunk_with_content]
mock_create.return_value = stream
chat_history.append(ChatMessage(text="hello world", role="user"))
chat_history.append(Message(text="hello world", role="user"))
azure_chat_client = AzureOpenAIChatClient()
results: list[ChatResponseUpdate] = []
@@ -653,11 +653,11 @@ def get_weather(location: str) -> str:
async def test_azure_openai_chat_client_response() -> None:
"""Test Azure OpenAI chat completion responses."""
azure_chat_client = AzureOpenAIChatClient(credential=AzureCliCredential())
assert isinstance(azure_chat_client, ChatClientProtocol)
assert isinstance(azure_chat_client, SupportsChatGetResponse)
messages: list[ChatMessage] = []
messages: list[Message] = []
messages.append(
ChatMessage(
Message(
role="user",
text="Emily and David, two passionate scientists, met during a research expedition to Antarctica. "
"Bonded by their love for the natural world and shared curiosity, they uncovered a "
@@ -665,7 +665,7 @@ async def test_azure_openai_chat_client_response() -> None:
"of climate change.",
)
)
messages.append(ChatMessage(role="user", text="who are Emily and David?"))
messages.append(Message(role="user", text="who are Emily and David?"))
# Test that the client can be used to get a response
response = await azure_chat_client.get_response(messages=messages)
@@ -683,10 +683,10 @@ async def test_azure_openai_chat_client_response() -> None:
async def test_azure_openai_chat_client_response_tools() -> None:
"""Test AzureOpenAI chat completion responses."""
azure_chat_client = AzureOpenAIChatClient(credential=AzureCliCredential())
assert isinstance(azure_chat_client, ChatClientProtocol)
assert isinstance(azure_chat_client, SupportsChatGetResponse)
messages: list[ChatMessage] = []
messages.append(ChatMessage(role="user", text="who are Emily and David?"))
messages: list[Message] = []
messages.append(Message(role="user", text="who are Emily and David?"))
# Test that the client can be used to get a response
response = await azure_chat_client.get_response(
@@ -704,11 +704,11 @@ async def test_azure_openai_chat_client_response_tools() -> None:
async def test_azure_openai_chat_client_streaming() -> None:
"""Test Azure OpenAI chat completion responses."""
azure_chat_client = AzureOpenAIChatClient(credential=AzureCliCredential())
assert isinstance(azure_chat_client, ChatClientProtocol)
assert isinstance(azure_chat_client, SupportsChatGetResponse)
messages: list[ChatMessage] = []
messages: list[Message] = []
messages.append(
ChatMessage(
Message(
role="user",
text="Emily and David, two passionate scientists, met during a research expedition to Antarctica. "
"Bonded by their love for the natural world and shared curiosity, they uncovered a "
@@ -716,7 +716,7 @@ async def test_azure_openai_chat_client_streaming() -> None:
"of climate change.",
)
)
messages.append(ChatMessage(role="user", text="who are Emily and David?"))
messages.append(Message(role="user", text="who are Emily and David?"))
# Test that the client can be used to get a response
response = azure_chat_client.get_response(messages=messages, stream=True)
@@ -739,10 +739,10 @@ async def test_azure_openai_chat_client_streaming() -> None:
async def test_azure_openai_chat_client_streaming_tools() -> None:
"""Test AzureOpenAI chat completion responses."""
azure_chat_client = AzureOpenAIChatClient(credential=AzureCliCredential())
assert isinstance(azure_chat_client, ChatClientProtocol)
assert isinstance(azure_chat_client, SupportsChatGetResponse)
messages: list[ChatMessage] = []
messages.append(ChatMessage(role="user", text="who are Emily and David?"))
messages: list[Message] = []
messages.append(Message(role="user", text="who are Emily and David?"))
# Test that the client can be used to get a response
response = azure_chat_client.get_response(
@@ -765,8 +765,8 @@ async def test_azure_openai_chat_client_streaming_tools() -> None:
@skip_if_azure_integration_tests_disabled
async def test_azure_openai_chat_client_agent_basic_run():
"""Test Azure OpenAI chat client agent basic run functionality with AzureOpenAIChatClient."""
async with ChatAgent(
chat_client=AzureOpenAIChatClient(credential=AzureCliCredential()),
async with Agent(
client=AzureOpenAIChatClient(credential=AzureCliCredential()),
) as agent:
# Test basic run
response = await agent.run("Please respond with exactly: 'This is a response test.'")
@@ -781,8 +781,8 @@ async def test_azure_openai_chat_client_agent_basic_run():
@skip_if_azure_integration_tests_disabled
async def test_azure_openai_chat_client_agent_basic_run_streaming():
"""Test Azure OpenAI chat client agent basic streaming functionality with AzureOpenAIChatClient."""
async with ChatAgent(
chat_client=AzureOpenAIChatClient(credential=AzureCliCredential()),
async with Agent(
client=AzureOpenAIChatClient(credential=AzureCliCredential()),
) as agent:
# Test streaming run
full_text = ""
@@ -799,8 +799,8 @@ async def test_azure_openai_chat_client_agent_basic_run_streaming():
@skip_if_azure_integration_tests_disabled
async def test_azure_openai_chat_client_agent_thread_persistence():
"""Test Azure OpenAI chat client agent thread persistence across runs with AzureOpenAIChatClient."""
async with ChatAgent(
chat_client=AzureOpenAIChatClient(credential=AzureCliCredential()),
async with Agent(
client=AzureOpenAIChatClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant with good memory.",
) as agent:
# Create a new thread that will be reused
@@ -827,8 +827,8 @@ async def test_azure_openai_chat_client_agent_existing_thread():
# First conversation - capture the thread
preserved_thread = None
async with ChatAgent(
chat_client=AzureOpenAIChatClient(credential=AzureCliCredential()),
async with Agent(
client=AzureOpenAIChatClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant with good memory.",
) as first_agent:
# Start a conversation and capture the thread
@@ -843,8 +843,8 @@ async def test_azure_openai_chat_client_agent_existing_thread():
# Second conversation - reuse the thread in a new agent instance
if preserved_thread:
async with ChatAgent(
chat_client=AzureOpenAIChatClient(credential=AzureCliCredential()),
async with Agent(
client=AzureOpenAIChatClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant with good memory.",
) as second_agent:
# Reuse the preserved thread
@@ -860,8 +860,8 @@ async def test_azure_openai_chat_client_agent_existing_thread():
async def test_azure_chat_client_agent_level_tool_persistence():
"""Test that agent-level tools persist across multiple runs with Azure Chat Client."""
async with ChatAgent(
chat_client=AzureOpenAIChatClient(credential=AzureCliCredential()),
async with Agent(
client=AzureOpenAIChatClient(credential=AzureCliCredential()),
instructions="You are a helpful assistant that uses available tools.",
tools=[get_weather], # Agent-level tool
) as agent:

Some files were not shown because too many files have changed in this diff Show More