diff --git a/python/packages/azure/pyproject.toml b/python/packages/azure/pyproject.toml index c2f2a718f1..7af20f799a 100644 --- a/python/packages/azure/pyproject.toml +++ b/python/packages/azure/pyproject.toml @@ -43,7 +43,7 @@ addopts = "-ra -q -r fEX" asyncio_mode = "auto" asyncio_default_fixture_loop_scope = "function" filterwarnings = [] -timeout = 120 +timeout = 300 [tool.ruff] extend = "../../pyproject.toml" diff --git a/python/packages/foundry/agent_framework_foundry/_chat_client.py b/python/packages/foundry/agent_framework_foundry/_chat_client.py index aee3e33feb..9c7739ff86 100644 --- a/python/packages/foundry/agent_framework_foundry/_chat_client.py +++ b/python/packages/foundry/agent_framework_foundry/_chat_client.py @@ -27,7 +27,7 @@ from agent_framework import ( use_tool_calling, ) from agent_framework._pydantic import AFBaseSettings -from agent_framework.exceptions import ServiceInitializationError +from agent_framework.exceptions import ServiceInitializationError, ServiceResponseException from agent_framework.telemetry import use_telemetry from azure.ai.agents.models import ( AgentsNamedToolChoice, @@ -48,6 +48,7 @@ from azure.ai.agents.models import ( RequiredFunctionToolCall, ResponseFormatJsonSchema, ResponseFormatJsonSchemaType, + RunError, RunStatus, RunStep, SubmitToolOutputsAction, @@ -427,6 +428,12 @@ class FoundryChatClient(ChatClientBase): raw_representation=event_data, response_id=response_id, ) + elif ( + event_type == AgentStreamEvent.THREAD_RUN_FAILED + and isinstance(event_data, ThreadRun) + and isinstance(event_data.last_error, RunError) + ): + raise ServiceResponseException(event_data.last_error.message) else: yield ChatResponseUpdate( contents=[], diff --git a/python/packages/main/agent_framework/__init__.py b/python/packages/main/agent_framework/__init__.py index 7a9a3d027f..da8d8ee482 100644 --- a/python/packages/main/agent_framework/__init__.py +++ b/python/packages/main/agent_framework/__init__.py @@ -12,5 +12,6 @@ from ._agents import * # noqa: F403 from ._clients import * # noqa: F403 from ._logging import * # noqa: F403 from ._mcp import * # noqa: F403 +from ._threads import * # noqa: F403 from ._tools import * # noqa: F403 from ._types import * # noqa: F403 diff --git a/python/packages/main/agent_framework/_agents.py b/python/packages/main/agent_framework/_agents.py index 948b6e5d59..f15a42d326 100644 --- a/python/packages/main/agent_framework/_agents.py +++ b/python/packages/main/agent_framework/_agents.py @@ -3,7 +3,6 @@ import sys from collections.abc import AsyncIterable, Callable, MutableMapping, Sequence from contextlib import AbstractAsyncContextManager, AsyncExitStack -from enum import Enum from itertools import chain from typing import Any, ClassVar, Literal, Protocol, TypeVar, runtime_checkable from uuid import uuid4 @@ -13,6 +12,7 @@ from pydantic import BaseModel, Field, PrivateAttr from ._clients import ChatClient from ._mcp import McpTool from ._pydantic import AFBaseModel +from ._threads import AgentThread, ChatMessageStore, deserialize_thread_state, thread_on_new_messages from ._tools import AITool from ._types import ( AgentRunResponse, @@ -34,47 +34,7 @@ else: TThreadType = TypeVar("TThreadType", bound="AgentThread") -# region AgentThread - -__all__ = [ - "AIAgent", - "AgentBase", - "AgentThread", - "ChatClientAgent", - "ChatClientAgentThread", - "ChatClientAgentThreadType", - "MessagesRetrievableThread", -] - - -class AgentThread(AFBaseModel): - """Base class for agent threads.""" - - id: str | None = None - - async def on_new_messages( - self, - new_messages: ChatMessage | Sequence[ChatMessage], - ) -> None: - """Invoked when a new message has been contributed to the chat by any participant.""" - await self._on_new_messages(new_messages=new_messages) - - async def _on_new_messages( - self, - new_messages: ChatMessage | Sequence[ChatMessage], - ) -> None: - """Invoked when a new message has been contributed to the chat by any participant.""" - pass - - -# region MessagesRetrievableThread - - -@runtime_checkable -class MessagesRetrievableThread(Protocol): - def get_messages(self) -> AsyncIterable[ChatMessage]: - """Asynchronously retrieves all messages from thread.""" - ... +__all__ = ["AIAgent", "AgentBase", "ChatClientAgent"] # region Agent Protocol @@ -186,7 +146,7 @@ class AgentBase(AFBaseModel): ) -> None: """Notify the thread of new messages.""" if isinstance(new_messages, ChatMessage) or len(new_messages) > 0: - await thread.on_new_messages(new_messages) + await thread_on_new_messages(thread, new_messages) @property def display_name(self) -> str: @@ -196,116 +156,17 @@ class AgentBase(AFBaseModel): """ return self.name or self.id - def _validate_or_create_thread_type( - self, - thread: AgentThread | None, - construct_thread: Callable[[], TThreadType], - expected_type: type[TThreadType], - ) -> TThreadType: - """Validate or create a AgentThread of the right type. - - Args: - thread: The thread to validate or create. - construct_thread: A callable that constructs a new thread if `thread` is None. - expected_type: The expected type of the thread. - - Returns: - The validated or newly created thread of the expected type. - - Raises: - AgentExecutionException: If the thread is not of the expected type. - """ - if thread is None: - return construct_thread() - - if not isinstance(thread, expected_type): - raise AgentExecutionException( - f"{self.__class__.__name__} currently only supports agent threads of type {expected_type.__name__}." - ) + def get_new_thread(self) -> AgentThread: + """Returns AgentThread instance that is compatible with the agent.""" + return AgentThread() + async def deserialize_thread(self, serialized_thread: Any, **kwargs: Any) -> AgentThread: + """Deserializes the thread.""" + thread: AgentThread = self.get_new_thread() + await deserialize_thread_state(thread, serialized_thread, **kwargs) return thread -# region ChatClientAgentThread - - -class ChatClientAgentThreadType(Enum): - """Defines the different supported storage locations for ChatClientAgentThread.""" - - IN_MEMORY_MESSAGES = "InMemoryMessages" - """Messages are stored in memory inside the thread object.""" - - CONVERSATION_ID = "ConversationId" - """Messages are stored in the service and the thread object just has an id reference to the service storage.""" - - -class ChatClientAgentThread(AgentThread): - """Chat client agent thread. - - This class manages chat threads either locally (in-memory) or via a service based on initialization. - """ - - chat_messages: list[ChatMessage] | None = None - storage_location: ChatClientAgentThreadType | None = None - - def __init__( - self, - id: str | None = None, - messages: Sequence[ChatMessage] | None = None, - **kwargs: Any, - ): - """Initialize the chat client agent thread. - - Args: - id: Service thread identifier. If provided, thread is managed by the service and messages are - not stored locally. Must not be empty or whitespace. - messages: Initial messages for local storage. If provided, thread is managed - locally in-memory. - kwargs: Additional keyword arguments. - - Raises: - ValueError: If both id and messages are provided, or if id is empty/whitespace. - - Notes: - - If id is set, _id is assigned and _chat_messages is None (service-managed). - - If messages is set, _chat_messages is populated and _id is None (local). - - If neither is provided, creates an empty local thread. - """ - processed_messages: list[ChatMessage] | None = None - storage_location: ChatClientAgentThreadType | None = None - - if id and messages: - raise ValueError("Cannot specify both id and messages") - - if id: - if not id.strip(): - raise ValueError("ID cannot be empty or whitespace") - storage_location = ChatClientAgentThreadType.CONVERSATION_ID - elif messages: - processed_messages = [] - processed_messages.extend(messages) - storage_location = ChatClientAgentThreadType.IN_MEMORY_MESSAGES - - super().__init__( - id=id, - chat_messages=processed_messages, # type: ignore[reportCallIssue] - storage_location=storage_location, # type: ignore[reportCallIssue] - **kwargs, - ) - - async def get_messages(self) -> AsyncIterable[ChatMessage]: - """Get all messages in the thread.""" - for message in self.chat_messages or []: - yield message - - async def _on_new_messages(self, new_messages: ChatMessage | Sequence[ChatMessage]) -> None: - """Handle new messages.""" - if self.storage_location == ChatClientAgentThreadType.IN_MEMORY_MESSAGES: - if self.chat_messages is None: - self.chat_messages = [] - self.chat_messages.extend([new_messages] if isinstance(new_messages, ChatMessage) else new_messages) - - # region ChatClientAgent @@ -317,6 +178,7 @@ class ChatClientAgent(AgentBase): chat_client: ChatClient instructions: str | None = None chat_options: ChatOptions + chat_message_store_factory: Callable[[], ChatMessageStore] | None = None _local_mcp_tools: list[McpTool] = PrivateAttr(default_factory=list) # type: ignore[reportUnknownVariableType] _async_exit_stack: AsyncExitStack = PrivateAttr(default_factory=AsyncExitStack) @@ -350,6 +212,7 @@ class ChatClientAgent(AgentBase): top_p: float | None = None, user: str | None = None, additional_properties: dict[str, Any] | None = None, + chat_message_store_factory: Callable[[], ChatMessageStore] | None = None, **kwargs: Any, ) -> None: """Create a ChatClientAgent. @@ -382,6 +245,8 @@ class ChatClientAgent(AgentBase): top_p: the nucleus sampling probability to use. user: the user to associate with the request. additional_properties: additional properties to include in the request. + chat_message_store_factory: factory function to create an instance of ChatMessageStore. If not provided, + the default in-memory store will be used. kwargs: any additional keyword arguments. Unused, can be used by subclasses of this Agent. """ @@ -394,6 +259,7 @@ class ChatClientAgent(AgentBase): final_tools = [tool for tool in normalized_tools if not isinstance(tool, McpTool)] args: dict[str, Any] = { "chat_client": chat_client, + "chat_message_store_factory": chat_message_store_factory, "chat_options": ChatOptions( ai_model_id=model, frequency_penalty=frequency_penalty, @@ -537,7 +403,7 @@ class ChatClientAgent(AgentBase): chat_options=self.chat_options & ChatOptions( ai_model_id=model, - conversation_id=thread.id, + conversation_id=thread.service_thread_id, frequency_penalty=frequency_penalty, logit_bias=logit_bias, max_tokens=max_tokens, @@ -660,7 +526,7 @@ class ChatClientAgent(AgentBase): messages=thread_messages, chat_options=self.chat_options & ChatOptions( - conversation_id=thread.id, + conversation_id=thread.service_thread_id, frequency_penalty=frequency_penalty, logit_bias=logit_bias, max_tokens=max_tokens, @@ -705,44 +571,50 @@ class ChatClientAgent(AgentBase): await self._notify_thread_of_new_messages(thread, input_messages) await self._notify_thread_of_new_messages(thread, response.messages) - def get_new_thread(self) -> ChatClientAgentThread: - return ChatClientAgentThread() + def get_new_thread(self) -> AgentThread: + message_store: ChatMessageStore | None = None + + if self.chat_message_store_factory: + message_store = self.chat_message_store_factory() + + return AgentThread() if message_store is None else AgentThread(message_store=message_store) def _update_thread_with_type_and_conversation_id( - self, chat_client_thread: ChatClientAgentThread, responseConversationId: str | None + self, thread: AgentThread, response_conversation_id: str | None ) -> None: """Update thread with storage type and conversation ID. Args: - chat_client_thread: The thread to update. - responseConversationId: The conversation ID from the response, if any. + thread: The thread to update. + response_conversation_id: The conversation ID from the response, if any. Raises: AgentExecutionException: If conversation ID is missing for service-managed thread. """ - # Set the thread's storage location, the first time that we use it. - if chat_client_thread.storage_location is None: - chat_client_thread.storage_location = ( - ChatClientAgentThreadType.CONVERSATION_ID - if responseConversationId is not None - else ChatClientAgentThreadType.IN_MEMORY_MESSAGES + if response_conversation_id is None and thread.service_thread_id is not None: + # We were passed a thread that is service managed, but we got no conversation id back from the chat client, + # meaning the service doesn't support service managed threads, + # so the thread cannot be used with this service. + raise AgentExecutionException( + "Service did not return a valid conversation id when using a service managed thread." ) - # If we got a conversation id back from the chat client, it means that the service supports server side thread - # storage so we should capture the id and update the thread with the new id. - if chat_client_thread.storage_location == ChatClientAgentThreadType.CONVERSATION_ID: - if responseConversationId is None: - raise AgentExecutionException( - "Service did not return a valid conversation id when using a service managed thread." - ) - chat_client_thread.id = responseConversationId + if response_conversation_id is not None: + # If we got a conversation id back from the chat client, it means that the service + # supports server side thread storage so we should update the thread with the new id. + thread.service_thread_id = response_conversation_id + elif thread.message_store is None and self.chat_message_store_factory is not None: + # If the service doesn't use service side thread storage (i.e. we got no id back from invocation), and + # the thread has no message_store yet, and we have a custom messages store, we should update the thread + # with the custom message_store so that it has somewhere to store the chat history. + thread.message_store = self.chat_message_store_factory() async def _prepare_thread_and_messages( self, *, thread: AgentThread | None, input_messages: list[ChatMessage] | None = None, - ) -> tuple[ChatClientAgentThread, list[ChatMessage]]: + ) -> tuple[AgentThread, list[ChatMessage]]: """Prepare the messages for agent execution. Args: @@ -755,19 +627,14 @@ class ChatClientAgent(AgentBase): Raises: AgentExecutionException: If the thread is not of the expected type. """ - validated_thread: ChatClientAgentThread = self._validate_or_create_thread_type( # type: ignore[reportAssignmentType] - thread=thread, - construct_thread=self.get_new_thread, - expected_type=ChatClientAgentThread, - ) + thread = thread or self.get_new_thread() + messages: list[ChatMessage] = [] if self.instructions: messages.append(ChatMessage(role=ChatRole.SYSTEM, text=self.instructions)) - if isinstance(validated_thread, MessagesRetrievableThread): - async for message in validated_thread.get_messages(): - messages.append(message) + messages.extend(await thread.list_messages() or []) messages.extend(input_messages or []) - return validated_thread, messages + return thread, messages def _normalize_messages( self, diff --git a/python/packages/main/agent_framework/_clients.py b/python/packages/main/agent_framework/_clients.py index 61a077f15b..68cfbd8ee1 100644 --- a/python/packages/main/agent_framework/_clients.py +++ b/python/packages/main/agent_framework/_clients.py @@ -10,6 +10,7 @@ from pydantic import BaseModel from ._logging import get_logger from ._pydantic import AFBaseModel +from ._threads import ChatMessageStore from ._tools import AIFunction, AITool from ._types import ( AIContents, @@ -645,6 +646,7 @@ class ChatClientBase(AFBaseModel, ABC): | MutableMapping[str, Any] | list[MutableMapping[str, Any]] | None = None, + chat_message_store_factory: Callable[[], ChatMessageStore] | None = None, **kwargs: Any, ) -> "ChatClientAgent": """Create an agent with the given name and instructions. @@ -653,6 +655,8 @@ class ChatClientBase(AFBaseModel, ABC): name: The name of the agent. instructions: The instructions for the agent. tools: Optional list of tools to associate with the agent. + chat_message_store_factory: Factory function to create an instance of ChatMessageStore. If not provided, + the default in-memory store will be used. **kwargs: Additional keyword arguments to pass to the agent. See ChatClientAgent for all the available options. @@ -661,7 +665,14 @@ class ChatClientBase(AFBaseModel, ABC): """ from ._agents import ChatClientAgent - return ChatClientAgent(chat_client=self, name=name, instructions=instructions, tools=tools, **kwargs) + return ChatClientAgent( + chat_client=self, + name=name, + instructions=instructions, + tools=tools, + chat_message_store_factory=chat_message_store_factory, + **kwargs, + ) # region Embedding Client diff --git a/python/packages/main/agent_framework/_threads.py b/python/packages/main/agent_framework/_threads.py new file mode 100644 index 0000000000..cb926825a5 --- /dev/null +++ b/python/packages/main/agent_framework/_threads.py @@ -0,0 +1,375 @@ +# Copyright (c) Microsoft. All rights reserved. + +from collections.abc import Sequence +from typing import Any, Protocol, overload + +from ._pydantic import AFBaseModel +from ._types import ChatMessage + +__all__ = ["AgentThread", "ChatMessageList", "ChatMessageStore"] + + +class ChatMessageStore(Protocol): + """Defines methods for storing and retrieving chat messages associated with a specific thread. + + Implementations of this protocol are responsible for managing the storage of chat messages, + including handling large volumes of data by truncating or summarizing messages as necessary. + """ + + async def list_messages(self) -> list[ChatMessage]: + """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. + + If the messages stored in the store become very large, it is up to the store to + truncate, summarize or otherwise limit the number of messages returned. + + When using implementations of ChatMessageStore, a new one should be created for each thread + since they may contain state that is specific to a thread. + """ + ... + + async def add_messages(self, messages: Sequence[ChatMessage]) -> None: + """Adds messages to the store.""" + ... + + async def deserialize_state(self, serialized_store_state: Any, **kwargs: Any) -> None: + """Deserializes the state into the properties on this store. + + This method, together with serialize_state can be used to save and load messages from a persistent store + if this store only has messages in memory. + """ + ... + + async def serialize_state(self, **kwargs: Any) -> Any: + """Serializes the current object's state. + + This method, together with deserialize_state can be used to save and load messages from a persistent store + if this store only has messages in memory. + """ + ... + + +class AgentThread(AFBaseModel): + """Base class for agent threads.""" + + _service_thread_id: str | None = None + _message_store: ChatMessageStore | None = None + + @overload + def __init__(self) -> None: + """Initialize an empty AgentThread with no service thread ID or message store.""" + ... + + @overload + def __init__(self, service_thread_id: str) -> None: + """Initialize an AgentThread with a service thread ID. + + Args: + service_thread_id: The ID of the thread managed by the agent service. + """ + ... + + @overload + def __init__(self, *, message_store: ChatMessageStore) -> None: + """Initialize an AgentThread with a custom message store. + + Args: + message_store: The ChatMessageStore implementation for managing chat messages. + """ + ... + + def __init__(self, service_thread_id: str | None = None, *, message_store: ChatMessageStore | None = None) -> None: + """Initialize an AgentThread. + + Args: + service_thread_id: Optional ID of the thread managed by the agent service. + message_store: Optional ChatMessageStore implementation for managing chat messages. + + Note: + Either service_thread_id or message_store may be set, but not both. + """ + super().__init__() + + self.service_thread_id = service_thread_id + self.message_store = message_store + + @property + def service_thread_id(self) -> str | None: + """Gets the ID of the current thread to support cases where the thread is owned by the agent service.""" + return self._service_thread_id + + @service_thread_id.setter + def service_thread_id(self, service_thread_id: str | None) -> None: + """Sets the ID of the current thread to support cases where the thread is owned by the agent service. + + Note that either service_thread_id or message_store may be set, but not both. + """ + if not self._service_thread_id and not service_thread_id: + return + + if self._message_store is not None: + raise ValueError( + "Only the service_thread_id or message_store may be set, " + "but not both and switching from one to another is not supported." + ) + + self._service_thread_id = service_thread_id + + @property + def message_store(self) -> ChatMessageStore | None: + """Gets the ChatMessageStore used by this thread, when messages should be stored in a custom location.""" + return self._message_store + + @message_store.setter + def message_store(self, message_store: ChatMessageStore | None) -> None: + """Sets the ChatMessageStore used by this thread, when messages should be stored in a custom location. + + Note that either service_thread_id or message_store may be set, but not both. + """ + if self._message_store is None and message_store is None: + return + + if self._service_thread_id: + raise ValueError( + "Only the service_thread_id or message_store may be set, " + "but not both and switching from one to another is not supported." + ) + + self._message_store = message_store + + async def list_messages(self) -> list[ChatMessage] | None: + """Retrieves any messages stored in ChatMessageStore of the thread, otherwise returns an empty collection.""" + return await self._message_store.list_messages() if self._message_store is not None else None + + async def serialize(self, **kwargs: Any) -> dict[str, Any]: + """Serializes the current object's state. + + Args: + **kwargs: Arguments for serialization. + """ + chat_message_store_state = None + if self._message_store is not None: + chat_message_store_state = await self._message_store.serialize_state(**kwargs) + + state = ThreadState( + service_thread_id=self._service_thread_id, chat_message_store_state=chat_message_store_state + ) + + return state.model_dump() + + +async def thread_on_new_messages(thread: AgentThread, new_messages: ChatMessage | Sequence[ChatMessage]) -> None: + """Invoked when a new message has been contributed to the chat by any participant.""" + if thread.service_thread_id is not None: + # If the thread messages are stored in the service there is nothing to do here, + # since invoking the service should already update the thread. + return + + if thread.message_store is None: + # If there is no conversation id, and no store we can + # create a default in memory store. + thread.message_store = ChatMessageList() + + # If a store has been provided, we need to add the messages to the store. + if isinstance(new_messages, ChatMessage): + new_messages = [new_messages] + + await thread.message_store.add_messages(new_messages) + + +async def deserialize_thread_state( + thread: AgentThread, + serialized_thread: dict[str, Any], + **kwargs: Any, +) -> None: + """Deserializes the state from a dictionary into the thread properties.""" + state = ThreadState.model_validate(serialized_thread) + + if state.service_thread_id: + thread.service_thread_id = state.service_thread_id + # Since we have an ID, we should not have a chat message store and we can return here. + return + + # If we don't have any ChatMessageStore state return here. + if state.chat_message_store_state is None: + return + + if thread.message_store is None: + # If we don't have a chat message store yet, create an in-memory one. + thread.message_store = ChatMessageList() + + await thread.message_store.deserialize_state(state.chat_message_store_state, **kwargs) + + +class ThreadState(AFBaseModel): + """State model for serializing and deserializing thread information. + + Attributes: + service_thread_id: Optional ID of the thread managed by the agent service. + chat_message_store_state: Optional serialized state of the chat message store. + """ + + service_thread_id: str | None = None + chat_message_store_state: Any | None = None + + +class StoreState(AFBaseModel): + """State model for serializing and deserializing chat message store data. + + Attributes: + messages: List of chat messages stored in the message store. + """ + + messages: list[ChatMessage] + + +class ChatMessageList: + """An in-memory implementation of ChatMessageStore that stores messages in a list. + + This implementation provides a simple, list-based storage for chat messages + with support for serialization and deserialization. It implements all the + required methods of the ChatMessageStore protocol and provides additional + list-like operations for direct message manipulation. + + The store maintains messages in memory and provides methods to serialize + and deserialize the state for persistence purposes. + """ + + def __init__(self, messages: Sequence[ChatMessage] | None = None) -> None: + """Initialize the message store with optional initial messages. + + Args: + messages: Optional collection of initial ChatMessage objects to store. + """ + self._messages: list[ChatMessage] = [] + if messages: + self._messages.extend(messages) + + async def add_messages(self, messages: Sequence[ChatMessage]) -> None: + """Add messages to the store. + + Args: + messages: Sequence of ChatMessage objects to add to the store. + """ + self._messages.extend(messages) + + async def list_messages(self) -> list[ChatMessage]: + """Get all messages from the store in chronological order. + + Returns: + List of ChatMessage objects, ordered from oldest to newest. + """ + return self._messages + + async def deserialize_state(self, serialized_store_state: Any, **kwargs: Any) -> None: + """Deserialize state data into this store instance. + + Args: + serialized_store_state: Previously serialized state data containing messages. + **kwargs: Additional arguments for deserialization. + """ + if serialized_store_state: + state = StoreState.model_validate(obj=serialized_store_state, **kwargs) + if state.messages: + self._messages.extend(state.messages) + + async def serialize_state(self, **kwargs: Any) -> Any: + """Serialize the current store state for persistence. + + Args: + **kwargs: Additional arguments for serialization. + + Returns: + Serialized state data that can be used with deserialize_state. + """ + state = StoreState(messages=self._messages) + return state.model_dump(**kwargs) + + def __len__(self) -> int: + """Return the number of messages in the store. + + Returns: + The count of messages currently stored. + """ + return len(self._messages) + + def __getitem__(self, index: int) -> ChatMessage: + """Get a message by index. + + Args: + index: The index of the message to retrieve. + + Returns: + The ChatMessage at the specified index. + """ + return self._messages[index] + + def __setitem__(self, index: int, item: ChatMessage) -> None: + """Set a message at the specified index. + + Args: + index: The index at which to set the message. + item: The ChatMessage to set at the specified index. + """ + self._messages[index] = item + + def append(self, item: ChatMessage) -> None: + """Append a message to the end of the store. + + Args: + item: The ChatMessage to append. + """ + self._messages.append(item) + + def clear(self) -> None: + """Remove all messages from the store.""" + self._messages.clear() + + def index(self, item: ChatMessage) -> int: + """Return the index of the first occurrence of the specified message. + + Args: + item: The ChatMessage to find. + + Returns: + The index of the first occurrence of the message. + + Raises: + ValueError: If the message is not found in the store. + """ + return self._messages.index(item) + + def insert(self, index: int, item: ChatMessage) -> None: + """Insert a message at the specified index. + + Args: + index: The index at which to insert the message. + item: The ChatMessage to insert. + """ + self._messages.insert(index, item) + + def remove(self, item: ChatMessage) -> None: + """Remove the first occurrence of the specified message from the store. + + Args: + item: The ChatMessage to remove. + + Raises: + ValueError: If the message is not found in the store. + """ + self._messages.remove(item) + + def pop(self, index: int = -1) -> ChatMessage: + """Remove and return a message at the specified index. + + Args: + index: The index of the message to remove and return. Defaults to -1 (last item). + + Returns: + The ChatMessage that was removed. + + Raises: + IndexError: If the index is out of range. + """ + return self._messages.pop(index) diff --git a/python/packages/main/agent_framework/telemetry.py b/python/packages/main/agent_framework/telemetry.py index 76c643205a..7e9745271a 100644 --- a/python/packages/main/agent_framework/telemetry.py +++ b/python/packages/main/agent_framework/telemetry.py @@ -18,8 +18,9 @@ from ._pydantic import AFBaseSettings if TYPE_CHECKING: # pragma: no cover from opentelemetry.util._decorator import _AgnosticContextManager # type: ignore[reportPrivateUsage] - from ._agents import AgentThread, AIAgent, ChatClientAgent + from ._agents import AIAgent, ChatClientAgent from ._clients import ChatClientBase + from ._threads import AgentThread from ._tools import AIFunction from ._types import ( AgentRunResponse, @@ -650,8 +651,8 @@ def _get_agent_run_span( span.set_attribute(GenAIAttributes.AGENT_NAME.value, agent.name) if agent.description: span.set_attribute(GenAIAttributes.AGENT_DESCRIPTION.value, agent.description) - if thread and thread.id: - span.set_attribute(GenAIAttributes.CONVERSATION_ID.value, thread.id) + if thread and thread.service_thread_id: + span.set_attribute(GenAIAttributes.CONVERSATION_ID.value, thread.service_thread_id) if "model" in kwargs: span.set_attribute(GenAIAttributes.MODEL.value, kwargs["model"]) if "seed" in kwargs: diff --git a/python/packages/main/tests/main/test_agents.py b/python/packages/main/tests/main/test_agents.py index 4bb3f874c8..8744baad91 100644 --- a/python/packages/main/tests/main/test_agents.py +++ b/python/packages/main/tests/main/test_agents.py @@ -1,6 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. -from collections.abc import AsyncIterable, MutableSequence, Sequence +from collections.abc import AsyncIterable, MutableSequence from typing import Any from uuid import uuid4 @@ -13,10 +13,9 @@ from agent_framework import ( AIAgent, ChatClient, ChatClientAgent, - ChatClientAgentThread, - ChatClientAgentThreadType, ChatClientBase, ChatMessage, + ChatMessageList, ChatOptions, ChatResponse, ChatResponseUpdate, @@ -28,8 +27,7 @@ from agent_framework.exceptions import AgentExecutionException # Mock AgentThread implementation for testing class MockAgentThread(AgentThread): - async def _on_new_messages(self, new_messages: ChatMessage | Sequence[ChatMessage]) -> None: - pass + pass # Mock Agent implementation for testing @@ -142,58 +140,6 @@ async def test_agent_run_streaming(agent: AIAgent) -> None: assert updates[0].text == "Response" -async def test_chat_client_agent_thread_init_in_memory() -> None: - messages = [ChatMessage(role=ChatRole.USER, contents=[TextContent("Hello")])] - thread = ChatClientAgentThread(messages=messages) - - assert thread.storage_location == ChatClientAgentThreadType.IN_MEMORY_MESSAGES - assert thread.id is None - assert thread.chat_messages == messages - - -async def test_chat_client_agent_thread_empty() -> None: - thread = ChatClientAgentThread() - - assert thread.storage_location is None - assert thread.id is None - assert thread.chat_messages is None - - -async def test_chat_client_agent_thread_init_invalid() -> None: - with raises(ValueError, match="Cannot specify both id and messages"): - ChatClientAgentThread(id="123", messages=[ChatMessage(role=ChatRole.USER, contents=[TextContent("Hello")])]) - - with raises(ValueError, match="ID cannot be empty or whitespace"): - ChatClientAgentThread(id=" ") - - -async def test_chat_client_agent_thread_init_conversation_id() -> None: - thread_id = str(uuid4()) - thread = ChatClientAgentThread(id=thread_id) - - assert thread.storage_location == ChatClientAgentThreadType.CONVERSATION_ID - assert thread.id == thread_id - assert thread.chat_messages is None - - -async def test_chat_client_agent_thread_get_messages() -> None: - messages = [ChatMessage(role=ChatRole.USER, contents=[TextContent("Hello")])] - thread = ChatClientAgentThread(messages=messages) - - result = [msg async for msg in thread.get_messages()] - assert result == messages - - -async def test_chat_client_agent_thread_on_new_messages_in_memory() -> None: - initial_message = ChatMessage(role=ChatRole.USER, contents=[TextContent("Initial message")]) - new_message = ChatMessage(role=ChatRole.USER, contents=[TextContent("New message")]) - - thread = ChatClientAgentThread(messages=[initial_message]) - - await thread._on_new_messages(new_message) # type: ignore[reportPrivateUsage] - assert thread.chat_messages == [initial_message, new_message] - - def test_chat_client_agent_type(chat_client: ChatClient) -> None: chat_client_agent = ChatClientAgent(chat_client=chat_client) assert isinstance(chat_client_agent, AIAgent) @@ -239,24 +185,16 @@ async def test_chat_client_agent_get_new_thread(chat_client: ChatClient) -> None agent = ChatClientAgent(chat_client=chat_client) thread = agent.get_new_thread() - assert isinstance(thread, ChatClientAgentThread) - assert thread.storage_location is None + assert isinstance(thread, AgentThread) async def test_chat_client_agent_prepare_thread_and_messages(chat_client: ChatClient) -> None: agent = ChatClientAgent(chat_client=chat_client) message = ChatMessage(role=ChatRole.USER, text="Hello") - thread = ChatClientAgentThread(messages=[message]) - - result_thread = agent._validate_or_create_thread_type( # type: ignore[reportPrivateUsage] - thread, lambda: ChatClientAgentThread(), expected_type=ChatClientAgentThread - ) # type: ignore[reportPrivateUsage] - - assert result_thread == thread - assert isinstance(result_thread, ChatClientAgentThread) + thread = AgentThread(message_store=ChatMessageList(messages=[message])) _, result_messages = await agent._prepare_thread_and_messages( # type: ignore[reportPrivateUsage] - thread=result_thread, + thread=thread, input_messages=[ChatMessage(role=ChatRole.USER, text="Test")], ) @@ -265,18 +203,6 @@ async def test_chat_client_agent_prepare_thread_and_messages(chat_client: ChatCl assert result_messages[1].text == "Test" -async def test_chat_client_agent_validate_or_create_thread(chat_client: ChatClient) -> None: - agent = ChatClientAgent(chat_client=chat_client) - thread = None - - result_thread = agent._validate_or_create_thread_type( # type: ignore[reportPrivateUsage] - thread, lambda: ChatClientAgentThread(), expected_type=ChatClientAgentThread - ) # type: ignore[reportPrivateUsage] - - assert result_thread != thread - assert isinstance(result_thread, ChatClientAgentThread) - - async def test_chat_client_agent_update_thread_id() -> None: chat_client = MockChatClient( mock_response=ChatResponse( @@ -290,9 +216,7 @@ async def test_chat_client_agent_update_thread_id() -> None: result = await agent.run("Hello", thread=thread) assert result.text == "test response" - assert thread.id == "123" - assert isinstance(thread, ChatClientAgentThread) - assert thread.storage_location == ChatClientAgentThreadType.CONVERSATION_ID + assert thread.service_thread_id == "123" async def test_chat_client_agent_update_thread_messages(chat_client: ChatClient) -> None: @@ -302,19 +226,19 @@ async def test_chat_client_agent_update_thread_messages(chat_client: ChatClient) result = await agent.run("Hello", thread=thread) assert result.text == "test response" - assert thread.id is None - assert isinstance(thread, ChatClientAgentThread) - assert thread.storage_location == ChatClientAgentThreadType.IN_MEMORY_MESSAGES + assert thread.service_thread_id is None - assert thread.chat_messages is not None - assert len(thread.chat_messages) == 2 - assert thread.chat_messages[0].text == "Hello" - assert thread.chat_messages[1].text == "test response" + chat_messages: list[ChatMessage] | None = await thread.list_messages() + + assert chat_messages is not None + assert len(chat_messages) == 2 + assert chat_messages[0].text == "Hello" + assert chat_messages[1].text == "test response" async def test_chat_client_agent_update_thread_conversation_id_missing(chat_client: ChatClient) -> None: agent = ChatClientAgent(chat_client=chat_client) - thread = ChatClientAgentThread(id="123") + thread = AgentThread(service_thread_id="123") with raises(AgentExecutionException, match="Service did not return a valid conversation id"): agent._update_thread_with_type_and_conversation_id(thread, None) # type: ignore[reportPrivateUsage] diff --git a/python/packages/main/tests/main/test_threads.py b/python/packages/main/tests/main/test_threads.py new file mode 100644 index 0000000000..83de7ac553 --- /dev/null +++ b/python/packages/main/tests/main/test_threads.py @@ -0,0 +1,482 @@ +# Copyright (c) Microsoft. All rights reserved. + +from collections.abc import Sequence +from typing import Any + +import pytest + +from agent_framework import AgentThread, ChatMessage, ChatMessageList, ChatRole +from agent_framework._threads import StoreState, ThreadState, deserialize_thread_state, thread_on_new_messages + + +class MockChatMessageStore: + """Mock implementation of ChatMessageStore for testing.""" + + def __init__(self, messages: list[ChatMessage] | None = None) -> None: + self._messages = messages or [] + self._serialize_calls = 0 + self._deserialize_calls = 0 + + async def list_messages(self) -> list[ChatMessage]: + return self._messages + + async def add_messages(self, messages: Sequence[ChatMessage]) -> None: + self._messages.extend(messages) + + async def serialize_state(self, **kwargs: Any) -> Any: + self._serialize_calls += 1 + return {"messages": [msg.__dict__ for msg in self._messages], "kwargs": kwargs} + + async def deserialize_state(self, serialized_store_state: Any, **kwargs: Any) -> None: + self._deserialize_calls += 1 + if serialized_store_state and "messages" in serialized_store_state: + self._messages = serialized_store_state["messages"] + + +@pytest.fixture +def sample_messages() -> list[ChatMessage]: + """Fixture providing sample chat messages for testing.""" + return [ + ChatMessage(role=ChatRole.USER, text="Hello", message_id="msg1"), + ChatMessage(role=ChatRole.ASSISTANT, text="Hi there!", message_id="msg2"), + ChatMessage(role=ChatRole.USER, text="How are you?", message_id="msg3"), + ] + + +@pytest.fixture +def sample_message() -> ChatMessage: + """Fixture providing a single sample chat message for testing.""" + return ChatMessage(role=ChatRole.USER, text="Test message", message_id="test1") + + +class TestAgentThread: + """Test cases for AgentThread class.""" + + def test_init_with_no_parameters(self) -> None: + """Test AgentThread initialization with no parameters.""" + thread = AgentThread() + assert thread.service_thread_id is None + assert thread.message_store is None + + def test_init_with_service_thread_id(self) -> None: + """Test AgentThread initialization with service_thread_id.""" + service_thread_id = "test-conversation-123" + thread = AgentThread(service_thread_id=service_thread_id) + assert thread.service_thread_id == service_thread_id + assert thread.message_store is None + + def test_init_with_message_store(self) -> None: + """Test AgentThread initialization with message_store.""" + store = ChatMessageList() + thread = AgentThread(message_store=store) + assert thread.service_thread_id is None + assert thread.message_store is store + + def test_service_thread_id_property_setter(self) -> None: + """Test service_thread_id property setter.""" + thread = AgentThread() + service_thread_id = "test-conversation-456" + + thread.service_thread_id = service_thread_id + assert thread.service_thread_id == service_thread_id + + def test_service_thread_id_setter_with_existing_message_store_raises_error(self) -> None: + """Test that setting service_thread_id when message_store exists raises ValueError.""" + store = ChatMessageList() + thread = AgentThread(message_store=store) + + with pytest.raises(ValueError, match="Only the service_thread_id or message_store may be set"): + thread.service_thread_id = "test-conversation-789" + + def test_service_thread_id_setter_with_none_values(self) -> None: + """Test service_thread_id setter with None values does nothing.""" + thread = AgentThread() + thread.service_thread_id = None # Should not raise error + assert thread.service_thread_id is None + + def test_message_store_property_setter(self) -> None: + """Test message_store property setter.""" + thread = AgentThread() + store = ChatMessageList() + + thread.message_store = store + assert thread.message_store is store + + def test_message_store_setter_with_existing_service_thread_id_raises_error(self) -> None: + """Test that setting message_store when service_thread_id exists raises ValueError.""" + service_thread_id = "test-conversation-999" + thread = AgentThread(service_thread_id=service_thread_id) + store = ChatMessageList() + + with pytest.raises(ValueError, match="Only the service_thread_id or message_store may be set"): + thread.message_store = store + + def test_message_store_setter_with_none_values(self) -> None: + """Test message_store setter with None values does nothing.""" + thread = AgentThread() + thread.message_store = None # Should not raise error + assert thread.message_store is None + + async def test_get_messages_with_message_store(self, sample_messages: list[ChatMessage]) -> None: + """Test get_messages when message_store is set.""" + store = ChatMessageList(sample_messages) + thread = AgentThread(message_store=store) + + messages: list[ChatMessage] | None = await thread.list_messages() + + assert messages is not None + assert len(messages) == 3 + assert messages[0].text == "Hello" + assert messages[1].text == "Hi there!" + assert messages[2].text == "How are you?" + + async def test_get_messages_with_no_message_store(self) -> None: + """Test get_messages when no message_store is set.""" + thread = AgentThread() + + messages: list[ChatMessage] | None = await thread.list_messages() + + assert messages is None + + async def test_on_new_messages_with_service_thread_id(self, sample_message: ChatMessage) -> None: + """Test _on_new_messages when service_thread_id is set (should do nothing).""" + thread = AgentThread(service_thread_id="test-conv") + + await thread_on_new_messages(thread, sample_message) + + # Should not create a message store + assert thread.message_store is None + + async def test_on_new_messages_single_message_creates_store(self, sample_message: ChatMessage) -> None: + """Test _on_new_messages with single message creates ChatMessageList.""" + thread = AgentThread() + + await thread_on_new_messages(thread, sample_message) + + assert thread.message_store is not None + assert isinstance(thread.message_store, ChatMessageList) + messages = await thread.message_store.list_messages() + assert len(messages) == 1 + assert messages[0].text == "Test message" + + async def test_on_new_messages_multiple_messages(self, sample_messages: list[ChatMessage]) -> None: + """Test _on_new_messages with multiple messages.""" + thread = AgentThread() + + await thread_on_new_messages(thread, sample_messages) + + assert thread.message_store is not None + messages = await thread.message_store.list_messages() + assert len(messages) == 3 + + async def test_on_new_messages_with_existing_store(self, sample_message: ChatMessage) -> None: + """Test _on_new_messages adds to existing message store.""" + initial_messages = [ChatMessage(role=ChatRole.USER, text="Initial", message_id="init1")] + store = ChatMessageList(initial_messages) + thread = AgentThread(message_store=store) + + await thread_on_new_messages(thread, sample_message) + + assert thread.message_store is not None + messages = await thread.message_store.list_messages() + assert len(messages) == 2 + assert messages[0].text == "Initial" + assert messages[1].text == "Test message" + + async def test_deserialize_with_service_thread_id(self) -> None: + """Test _deserialize with service_thread_id.""" + thread = AgentThread() + serialized_data = {"service_thread_id": "test-conv-123", "chat_message_store_state": None} + + await deserialize_thread_state(thread, serialized_data) + + assert thread.service_thread_id == "test-conv-123" + assert thread.message_store is None + + async def test_deserialize_with_store_state(self, sample_messages: list[ChatMessage]) -> None: + """Test _deserialize with chat_message_store_state.""" + thread = AgentThread() + store_state = {"messages": sample_messages} + serialized_data = {"service_thread_id": None, "chat_message_store_state": store_state} + + await deserialize_thread_state(thread, serialized_data) + + assert thread.service_thread_id is None + assert thread.message_store is not None + assert isinstance(thread.message_store, ChatMessageList) + + async def test_deserialize_with_no_state(self) -> None: + """Test _deserialize with no state.""" + thread = AgentThread() + serialized_data = {"service_thread_id": None, "chat_message_store_state": None} + + await deserialize_thread_state(thread, serialized_data) + + assert thread.service_thread_id is None + assert thread.message_store is None + + async def test_deserialize_with_existing_store(self) -> None: + """Test _deserialize with existing message store.""" + store = MockChatMessageStore() + thread = AgentThread(message_store=store) + serialized_data: dict[str, Any] = {"service_thread_id": None, "chat_message_store_state": {"messages": []}} + + await deserialize_thread_state(thread, serialized_data) + + assert store._deserialize_calls == 1 # pyright: ignore[reportPrivateUsage] + + async def test_serialize_with_service_thread_id(self) -> None: + """Test serialize with service_thread_id.""" + thread = AgentThread(service_thread_id="test-conv-456") + + result = await thread.serialize() + + assert result["service_thread_id"] == "test-conv-456" + assert result["chat_message_store_state"] is None + + async def test_serialize_with_message_store(self) -> None: + """Test serialize with message_store.""" + store = MockChatMessageStore() + thread = AgentThread(message_store=store) + + result = await thread.serialize() + + assert result["service_thread_id"] is None + assert result["chat_message_store_state"] is not None + assert store._serialize_calls == 1 # pyright: ignore[reportPrivateUsage] + + async def test_serialize_with_no_state(self) -> None: + """Test serialize with no state.""" + thread = AgentThread() + + result = await thread.serialize() + + assert result["service_thread_id"] is None + assert result["chat_message_store_state"] is None + + async def test_serialize_with_kwargs(self) -> None: + """Test serialize passes kwargs to message store.""" + store = MockChatMessageStore() + thread = AgentThread(message_store=store) + + await thread.serialize(custom_param="test_value") + + assert store._serialize_calls == 1 # pyright: ignore[reportPrivateUsage] + + +class TestChatMessageList: + """Test cases for ChatMessageList class.""" + + def test_init_empty(self) -> None: + """Test ChatMessageList initialization with no messages.""" + store = ChatMessageList() + assert len(store) == 0 + + def test_init_with_messages(self, sample_messages: list[ChatMessage]) -> None: + """Test ChatMessageList initialization with messages.""" + store = ChatMessageList(sample_messages) + assert len(store) == 3 + + async def test_add_messages(self, sample_messages: list[ChatMessage]) -> None: + """Test adding messages to the store.""" + store = ChatMessageList() + + await store.add_messages(sample_messages) + + assert len(store) == 3 + messages = await store.list_messages() + assert messages[0].text == "Hello" + + async def test_get_messages(self, sample_messages: list[ChatMessage]) -> None: + """Test getting messages from the store.""" + store = ChatMessageList(sample_messages) + + messages = await store.list_messages() + + assert len(messages) == 3 + assert messages[0].message_id == "msg1" + + async def test_serialize_state(self, sample_messages: list[ChatMessage]) -> None: + """Test serializing store state.""" + store = ChatMessageList(sample_messages) + + result = await store.serialize_state() + + assert "messages" in result + assert len(result["messages"]) == 3 + + async def test_serialize_state_empty(self) -> None: + """Test serializing empty store state.""" + store = ChatMessageList() + + result = await store.serialize_state() + + assert "messages" in result + assert len(result["messages"]) == 0 + + async def test_deserialize_state(self, sample_messages: list[ChatMessage]) -> None: + """Test deserializing store state.""" + store = ChatMessageList() + state_data = {"messages": sample_messages} + + await store.deserialize_state(state_data) + + messages = await store.list_messages() + assert len(messages) == 3 + assert messages[0].text == "Hello" + + async def test_deserialize_state_none(self) -> None: + """Test deserializing None state.""" + store = ChatMessageList() + + await store.deserialize_state(None) + + assert len(store) == 0 + + async def test_deserialize_state_empty(self) -> None: + """Test deserializing empty state.""" + store = ChatMessageList() + + await store.deserialize_state({}) + + assert len(store) == 0 + + def test_len(self, sample_messages: list[ChatMessage]) -> None: + """Test __len__ method.""" + store = ChatMessageList(sample_messages) + assert len(store) == 3 + + empty_store = ChatMessageList() + assert len(empty_store) == 0 + + def test_getitem(self, sample_messages: list[ChatMessage]) -> None: + """Test __getitem__ method.""" + store = ChatMessageList(sample_messages) + + assert store[0].text == "Hello" + assert store[1].text == "Hi there!" + assert store[2].text == "How are you?" + + def test_setitem(self, sample_messages: list[ChatMessage], sample_message: ChatMessage) -> None: + """Test __setitem__ method.""" + store = ChatMessageList(sample_messages) + + store[1] = sample_message + assert store[1].text == "Test message" + assert store[1].message_id == "test1" + + def test_append(self, sample_message: ChatMessage) -> None: + """Test append method.""" + store = ChatMessageList() + + store.append(sample_message) + + assert len(store) == 1 + assert store[0].text == "Test message" + + def test_clear(self, sample_messages: list[ChatMessage]) -> None: + """Test clear method.""" + store = ChatMessageList(sample_messages) + assert len(store) == 3 + + store.clear() + assert len(store) == 0 + + def test_index(self, sample_messages: list[ChatMessage]) -> None: + """Test index method.""" + store = ChatMessageList(sample_messages) + + index = store.index(sample_messages[1]) + assert index == 1 + + def test_insert(self, sample_messages: list[ChatMessage], sample_message: ChatMessage) -> None: + """Test insert method.""" + store = ChatMessageList(sample_messages) + + store.insert(1, sample_message) + + assert len(store) == 4 + assert store[1].text == "Test message" + assert store[2].text == "Hi there!" # Original message at index 1 is now at index 2 + + def test_remove(self, sample_messages: list[ChatMessage]) -> None: + """Test remove method.""" + store = ChatMessageList(sample_messages) + message_to_remove = sample_messages[1] + + store.remove(message_to_remove) + + assert len(store) == 2 + assert store[0].text == "Hello" + assert store[1].text == "How are you?" + + def test_pop_default(self, sample_messages: list[ChatMessage]) -> None: + """Test pop method with default index.""" + store = ChatMessageList(sample_messages) + + popped_message = store.pop() + + assert len(store) == 2 + assert popped_message.text == "How are you?" # Last message + + def test_pop_with_index(self, sample_messages: list[ChatMessage]) -> None: + """Test pop method with specific index.""" + store = ChatMessageList(sample_messages) + + popped_message = store.pop(1) + + assert len(store) == 2 + assert popped_message.text == "Hi there!" + assert store[0].text == "Hello" + assert store[1].text == "How are you?" + + +class TestStoreState: + """Test cases for StoreState class.""" + + def test_init(self, sample_messages: list[ChatMessage]) -> None: + """Test StoreState initialization.""" + state = StoreState(messages=sample_messages) + + assert len(state.messages) == 3 + assert state.messages[0].text == "Hello" + + def test_init_empty(self) -> None: + """Test StoreState initialization with empty messages.""" + state = StoreState(messages=[]) + + assert len(state.messages) == 0 + + +class TestThreadState: + """Test cases for ThreadState class.""" + + def test_init_with_service_thread_id(self) -> None: + """Test ThreadState initialization with service_thread_id.""" + state = ThreadState(service_thread_id="test-conv-123") + + assert state.service_thread_id == "test-conv-123" + assert state.chat_message_store_state is None + + def test_init_with_chat_message_store_state(self) -> None: + """Test ThreadState initialization with chat_message_store_state.""" + store_data: dict[str, Any] = {"messages": []} + state = ThreadState(chat_message_store_state=store_data) + + assert state.service_thread_id is None + assert state.chat_message_store_state == store_data + + def test_init_with_both(self) -> None: + """Test ThreadState initialization with both parameters.""" + store_data: dict[str, Any] = {"messages": []} + state = ThreadState(service_thread_id="test-conv-456", chat_message_store_state=store_data) + + assert state.service_thread_id == "test-conv-456" + assert state.chat_message_store_state == store_data + + def test_init_defaults(self) -> None: + """Test ThreadState initialization with defaults.""" + state = ThreadState() + + assert state.service_thread_id is None + assert state.chat_message_store_state is None diff --git a/python/samples/getting_started/agents/azure_assistants_client/azure_assistants_with_thread.py b/python/samples/getting_started/agents/azure_assistants_client/azure_assistants_with_thread.py index 749eee597e..2f79c007e2 100644 --- a/python/samples/getting_started/agents/azure_assistants_client/azure_assistants_with_thread.py +++ b/python/samples/getting_started/agents/azure_assistants_client/azure_assistants_with_thread.py @@ -4,7 +4,7 @@ import asyncio from random import randint from typing import Annotated -from agent_framework import ChatClientAgent, ChatClientAgentThread +from agent_framework import AgentThread, ChatClientAgent from agent_framework.azure import AzureAssistantsClient from azure.identity import DefaultAzureCredential from pydantic import Field @@ -95,7 +95,7 @@ async def example_with_existing_thread_id() -> None: print(f"Agent: {result1.text}") # The thread ID is set after the first response - existing_thread_id = thread.id + existing_thread_id = thread.service_thread_id print(f"Thread ID: {existing_thread_id}") if existing_thread_id: @@ -108,7 +108,7 @@ async def example_with_existing_thread_id() -> None: tools=get_weather, ) as agent: # Create a thread with the existing ID - thread = ChatClientAgentThread(id=existing_thread_id) + thread = AgentThread(service_thread_id=existing_thread_id) query2 = "What was the last city I asked about?" print(f"User: {query2}") diff --git a/python/samples/getting_started/agents/azure_chat_client/azure_chat_client_with_thread.py b/python/samples/getting_started/agents/azure_chat_client/azure_chat_client_with_thread.py index 7a256286ab..e9f26a820b 100644 --- a/python/samples/getting_started/agents/azure_chat_client/azure_chat_client_with_thread.py +++ b/python/samples/getting_started/agents/azure_chat_client/azure_chat_client_with_thread.py @@ -4,7 +4,7 @@ import asyncio from random import randint from typing import Annotated -from agent_framework import ChatClientAgent, ChatClientAgentThread +from agent_framework import AgentThread, ChatClientAgent, ChatMessageList from agent_framework.azure import AzureChatClient from azure.identity import DefaultAzureCredential from pydantic import Field @@ -88,7 +88,6 @@ async def example_with_existing_thread_messages() -> None: # Start a conversation and build up message history thread = agent.get_new_thread() - assert isinstance(thread, ChatClientAgentThread) # Ensure we have the right type query1 = "What's the weather in Paris?" print(f"User: {query1}") @@ -96,7 +95,9 @@ async def example_with_existing_thread_messages() -> None: print(f"Agent: {result1.text}") # The thread now contains the conversation history in memory - message_count = len(thread.chat_messages or []) + messages = await thread.list_messages() + + message_count = len(messages or []) print(f"Thread contains {message_count} messages") print("\n--- Continuing with the same thread in a new agent instance ---") @@ -118,8 +119,9 @@ async def example_with_existing_thread_messages() -> None: print("\n--- Alternative: Creating a new thread from existing messages ---") # You can also create a new thread from existing messages - existing_messages = thread.chat_messages or [] - new_thread = ChatClientAgentThread(messages=existing_messages) + messages = await thread.list_messages() + + new_thread = AgentThread(message_store=ChatMessageList(messages)) query3 = "How does the Paris weather compare to London?" print(f"User: {query3}") diff --git a/python/samples/getting_started/agents/azure_responses_client/azure_responses_client_with_thread.py b/python/samples/getting_started/agents/azure_responses_client/azure_responses_client_with_thread.py index 0fd3284738..e8ae1aa3d4 100644 --- a/python/samples/getting_started/agents/azure_responses_client/azure_responses_client_with_thread.py +++ b/python/samples/getting_started/agents/azure_responses_client/azure_responses_client_with_thread.py @@ -4,7 +4,7 @@ import asyncio from random import randint from typing import Annotated -from agent_framework import ChatClientAgent, ChatClientAgentThread +from agent_framework import AgentThread, ChatClientAgent from agent_framework.azure import AzureResponsesClient from azure.identity import DefaultAzureCredential from pydantic import Field @@ -57,28 +57,27 @@ async def example_with_thread_persistence_in_memory() -> None: # Create a new thread that will be reused thread = agent.get_new_thread() - assert isinstance(thread, ChatClientAgentThread) # First conversation query1 = "What's the weather like in Tokyo?" print(f"User: {query1}") result1 = await agent.run(query1, thread=thread) print(f"Agent: {result1.text}") - print(f"Thread contains {len(thread.chat_messages or [])} messages in-memory.") + print(f"Thread contains {len(await thread.list_messages() or [])} messages in-memory.") # Second conversation using the same thread - maintains context query2 = "How about London?" print(f"\nUser: {query2}") result2 = await agent.run(query2, thread=thread) print(f"Agent: {result2.text}") - print(f"Thread contains {len(thread.chat_messages or [])} messages in-memory.") + print(f"Thread contains {len(await thread.list_messages() or [])} messages in-memory.") # Third conversation - agent should remember both previous cities query3 = "Which of the cities I asked about has better weather?" print(f"\nUser: {query3}") result3 = await agent.run(query3, thread=thread) print(f"Agent: {result3.text}") - print(f"Thread contains {len(thread.chat_messages or [])} messages in-memory.") + print(f"Thread contains {len(await thread.list_messages() or [])} messages in-memory.") print("Note: The agent remembers context from previous messages in the same thread.\n") @@ -100,17 +99,16 @@ async def example_with_existing_thread_id() -> None: # Start a conversation and get the thread ID thread = agent.get_new_thread() - assert isinstance(thread, ChatClientAgentThread) query1 = "What's the weather in Paris?" print(f"User: {query1}") # Enable Azure OpenAI conversation state by setting `store` parameter to True result1 = await agent.run(query1, thread=thread, store=True) print(f"Agent: {result1.text}") - print(f"Thread contains {len(thread.chat_messages or [])} messages in-memory.") + print(f"Thread contains {len(await thread.list_messages() or [])} messages in-memory.") # The thread ID is set after the first response - existing_thread_id = thread.id + existing_thread_id = thread.service_thread_id print(f"Thread ID: {existing_thread_id}") if existing_thread_id: @@ -123,13 +121,13 @@ async def example_with_existing_thread_id() -> None: ) # Create a thread with the existing ID - thread = ChatClientAgentThread(id=existing_thread_id) + thread = AgentThread(service_thread_id=existing_thread_id) query2 = "What was the last city I asked about?" print(f"User: {query2}") result2 = await agent.run(query2, thread=thread, store=True) print(f"Agent: {result2.text}") - print(f"Thread contains {len(thread.chat_messages or [])} messages in-memory.") + print(f"Thread contains {len(await thread.list_messages() or [])} messages in-memory.") print("Note: The agent continues the conversation from the previous thread by using thread ID.\n") diff --git a/python/samples/getting_started/agents/foundry/foundry_with_thread.py b/python/samples/getting_started/agents/foundry/foundry_with_thread.py index 9e7a6a60fb..6a4d9a046f 100644 --- a/python/samples/getting_started/agents/foundry/foundry_with_thread.py +++ b/python/samples/getting_started/agents/foundry/foundry_with_thread.py @@ -4,7 +4,7 @@ import asyncio from random import randint from typing import Annotated -from agent_framework import ChatClientAgent, ChatClientAgentThread +from agent_framework import AgentThread, ChatClientAgent from agent_framework.foundry import FoundryChatClient from azure.identity.aio import DefaultAzureCredential from pydantic import Field @@ -104,7 +104,7 @@ async def example_with_existing_thread_id() -> None: print(f"Agent: {result1.text}") # The thread ID is set after the first response - existing_thread_id = thread.id + existing_thread_id = thread.service_thread_id print(f"Thread ID: {existing_thread_id}") if existing_thread_id: @@ -120,7 +120,7 @@ async def example_with_existing_thread_id() -> None: ) as agent, ): # Create a thread with the existing ID - thread = ChatClientAgentThread(id=existing_thread_id) + thread = AgentThread(service_thread_id=existing_thread_id) query2 = "What was the last city I asked about?" print(f"User: {query2}") diff --git a/python/samples/getting_started/agents/openai_assistants_client/openai_assistants_with_thread.py b/python/samples/getting_started/agents/openai_assistants_client/openai_assistants_with_thread.py index 71dd5efff7..875e276585 100644 --- a/python/samples/getting_started/agents/openai_assistants_client/openai_assistants_with_thread.py +++ b/python/samples/getting_started/agents/openai_assistants_client/openai_assistants_with_thread.py @@ -4,7 +4,7 @@ import asyncio from random import randint from typing import Annotated -from agent_framework import ChatClientAgent, ChatClientAgentThread +from agent_framework import AgentThread, ChatClientAgent from agent_framework.openai import OpenAIAssistantsClient from pydantic import Field @@ -94,7 +94,7 @@ async def example_with_existing_thread_id() -> None: print(f"Agent: {result1.text}") # The thread ID is set after the first response - existing_thread_id = thread.id + existing_thread_id = thread.service_thread_id print(f"Thread ID: {existing_thread_id}") if existing_thread_id: @@ -107,7 +107,7 @@ async def example_with_existing_thread_id() -> None: tools=get_weather, ) as agent: # Create a thread with the existing ID - thread = ChatClientAgentThread(id=existing_thread_id) + thread = AgentThread(service_thread_id=existing_thread_id) query2 = "What was the last city I asked about?" print(f"User: {query2}") diff --git a/python/samples/getting_started/agents/openai_chat_client/openai_chat_client_with_thread.py b/python/samples/getting_started/agents/openai_chat_client/openai_chat_client_with_thread.py index cb48ef7b75..05701939fd 100644 --- a/python/samples/getting_started/agents/openai_chat_client/openai_chat_client_with_thread.py +++ b/python/samples/getting_started/agents/openai_chat_client/openai_chat_client_with_thread.py @@ -4,7 +4,7 @@ import asyncio from random import randint from typing import Annotated -from agent_framework import ChatClientAgent, ChatClientAgentThread +from agent_framework import AgentThread, ChatClientAgent, ChatMessageList from agent_framework.openai import OpenAIChatClient from pydantic import Field @@ -87,7 +87,6 @@ async def example_with_existing_thread_messages() -> None: # Start a conversation and build up message history thread = agent.get_new_thread() - assert isinstance(thread, ChatClientAgentThread) # Ensure we have the right type query1 = "What's the weather in Paris?" print(f"User: {query1}") @@ -95,7 +94,9 @@ async def example_with_existing_thread_messages() -> None: print(f"Agent: {result1.text}") # The thread now contains the conversation history in memory - message_count = len(thread.chat_messages or []) + messages = await thread.list_messages() + + message_count = len(messages or []) print(f"Thread contains {message_count} messages") print("\n--- Continuing with the same thread in a new agent instance ---") @@ -117,8 +118,9 @@ async def example_with_existing_thread_messages() -> None: print("\n--- Alternative: Creating a new thread from existing messages ---") # You can also create a new thread from existing messages - existing_messages = thread.chat_messages or [] - new_thread = ChatClientAgentThread(messages=existing_messages) + messages = await thread.list_messages() + + new_thread = AgentThread(message_store=ChatMessageList(messages)) query3 = "How does the Paris weather compare to London?" print(f"User: {query3}") diff --git a/python/samples/getting_started/agents/openai_responses_client/openai_responses_client_with_thread.py b/python/samples/getting_started/agents/openai_responses_client/openai_responses_client_with_thread.py index 8c735c9361..da050ac9ce 100644 --- a/python/samples/getting_started/agents/openai_responses_client/openai_responses_client_with_thread.py +++ b/python/samples/getting_started/agents/openai_responses_client/openai_responses_client_with_thread.py @@ -4,7 +4,7 @@ import asyncio from random import randint from typing import Annotated -from agent_framework import ChatClientAgent, ChatClientAgentThread +from agent_framework import AgentThread, ChatClientAgent from agent_framework.openai import OpenAIResponsesClient from pydantic import Field @@ -56,28 +56,29 @@ async def example_with_thread_persistence_in_memory() -> None: # Create a new thread that will be reused thread = agent.get_new_thread() - assert isinstance(thread, ChatClientAgentThread) # First conversation query1 = "What's the weather like in Tokyo?" print(f"User: {query1}") result1 = await agent.run(query1, thread=thread) print(f"Agent: {result1.text}") - print(f"Thread contains {len(thread.chat_messages or [])} messages in-memory.") + + print(f"Thread contains {len(await thread.list_messages() or [])} messages in-memory.") # Second conversation using the same thread - maintains context query2 = "How about London?" print(f"\nUser: {query2}") result2 = await agent.run(query2, thread=thread) print(f"Agent: {result2.text}") - print(f"Thread contains {len(thread.chat_messages or [])} messages in-memory.") + + print(f"Thread contains {len(await thread.list_messages() or [])} messages in-memory.") # Third conversation - agent should remember both previous cities query3 = "Which of the cities I asked about has better weather?" print(f"\nUser: {query3}") result3 = await agent.run(query3, thread=thread) print(f"Agent: {result3.text}") - print(f"Thread contains {len(thread.chat_messages or [])} messages in-memory.") + print(f"Thread contains {len(await thread.list_messages() or [])} messages in-memory.") print("Note: The agent remembers context from previous messages in the same thread.\n") @@ -99,17 +100,16 @@ async def example_with_existing_thread_id() -> None: # Start a conversation and get the thread ID thread = agent.get_new_thread() - assert isinstance(thread, ChatClientAgentThread) query1 = "What's the weather in Paris?" print(f"User: {query1}") # Enable OpenAI conversation state by setting `store` parameter to True result1 = await agent.run(query1, thread=thread, store=True) print(f"Agent: {result1.text}") - print(f"Thread contains {len(thread.chat_messages or [])} messages in-memory.") + print(f"Thread contains {len(await thread.list_messages() or [])} messages in-memory.") # The thread ID is set after the first response - existing_thread_id = thread.id + existing_thread_id = thread.service_thread_id print(f"Thread ID: {existing_thread_id}") if existing_thread_id: @@ -122,13 +122,13 @@ async def example_with_existing_thread_id() -> None: ) # Create a thread with the existing ID - thread = ChatClientAgentThread(id=existing_thread_id) + thread = AgentThread(service_thread_id=existing_thread_id) query2 = "What was the last city I asked about?" print(f"User: {query2}") result2 = await agent.run(query2, thread=thread, store=True) print(f"Agent: {result2.text}") - print(f"Thread contains {len(thread.chat_messages or [])} messages in-memory.") + print(f"Thread contains {len(await thread.list_messages() or [])} messages in-memory.") print("Note: The agent continues the conversation from the previous thread by using thread ID.\n") diff --git a/python/samples/getting_started/threads/custom_chat_message_store_thread.py b/python/samples/getting_started/threads/custom_chat_message_store_thread.py new file mode 100644 index 0000000000..497e277f26 --- /dev/null +++ b/python/samples/getting_started/threads/custom_chat_message_store_thread.py @@ -0,0 +1,82 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +from collections.abc import Collection +from typing import Any + +from agent_framework import ChatMessage, ChatMessageStore +from agent_framework.openai import OpenAIChatClient +from pydantic import BaseModel + + +class CustomStoreState(BaseModel): + """Implementation of custom chat message store state.""" + + messages: list[ChatMessage] + + +class CustomChatMessageStore(ChatMessageStore): + """Implementation of custom chat message store. + In real applications, this can be an implementation of relational database or vector store.""" + + def __init__(self, messages: Collection[ChatMessage] | None = None) -> None: + self._messages: list[ChatMessage] = [] + if messages: + self._messages.extend(messages) + + async def add_messages(self, messages: Collection[ChatMessage]) -> None: + self._messages.extend(messages) + + async def list_messages(self) -> list[ChatMessage]: + return self._messages + + async def deserialize_state(self, serialized_store_state: Any, **kwargs: Any) -> None: + if serialized_store_state: + state = CustomStoreState.model_validate(serialized_store_state, **kwargs) + if state.messages: + self._messages.extend(state.messages) + + async def serialize_state(self, **kwargs: Any) -> Any: + state = CustomStoreState(messages=self._messages) + return state.model_dump(**kwargs) + + +async def main() -> None: + """Demonstrates how to use 3rd party or custom chat message store for threads.""" + print("=== Thread with 3rd party or custom chat message store ===") + + # OpenAI Chat Client is used as an example here, + # other chat clients can be used as well. + agent = OpenAIChatClient().create_agent( + name="Joker", + instructions="You are good at telling jokes.", + # Use custom chat message store. + # If not provided, the default in-memory store will be used. + chat_message_store_factory=CustomChatMessageStore, + ) + + # Start a new thread for the agent conversation. + thread = agent.get_new_thread() + + # Respond to user input. + query = "Tell me a joke about a pirate." + print(f"User: {query}") + print(f"Agent: {await agent.run(query, thread=thread)}\n") + + # Serialize the thread state, so it can be stored for later use. + serialized_thread = await thread.serialize() + + # The thread can now be saved to a database, file, or any other storage mechanism and loaded again later. + print(f"Serialized thread: {serialized_thread}\n") + + # Deserialize the thread state after loading from storage. + resumed_thread = await agent.deserialize_thread(serialized_thread) + + # Respond to user input. + query = "Now tell the same joke in the voice of a pirate, and add some emojis to the joke." + print(f"User: {query}") + print(f"Agent: {await agent.run(query, thread=resumed_thread)}\n") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/python/samples/getting_started/threads/suspend_resume_thread.py b/python/samples/getting_started/threads/suspend_resume_thread.py new file mode 100644 index 0000000000..cbb8a9b953 --- /dev/null +++ b/python/samples/getting_started/threads/suspend_resume_thread.py @@ -0,0 +1,83 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio + +from agent_framework.foundry import FoundryChatClient +from agent_framework.openai import OpenAIChatClient +from azure.identity.aio import DefaultAzureCredential + + +async def suspend_resume_service_managed_thread() -> None: + """Demonstrates how to suspend and resume a service-managed thread.""" + print("=== Suspend-Resume Service-Managed Thread ===") + + # Foundry Chat Client is used as an example here, + # other chat clients can be used as well. + async with ( + DefaultAzureCredential() as credential, + FoundryChatClient(async_ad_credential=credential).create_agent( + name="Joker", instructions="You are good at telling jokes." + ) as agent, + ): + # Start a new thread for the agent conversation. + thread = agent.get_new_thread() + + # Respond to user input. + query = "Tell me a joke about a pirate." + print(f"User: {query}") + print(f"Agent: {await agent.run(query, thread=thread)}\n") + + # Serialize the thread state, so it can be stored for later use. + serialized_thread = await thread.serialize() + + # The thread can now be saved to a database, file, or any other storage mechanism and loaded again later. + print(f"Serialized thread: {serialized_thread}\n") + + # Deserialize the thread state after loading from storage. + resumed_thread = await agent.deserialize_thread(serialized_thread) + + # Respond to user input. + query = "Now tell the same joke in the voice of a pirate, and add some emojis to the joke." + print(f"User: {query}") + print(f"Agent: {await agent.run(query, thread=resumed_thread)}\n") + + +async def suspend_resume_in_memory_thread() -> None: + """Demonstrates how to suspend and resume an in-memory thread.""" + print("=== Suspend-Resume In-Memory Thread ===") + + # OpenAI Chat Client is used as an example here, + # other chat clients can be used as well. + agent = OpenAIChatClient().create_agent(name="Joker", instructions="You are good at telling jokes.") + + # Start a new thread for the agent conversation. + thread = agent.get_new_thread() + + # Respond to user input. + query = "Tell me a joke about a pirate." + print(f"User: {query}") + print(f"Agent: {await agent.run(query, thread=thread)}\n") + + # Serialize the thread state, so it can be stored for later use. + serialized_thread = await thread.serialize() + + # The thread can now be saved to a database, file, or any other storage mechanism and loaded again later. + print(f"Serialized thread: {serialized_thread}\n") + + # Deserialize the thread state after loading from storage. + resumed_thread = await agent.deserialize_thread(serialized_thread) + + # Respond to user input. + query = "Now tell the same joke in the voice of a pirate, and add some emojis to the joke." + print(f"User: {query}") + print(f"Agent: {await agent.run(query, thread=resumed_thread)}\n") + + +async def main() -> None: + print("=== Suspend-Resume Thread Examples ===") + await suspend_resume_service_managed_thread() + await suspend_resume_in_memory_thread() + + +if __name__ == "__main__": + asyncio.run(main())