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/packages/main/tests/main/test_agents.py

@@ -1,6 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-from collections.abc import AsyncIterable, MutableSequence, Sequence
+from collections.abc import AsyncIterable, MutableSequence
 from typing import Any
 from uuid import uuid4
 
@@ -13,10 +13,9 @@ from agent_framework import (
     AIAgent,
     ChatClient,
     ChatClientAgent,
-    ChatClientAgentThread,
-    ChatClientAgentThreadType,
     ChatClientBase,
     ChatMessage,
+    ChatMessageList,
     ChatOptions,
     ChatResponse,
     ChatResponseUpdate,
@@ -28,8 +27,7 @@ from agent_framework.exceptions import AgentExecutionException
 
 # Mock AgentThread implementation for testing
 class MockAgentThread(AgentThread):
-    async def _on_new_messages(self, new_messages: ChatMessage | Sequence[ChatMessage]) -> None:
-        pass
+    pass
 
 
 # Mock Agent implementation for testing
@@ -142,58 +140,6 @@ async def test_agent_run_streaming(agent: AIAgent) -> None:
     assert updates[0].text == "Response"
 
 
-async def test_chat_client_agent_thread_init_in_memory() -> None:
-    messages = [ChatMessage(role=ChatRole.USER, contents=[TextContent("Hello")])]
-    thread = ChatClientAgentThread(messages=messages)
-
-    assert thread.storage_location == ChatClientAgentThreadType.IN_MEMORY_MESSAGES
-    assert thread.id is None
-    assert thread.chat_messages == messages
-
-
-async def test_chat_client_agent_thread_empty() -> None:
-    thread = ChatClientAgentThread()
-
-    assert thread.storage_location is None
-    assert thread.id is None
-    assert thread.chat_messages is None
-
-
-async def test_chat_client_agent_thread_init_invalid() -> None:
-    with raises(ValueError, match="Cannot specify both id and messages"):
-        ChatClientAgentThread(id="123", messages=[ChatMessage(role=ChatRole.USER, contents=[TextContent("Hello")])])
-
-    with raises(ValueError, match="ID cannot be empty or whitespace"):
-        ChatClientAgentThread(id=" ")
-
-
-async def test_chat_client_agent_thread_init_conversation_id() -> None:
-    thread_id = str(uuid4())
-    thread = ChatClientAgentThread(id=thread_id)
-
-    assert thread.storage_location == ChatClientAgentThreadType.CONVERSATION_ID
-    assert thread.id == thread_id
-    assert thread.chat_messages is None
-
-
-async def test_chat_client_agent_thread_get_messages() -> None:
-    messages = [ChatMessage(role=ChatRole.USER, contents=[TextContent("Hello")])]
-    thread = ChatClientAgentThread(messages=messages)
-
-    result = [msg async for msg in thread.get_messages()]
-    assert result == messages
-
-
-async def test_chat_client_agent_thread_on_new_messages_in_memory() -> None:
-    initial_message = ChatMessage(role=ChatRole.USER, contents=[TextContent("Initial message")])
-    new_message = ChatMessage(role=ChatRole.USER, contents=[TextContent("New message")])
-
-    thread = ChatClientAgentThread(messages=[initial_message])
-
-    await thread._on_new_messages(new_message)  # type: ignore[reportPrivateUsage]
-    assert thread.chat_messages == [initial_message, new_message]
-
-
 def test_chat_client_agent_type(chat_client: ChatClient) -> None:
     chat_client_agent = ChatClientAgent(chat_client=chat_client)
     assert isinstance(chat_client_agent, AIAgent)
@@ -239,24 +185,16 @@ async def test_chat_client_agent_get_new_thread(chat_client: ChatClient) -> None
     agent = ChatClientAgent(chat_client=chat_client)
     thread = agent.get_new_thread()
 
-    assert isinstance(thread, ChatClientAgentThread)
-    assert thread.storage_location is None
+    assert isinstance(thread, AgentThread)
 
 
 async def test_chat_client_agent_prepare_thread_and_messages(chat_client: ChatClient) -> None:
     agent = ChatClientAgent(chat_client=chat_client)
     message = ChatMessage(role=ChatRole.USER, text="Hello")
-    thread = ChatClientAgentThread(messages=[message])
-
-    result_thread = agent._validate_or_create_thread_type(  # type: ignore[reportPrivateUsage]
-        thread, lambda: ChatClientAgentThread(), expected_type=ChatClientAgentThread
-    )  # type: ignore[reportPrivateUsage]
-
-    assert result_thread == thread
-    assert isinstance(result_thread, ChatClientAgentThread)
+    thread = AgentThread(message_store=ChatMessageList(messages=[message]))
 
     _, result_messages = await agent._prepare_thread_and_messages(  # type: ignore[reportPrivateUsage]
-        thread=result_thread,
+        thread=thread,
         input_messages=[ChatMessage(role=ChatRole.USER, text="Test")],
     )
 
@@ -265,18 +203,6 @@ async def test_chat_client_agent_prepare_thread_and_messages(chat_client: ChatCl
     assert result_messages[1].text == "Test"
 
 
-async def test_chat_client_agent_validate_or_create_thread(chat_client: ChatClient) -> None:
-    agent = ChatClientAgent(chat_client=chat_client)
-    thread = None
-
-    result_thread = agent._validate_or_create_thread_type(  # type: ignore[reportPrivateUsage]
-        thread, lambda: ChatClientAgentThread(), expected_type=ChatClientAgentThread
-    )  # type: ignore[reportPrivateUsage]
-
-    assert result_thread != thread
-    assert isinstance(result_thread, ChatClientAgentThread)
-
-
 async def test_chat_client_agent_update_thread_id() -> None:
     chat_client = MockChatClient(
         mock_response=ChatResponse(
@@ -290,9 +216,7 @@ async def test_chat_client_agent_update_thread_id() -> None:
     result = await agent.run("Hello", thread=thread)
     assert result.text == "test response"
 
-    assert thread.id == "123"
-    assert isinstance(thread, ChatClientAgentThread)
-    assert thread.storage_location == ChatClientAgentThreadType.CONVERSATION_ID
+    assert thread.service_thread_id == "123"
 
 
 async def test_chat_client_agent_update_thread_messages(chat_client: ChatClient) -> None:
@@ -302,19 +226,19 @@ async def test_chat_client_agent_update_thread_messages(chat_client: ChatClient)
     result = await agent.run("Hello", thread=thread)
     assert result.text == "test response"
 
-    assert thread.id is None
-    assert isinstance(thread, ChatClientAgentThread)
-    assert thread.storage_location == ChatClientAgentThreadType.IN_MEMORY_MESSAGES
+    assert thread.service_thread_id is None
 
-    assert thread.chat_messages is not None
-    assert len(thread.chat_messages) == 2
-    assert thread.chat_messages[0].text == "Hello"
-    assert thread.chat_messages[1].text == "test response"
+    chat_messages: list[ChatMessage] | None = await thread.list_messages()
+
+    assert chat_messages is not None
+    assert len(chat_messages) == 2
+    assert chat_messages[0].text == "Hello"
+    assert chat_messages[1].text == "test response"
 
 
 async def test_chat_client_agent_update_thread_conversation_id_missing(chat_client: ChatClient) -> None:
     agent = ChatClientAgent(chat_client=chat_client)
-    thread = ChatClientAgentThread(id="123")
+    thread = AgentThread(service_thread_id="123")
 
     with raises(AgentExecutionException, match="Service did not return a valid conversation id"):
         agent._update_thread_with_type_and_conversation_id(thread, None)  # type: ignore[reportPrivateUsage]