Python: Thread storage and serialization (#394)

* 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

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}")

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}")

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")
 
 

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}")

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}")

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}")

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")
 
 

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())

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())