From df84675c0f4b91c6cee696426328097a05ce42fa Mon Sep 17 00:00:00 2001 From: Dmytro Struk <13853051+dmytrostruk@users.noreply.github.com> Date: Fri, 11 Jul 2025 08:13:06 -0700 Subject: [PATCH] Python: Added AgentRunResponse and AgentRunResponseUpdate types (#157) * Removed instructions property from Agent * Added AgentRunResponse and AgentRunResponseUpdate types * Added unit tests for agent response types * Small fix * Addressed PR feedback * Small improvement * Small fix --- .../packages/main/agent_framework/__init__.py | 2 + .../main/agent_framework/__init__.pyi | 4 + .../packages/main/agent_framework/_agents.py | 22 +-- .../packages/main/agent_framework/_types.py | 127 ++++++++++++++++-- .../packages/main/tests/unit/test_agents.py | 21 ++- python/packages/main/tests/unit/test_types.py | 97 ++++++++++++- 6 files changed, 241 insertions(+), 32 deletions(-) diff --git a/python/packages/main/agent_framework/__init__.py b/python/packages/main/agent_framework/__init__.py index 7bb938a5df..408fa83351 100644 --- a/python/packages/main/agent_framework/__init__.py +++ b/python/packages/main/agent_framework/__init__.py @@ -14,6 +14,8 @@ _IMPORTS = { "AFBaseModel": "._pydantic", "AFBaseSettings": "._pydantic", "Agent": "._agents", + "AgentRunResponse": "._types", + "AgentRunResponseUpdate": "._types", "AgentThread": "._agents", "AITool": "._tools", "ai_function": "._tools", diff --git a/python/packages/main/agent_framework/__init__.pyi b/python/packages/main/agent_framework/__init__.pyi index 28b799be37..834a7524e1 100644 --- a/python/packages/main/agent_framework/__init__.pyi +++ b/python/packages/main/agent_framework/__init__.pyi @@ -7,6 +7,8 @@ from ._logging import get_logger from ._pydantic import AFBaseModel, AFBaseSettings from ._tools import AITool, ai_function from ._types import ( + AgentRunResponse, + AgentRunResponseUpdate, AIContent, AIContents, ChatFinishReason, @@ -39,6 +41,8 @@ __all__ = [ "AIContents", "AITool", "Agent", + "AgentRunResponse", + "AgentRunResponseUpdate", "AgentThread", "ChatClient", "ChatClientBase", diff --git a/python/packages/main/agent_framework/_agents.py b/python/packages/main/agent_framework/_agents.py index d993479140..a74b1cb920 100644 --- a/python/packages/main/agent_framework/_agents.py +++ b/python/packages/main/agent_framework/_agents.py @@ -5,7 +5,7 @@ from collections.abc import AsyncIterable, Sequence from typing import Any, Protocol, runtime_checkable from ._pydantic import AFBaseModel -from ._types import ChatMessage, ChatResponse, ChatResponseUpdate +from ._types import AgentRunResponse, AgentRunResponseUpdate, ChatMessage # region AgentThread @@ -82,26 +82,21 @@ class Agent(Protocol): """Returns the description of the agent.""" ... - @property - def instructions(self) -> str | None: - """Returns the instructions for the agent.""" - ... - async def run( self, messages: str | ChatMessage | list[ChatMessage] | None = None, *, thread: AgentThread | None = None, **kwargs: Any, - ) -> ChatResponse: + ) -> AgentRunResponse: """Get a response from the agent. This method returns the final result of the agent's execution - as a single ChatResponse object. The caller is blocked until + as a single AgentRunResponse object. The caller is blocked until the final result is available. Note: For streaming responses, use the run_stream method, which returns - intermediate steps and the final result as a stream of ChatResponseUpdate + intermediate steps and the final result as a stream of AgentRunResponseUpdate objects. Streaming only the final result is not feasible because the timing of the final result's availability is unknown, and blocking the caller until then is undesirable in streaming scenarios. @@ -122,16 +117,13 @@ class Agent(Protocol): *, thread: AgentThread | None = None, **kwargs: Any, - ) -> AsyncIterable[ChatResponseUpdate]: + ) -> AsyncIterable[AgentRunResponseUpdate]: """Run the agent as a stream. This method will return the intermediate steps and final results of the - agent's execution as a stream of ChatResponseUpdate objects to the caller. + agent's execution as a stream of AgentRunResponseUpdate objects to the caller. - To get the intermediate steps of the agent's execution as fully formed messages, - use the on_intermediate_message callback. - - Note: A ChatResponseUpdate object contains a chunk of a message. + Note: An AgentRunResponseUpdate object contains a chunk of a message. Args: messages: The message(s) to send to the agent. diff --git a/python/packages/main/agent_framework/_types.py b/python/packages/main/agent_framework/_types.py index 65d76fde6a..f3f6a92d13 100644 --- a/python/packages/main/agent_framework/_types.py +++ b/python/packages/main/agent_framework/_types.py @@ -25,6 +25,7 @@ TValue = TypeVar("TValue") TEmbedding = TypeVar("TEmbedding") TChatResponse = TypeVar("TChatResponse", bound="ChatResponse") TChatToolMode = TypeVar("TChatToolMode", bound="ChatToolMode") +TAgentRunResponse = TypeVar("TAgentRunResponse", bound="AgentRunResponse") CreatedAtT = str # Use a datetimeoffset type? Or a more specific type like datetime.datetime? @@ -152,7 +153,9 @@ class UsageDetails(AFBaseModel): return self -def _process_update(response: "ChatResponse", update: "ChatResponseUpdate") -> None: +def _process_update( + response: "ChatResponse | AgentRunResponse", update: "ChatResponseUpdate | AgentRunResponseUpdate" +) -> None: """Processes a single update and modifies the response in place.""" is_new_message = False if not response.messages or (update.message_id and response.messages[-1].message_id != update.message_id): @@ -189,19 +192,21 @@ def _process_update(response: "ChatResponse", update: "ChatResponseUpdate") -> N # Incorporate the update's properties into the response. if update.response_id: response.response_id = update.response_id - if update.conversation_id is not None: - response.conversation_id = update.conversation_id if update.created_at is not None: response.created_at = update.created_at - if update.finish_reason is not None: - response.finish_reason = update.finish_reason - if update.ai_model_id is not None: - response.ai_model_id = update.ai_model_id if update.additional_properties is not None: if response.additional_properties is None: response.additional_properties = {} response.additional_properties.update(update.additional_properties) + if isinstance(response, ChatResponse) and isinstance(update, ChatResponseUpdate): + if update.conversation_id is not None: + response.conversation_id = update.conversation_id + if update.finish_reason is not None: + response.finish_reason = update.finish_reason + if update.ai_model_id is not None: + response.ai_model_id = update.ai_model_id + def _coalesce_text_content( contents: list["AIContents"], type_: type["TextContent"] | type["TextReasoningContent"] @@ -235,8 +240,8 @@ def _coalesce_text_content( contents.extend(coalesced_contents) -def _finalize_response(response: "ChatResponse") -> None: - """Finalizes the chat response by performing any necessary post-processing.""" +def _finalize_response(response: "ChatResponse | AgentRunResponse") -> None: + """Finalizes the response by performing any necessary post-processing.""" for msg in response.messages: _coalesce_text_content(msg.contents, TextContent) _coalesce_text_content(msg.contents, TextReasoningContent) @@ -1554,6 +1559,110 @@ class GeneratedEmbeddings(AFBaseModel, MutableSequence[TEmbedding], Generic[TEmb return self +# region AgentRunResponse + + +class AgentRunResponse(AFBaseModel): + """Represents the response to an Agent run request. + + Provides one or more response messages and metadata about the response. + A typical response will contain a single message, but may contain multiple + messages in scenarios involving function calls, RAG retrievals, or complex logic. + """ + + messages: list[ChatMessage] = Field(default_factory=list[ChatMessage]) + response_id: str | None = None + created_at: CreatedAtT | None = None # use a datetimeoffset type? + usage_details: UsageDetails | None = None + raw_representation: Any | None = None + additional_properties: dict[str, Any] | None = None + + def __init__( + self, + messages: ChatMessage | list[ChatMessage] | None = None, + response_id: str | None = None, + created_at: CreatedAtT | None = None, + usage_details: UsageDetails | None = None, + raw_representation: Any | None = None, + additional_properties: dict[str, Any] | None = None, + **kwargs: Any, + ) -> None: + """Initialize an AgentRunResponse. + + Attributes: + messages: The list of chat messages in the response. + response_id: The ID of the chat response. + created_at: A timestamp for the chat response. + usage_details: The usage details for the chat response. + additional_properties: Any additional properties associated with the chat response. + raw_representation: The raw representation of the chat response from an underlying implementation. + **kwargs: Additional properties to set on the response. + """ + processed_messages: list[ChatMessage] = [] + if messages is not None: + if isinstance(messages, ChatMessage): + processed_messages.append(messages) + elif isinstance(messages, list): + processed_messages.extend(messages) + + super().__init__( + messages=processed_messages, # type: ignore[reportCallIssue] + response_id=response_id, # type: ignore[reportCallIssue] + created_at=created_at, # type: ignore[reportCallIssue] + usage_details=usage_details, # type: ignore[reportCallIssue] + additional_properties=additional_properties, # type: ignore[reportCallIssue] + raw_representation=raw_representation, # type: ignore[reportCallIssue] + **kwargs, + ) + + @property + def text(self) -> str: + """Get the concatenated text of all messages.""" + return "".join(msg.text for msg in self.messages) if self.messages else "" + + @classmethod + def from_agent_run_response_updates( + cls: type[TAgentRunResponse], updates: Sequence["AgentRunResponseUpdate"] + ) -> TAgentRunResponse: + """Joins multiple updates into a single AgentRunResponse.""" + msg = cls(messages=[]) + for update in updates: + _process_update(msg, update) + _finalize_response(msg) + return msg + + def __str__(self) -> str: + return self.text + + +# region AgentRunResponseUpdate + + +class AgentRunResponseUpdate(AFBaseModel): + """Represents a single streaming response chunk from an Agent.""" + + contents: list[AIContents] = Field(default_factory=list[AIContents]) + role: ChatRole | None = None + author_name: str | None = None + response_id: str | None = None + message_id: str | None = None + created_at: CreatedAtT | None = None # use a datetimeoffset type? + additional_properties: dict[str, Any] | None = None + raw_representation: Any | None = None + + @property + def text(self) -> str: + """Get the concatenated text of all TextContent objects in contents.""" + return ( + "".join(content.text for content in self.contents if isinstance(content, TextContent)) + if self.contents + else "" + ) + + def __str__(self) -> str: + return self.text + + # region: SpeechToTextOptions diff --git a/python/packages/main/tests/unit/test_agents.py b/python/packages/main/tests/unit/test_agents.py index 6b6af9731b..ee4435e6bf 100644 --- a/python/packages/main/tests/unit/test_agents.py +++ b/python/packages/main/tests/unit/test_agents.py @@ -7,7 +7,15 @@ from uuid import uuid4 from pydantic import BaseModel, Field from pytest import fixture -from agent_framework import Agent, AgentThread, ChatMessage, ChatResponse, ChatResponseUpdate, ChatRole, TextContent +from agent_framework import ( + Agent, + AgentRunResponse, + AgentRunResponseUpdate, + AgentThread, + ChatMessage, + ChatRole, + TextContent, +) TThreadType = TypeVar("TThreadType", bound=AgentThread) @@ -29,7 +37,6 @@ class MockAgent(BaseModel): id: str = Field(default_factory=lambda: str(uuid4())) name: str | None = None description: str | None = None - instructions: str | None = None async def run( self, @@ -37,8 +44,8 @@ class MockAgent(BaseModel): *, thread: AgentThread | None = None, **kwargs: Any, - ) -> ChatResponse: - return ChatResponse(messages=[ChatMessage(role=ChatRole.ASSISTANT, contents=[TextContent("Response")])]) + ) -> AgentRunResponse: + return AgentRunResponse(messages=[ChatMessage(role=ChatRole.ASSISTANT, contents=[TextContent("Response")])]) async def run_stream( self, @@ -46,8 +53,8 @@ class MockAgent(BaseModel): *, thread: AgentThread | None = None, **kwargs: Any, - ) -> AsyncIterable[ChatResponseUpdate]: - yield ChatResponseUpdate(contents=[TextContent("Response")]) + ) -> AsyncIterable[AgentRunResponseUpdate]: + yield AgentRunResponseUpdate(contents=[TextContent("Response")]) def get_new_thread(self) -> AgentThread: return MockAgentThread() @@ -107,7 +114,7 @@ async def test_agent_run(agent: Agent) -> None: async def test_agent_run_stream(agent: Agent) -> None: - async def collect_updates(updates: AsyncIterable[ChatResponseUpdate]) -> list[ChatResponseUpdate]: + async def collect_updates(updates: AsyncIterable[AgentRunResponseUpdate]) -> list[AgentRunResponseUpdate]: return [u async for u in updates] updates = await collect_updates(agent.run_stream(messages="test")) diff --git a/python/packages/main/tests/unit/test_types.py b/python/packages/main/tests/unit/test_types.py index 79bc570c46..a16fb02df9 100644 --- a/python/packages/main/tests/unit/test_types.py +++ b/python/packages/main/tests/unit/test_types.py @@ -3,9 +3,11 @@ from collections.abc import MutableSequence from pydantic import BaseModel, ValidationError -from pytest import mark, raises +from pytest import fixture, mark, raises from agent_framework import ( + AgentRunResponse, + AgentRunResponseUpdate, AIContent, AIContents, ChatMessage, @@ -480,3 +482,96 @@ def test_generated_embeddings(): # Ensure the instance is of type GeneratedEmbeddings assert isinstance(embeddings, GeneratedEmbeddings) assert issubclass(GeneratedEmbeddings, MutableSequence) + + +# region Agent Response Fixtures + + +@fixture +def chat_message() -> ChatMessage: + return ChatMessage(role=ChatRole.USER, text="Hello") + + +@fixture +def text_content() -> TextContent: + return TextContent(text="Test content") + + +@fixture +def agent_run_response(chat_message: ChatMessage) -> AgentRunResponse: + return AgentRunResponse(messages=chat_message) + + +@fixture +def agent_run_response_update(text_content: TextContent) -> AgentRunResponseUpdate: + return AgentRunResponseUpdate(role=ChatRole.ASSISTANT, contents=[text_content]) + + +# region AgentRunResponse + + +def test_agent_run_response_init_single_message(chat_message: ChatMessage) -> None: + response = AgentRunResponse(messages=chat_message) + assert response.messages == [chat_message] + + +def test_agent_run_response_init_list_messages(chat_message: ChatMessage) -> None: + response = AgentRunResponse(messages=[chat_message, chat_message]) + assert len(response.messages) == 2 + assert response.messages[0] == chat_message + + +def test_agent_run_response_init_none_messages() -> None: + response = AgentRunResponse() + assert response.messages == [] + + +def test_agent_run_response_text_property(chat_message: ChatMessage) -> None: + response = AgentRunResponse(messages=[chat_message, chat_message]) + assert response.text == "HelloHello" + + +def test_agent_run_response_text_property_empty() -> None: + response = AgentRunResponse() + assert response.text == "" + + +def test_agent_run_response_from_updates(agent_run_response_update: AgentRunResponseUpdate) -> None: + updates = [agent_run_response_update, agent_run_response_update] + response = AgentRunResponse.from_agent_run_response_updates(updates) + assert len(response.messages) > 0 + assert response.text == "Test content\nTest content" + + +def test_agent_run_response_str_method(chat_message: ChatMessage) -> None: + response = AgentRunResponse(messages=chat_message) + assert str(response) == "Hello" + + +# region AgentRunResponseUpdate + + +def test_agent_run_response_update_init_content_list(text_content: TextContent) -> None: + update = AgentRunResponseUpdate(contents=[text_content, text_content]) + assert len(update.contents) == 2 + assert update.contents[0] == text_content + + +def test_agent_run_response_update_init_none_content() -> None: + update = AgentRunResponseUpdate() + assert update.contents == [] + + +def test_agent_run_response_update_text_property(text_content: TextContent) -> None: + update = AgentRunResponseUpdate(contents=[text_content, text_content]) + assert update.text == "Test contentTest content" + + +def test_agent_run_response_update_text_property_empty() -> None: + update = AgentRunResponseUpdate() + assert update.text == "" + + +def test_agent_run_response_update_str_method(text_content: TextContent) -> None: + update = AgentRunResponseUpdate(contents=[text_content]) + assert str(update) == "Test content"