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

Python: Thread storage and serialization (#394)

Browse Source 
* Updates for message store support

* Added unit tests

* Added suspend-resume example

* Added example with custom chat message store

* Small fix

* Addressed PR feedback

* Renaming and documentation

* More renaming

* Addressed more PR feedback

* Small fixes in Foundry chat client and examples

* Small update

* Addressed PR feedback

* Increased timeout for Azure tests
This commit is contained in:
Dmytro Struk
2025-08-14 14:15:42 -07:00
committed by GitHub
Unverified
parent 95cb20ca40
commit dea736e550
18 changed files with 1151 additions and 316 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/main/agent_framework/_agents.py
+47 -180
View File
@@ -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,