agent/agent-framework
1
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/azure/pyproject.toml
+1 -1
View File
@@ -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"
python/packages/foundry/agent_framework_foundry/_chat_client.py
+8 -1
View File
@@ -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=[],
python/packages/main/agent_framework/__init__.py
+1
View File
@@ -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
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,
python/packages/main/agent_framework/_clients.py
+12 -1
View File
@@ -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
python/packages/main/agent_framework/_threads.py
+375
View File
@@ -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)
python/packages/main/agent_framework/telemetry.py
+4 -3
View File
@@ -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:
python/packages/main/tests/main/test_agents.py
+15 -91
View File
@@ -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]
python/packages/main/tests/main/test_threads.py
+482
View File
@@ -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
python/samples/getting_started/agents/azure_assistants_client/azure_assistants_with_thread.py
+3 -3
View File
@@ -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}")
python/samples/getting_started/agents/azure_chat_client/azure_chat_client_with_thread.py
+7 -5
View File
@@ -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}")
python/samples/getting_started/agents/azure_responses_client/azure_responses_client_with_thread.py
+8 -10
View File
@@ -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")
python/samples/getting_started/agents/foundry/foundry_with_thread.py
+3 -3
View File
@@ -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}")
python/samples/getting_started/agents/openai_assistants_client/openai_assistants_with_thread.py
+3 -3
View File
@@ -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}")
python/samples/getting_started/agents/openai_chat_client/openai_chat_client_with_thread.py
+7 -5
View File
@@ -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}")
python/samples/getting_started/agents/openai_responses_client/openai_responses_client_with_thread.py
+10 -10
View File
@@ -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")
python/samples/getting_started/threads/custom_chat_message_store_thread.py
+82
View File
@@ -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())
python/samples/getting_started/threads/suspend_resume_thread.py
+83
View File
@@ -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())