mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: [BREAKING] cleanup of thread API and serialization (#893)
* cleanup of threads and serialization * fix for sliding window * fix redis test * updated from comments * updated context provider and threads * updated lock * add asyncio default * fix redis tests * fix tests * fix tests * renamed to invoking * fixed tests * fix for instructions
This commit is contained in:
Eduard van Valkenburg
committed by
GitHub
Unverified
GitHub
parent
bf5931932e
commit
10d10364a9
@@ -141,7 +141,7 @@ class MockBaseChatClient(BaseChatClient):
|
||||
logger.debug(f"Running base chat client inner, with: {messages=}, {chat_options=}, {kwargs=}")
|
||||
self.call_count += 1
|
||||
if not self.run_responses:
|
||||
return ChatResponse(messages=ChatMessage(role="assistant", text=f"test response - {messages[0].text}"))
|
||||
return ChatResponse(messages=ChatMessage(role="assistant", text=f"test response - {messages[-1].text}"))
|
||||
|
||||
response = self.run_responses.pop(0)
|
||||
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from collections.abc import AsyncIterable, MutableSequence, Sequence
|
||||
from typing import Any
|
||||
from uuid import uuid4
|
||||
|
||||
from pytest import raises
|
||||
@@ -10,17 +11,18 @@ from agent_framework import (
|
||||
AgentRunResponse,
|
||||
AgentRunResponseUpdate,
|
||||
AgentThread,
|
||||
AggregateContextProvider,
|
||||
ChatAgent,
|
||||
ChatClientProtocol,
|
||||
ChatMessage,
|
||||
ChatMessageList,
|
||||
ChatMessageStore,
|
||||
ChatResponse,
|
||||
Contents,
|
||||
Context,
|
||||
ContextProvider,
|
||||
HostedCodeInterpreterTool,
|
||||
Role,
|
||||
TextContent,
|
||||
)
|
||||
from agent_framework._memory import AggregateContextProvider, Context, ContextProvider
|
||||
from agent_framework.exceptions import AgentExecutionException
|
||||
|
||||
|
||||
@@ -98,11 +100,10 @@ async def test_chat_client_agent_get_new_thread(chat_client: ChatClientProtocol)
|
||||
async def test_chat_client_agent_prepare_thread_and_messages(chat_client: ChatClientProtocol) -> None:
|
||||
agent = ChatAgent(chat_client=chat_client)
|
||||
message = ChatMessage(role=Role.USER, text="Hello")
|
||||
thread = AgentThread(message_store=ChatMessageList(messages=[message]))
|
||||
thread = AgentThread(message_store=ChatMessageStore(messages=[message]))
|
||||
|
||||
_, result_messages = await agent._prepare_thread_and_messages( # type: ignore[reportPrivateUsage]
|
||||
_, _, result_messages = await agent._prepare_thread_and_messages( # type: ignore[reportPrivateUsage]
|
||||
thread=thread,
|
||||
context=Context(),
|
||||
input_messages=[ChatMessage(role=Role.USER, text="Test")],
|
||||
)
|
||||
|
||||
@@ -152,7 +153,7 @@ async def test_chat_client_agent_update_thread_conversation_id_missing(chat_clie
|
||||
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]
|
||||
await agent._update_thread_with_type_and_conversation_id(thread, None) # type: ignore[reportPrivateUsage]
|
||||
|
||||
|
||||
async def test_chat_client_agent_default_author_name(chat_client: ChatClientProtocol) -> None:
|
||||
@@ -191,49 +192,49 @@ async def test_chat_client_agent_author_name_is_used_from_response(chat_client_b
|
||||
|
||||
# Mock context provider for testing
|
||||
class MockContextProvider(ContextProvider):
|
||||
context_contents: list[Contents] | None = None
|
||||
thread_created_called: bool = False
|
||||
messages_adding_called: bool = False
|
||||
model_invoking_called: bool = False
|
||||
thread_created_thread_id: str | None = None
|
||||
messages_adding_thread_id: str | None = None
|
||||
new_messages: list[ChatMessage] = []
|
||||
|
||||
def __init__(self, contents: list[Contents] | None = None) -> None:
|
||||
super().__init__()
|
||||
self.context_contents = contents
|
||||
def __init__(self, messages: list[ChatMessage] | None = None) -> None:
|
||||
self.context_messages = messages
|
||||
self.thread_created_called = False
|
||||
self.messages_adding_called = False
|
||||
self.model_invoking_called = False
|
||||
self.invoked_called = False
|
||||
self.invoking_called = False
|
||||
self.thread_created_thread_id = None
|
||||
self.messages_adding_thread_id = None
|
||||
self.new_messages = []
|
||||
self.invoked_thread_id = None
|
||||
self.new_messages: list[ChatMessage] = []
|
||||
|
||||
async def thread_created(self, thread_id: str | None) -> None:
|
||||
self.thread_created_called = True
|
||||
self.thread_created_thread_id = thread_id
|
||||
|
||||
async def messages_adding(self, thread_id: str | None, new_messages: ChatMessage | Sequence[ChatMessage]) -> None:
|
||||
self.messages_adding_called = True
|
||||
self.messages_adding_thread_id = thread_id
|
||||
if isinstance(new_messages, ChatMessage):
|
||||
self.new_messages.append(new_messages)
|
||||
async def invoked(
|
||||
self,
|
||||
request_messages: ChatMessage | Sequence[ChatMessage],
|
||||
response_messages: ChatMessage | Sequence[ChatMessage] | None = None,
|
||||
invoke_exception: Any = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
self.invoked_called = True
|
||||
if isinstance(request_messages, ChatMessage):
|
||||
self.new_messages.append(request_messages)
|
||||
else:
|
||||
self.new_messages.extend(new_messages)
|
||||
self.new_messages.extend(request_messages)
|
||||
if isinstance(response_messages, ChatMessage):
|
||||
self.new_messages.append(response_messages)
|
||||
else:
|
||||
self.new_messages.extend(response_messages)
|
||||
|
||||
async def model_invoking(self, messages: ChatMessage | MutableSequence[ChatMessage]) -> Context:
|
||||
self.model_invoking_called = True
|
||||
return Context(contents=self.context_contents)
|
||||
async def invoking(self, messages: ChatMessage | MutableSequence[ChatMessage], **kwargs: Any) -> Context:
|
||||
self.invoking_called = True
|
||||
return Context(messages=self.context_messages)
|
||||
|
||||
|
||||
async def test_chat_agent_context_providers_model_invoking(chat_client: ChatClientProtocol) -> None:
|
||||
"""Test that context providers' model_invoking is called during agent run."""
|
||||
mock_provider = MockContextProvider(contents=[TextContent("Test context instructions")])
|
||||
"""Test that context providers' invoking is called during agent run."""
|
||||
mock_provider = MockContextProvider(messages=[ChatMessage(role=Role.SYSTEM, text="Test context instructions")])
|
||||
agent = ChatAgent(chat_client=chat_client, context_providers=mock_provider)
|
||||
|
||||
await agent.run("Hello")
|
||||
|
||||
assert mock_provider.model_invoking_called
|
||||
assert mock_provider.invoking_called
|
||||
|
||||
|
||||
async def test_chat_agent_context_providers_thread_created(chat_client_base: ChatClientProtocol) -> None:
|
||||
@@ -255,75 +256,54 @@ async def test_chat_agent_context_providers_thread_created(chat_client_base: Cha
|
||||
|
||||
|
||||
async def test_chat_agent_context_providers_messages_adding(chat_client: ChatClientProtocol) -> None:
|
||||
"""Test that context providers' messages_adding is called during agent run."""
|
||||
"""Test that context providers' invoked is called during agent run."""
|
||||
mock_provider = MockContextProvider()
|
||||
agent = ChatAgent(chat_client=chat_client, context_providers=mock_provider)
|
||||
|
||||
await agent.run("Hello")
|
||||
|
||||
assert mock_provider.messages_adding_called
|
||||
assert mock_provider.invoked_called
|
||||
# Should be called with both input and response messages
|
||||
assert len(mock_provider.new_messages) >= 2
|
||||
|
||||
|
||||
async def test_chat_agent_context_instructions_in_messages(chat_client: ChatClientProtocol) -> None:
|
||||
"""Test that AI context instructions are included in messages."""
|
||||
mock_provider = MockContextProvider(contents=[TextContent("Context-specific instructions")])
|
||||
mock_provider = MockContextProvider(messages=[ChatMessage(role="system", text="Context-specific instructions")])
|
||||
agent = ChatAgent(chat_client=chat_client, instructions="Agent instructions", context_providers=mock_provider)
|
||||
|
||||
# We need to test the _prepare_thread_and_messages method directly
|
||||
context = Context(contents=[TextContent("Context-specific instructions")])
|
||||
_, messages = await agent._prepare_thread_and_messages( # type: ignore[reportPrivateUsage]
|
||||
thread=None, context=context, input_messages=[ChatMessage(role=Role.USER, text="Hello")]
|
||||
_, _, messages = await agent._prepare_thread_and_messages( # type: ignore[reportPrivateUsage]
|
||||
thread=None, input_messages=[ChatMessage(role=Role.USER, text="Hello")]
|
||||
)
|
||||
|
||||
# Should have agent instructions, context instructions, and user message
|
||||
assert len(messages) == 3
|
||||
assert messages[0].role == Role.SYSTEM
|
||||
assert messages[0].text == "Agent instructions"
|
||||
assert messages[1].role == Role.SYSTEM
|
||||
assert messages[1].text == "Context-specific instructions"
|
||||
assert messages[2].role == Role.USER
|
||||
assert messages[2].text == "Hello"
|
||||
|
||||
|
||||
async def test_chat_agent_context_instructions_without_agent_instructions(chat_client: ChatClientProtocol) -> None:
|
||||
"""Test that AI context instructions work when agent has no instructions."""
|
||||
agent = ChatAgent(chat_client=chat_client) # No instructions
|
||||
context = Context(contents=[TextContent("Context-only instructions")])
|
||||
|
||||
_, messages = await agent._prepare_thread_and_messages( # type: ignore[reportPrivateUsage]
|
||||
thread=None, context=context, input_messages=[ChatMessage(role=Role.USER, text="Hello")]
|
||||
)
|
||||
|
||||
# Should have context instructions and user message only
|
||||
# Should have context instructions, and user message
|
||||
assert len(messages) == 2
|
||||
assert messages[0].role == Role.SYSTEM
|
||||
assert messages[0].text == "Context-only instructions"
|
||||
assert messages[0].text == "Context-specific instructions"
|
||||
assert messages[1].role == Role.USER
|
||||
assert messages[1].text == "Hello"
|
||||
# instructions system message is added by a chat_client
|
||||
|
||||
|
||||
async def test_chat_agent_no_context_instructions(chat_client: ChatClientProtocol) -> None:
|
||||
"""Test behavior when AI context has no instructions."""
|
||||
agent = ChatAgent(chat_client=chat_client, instructions="Agent instructions")
|
||||
context = Context() # No instructions
|
||||
mock_provider = MockContextProvider()
|
||||
agent = ChatAgent(chat_client=chat_client, instructions="Agent instructions", context_providers=mock_provider)
|
||||
|
||||
_, messages = await agent._prepare_thread_and_messages( # type: ignore[reportPrivateUsage]
|
||||
thread=None, context=context, input_messages=[ChatMessage(role=Role.USER, text="Hello")]
|
||||
_, _, messages = await agent._prepare_thread_and_messages( # type: ignore[reportPrivateUsage]
|
||||
thread=None, input_messages=[ChatMessage(role=Role.USER, text="Hello")]
|
||||
)
|
||||
|
||||
# Should have agent instructions and user message only
|
||||
assert len(messages) == 2
|
||||
assert messages[0].role == Role.SYSTEM
|
||||
assert messages[0].text == "Agent instructions"
|
||||
assert messages[1].role == Role.USER
|
||||
assert messages[1].text == "Hello"
|
||||
assert len(messages) == 1
|
||||
assert messages[0].role == Role.USER
|
||||
assert messages[0].text == "Hello"
|
||||
|
||||
|
||||
async def test_chat_agent_run_stream_context_providers(chat_client: ChatClientProtocol) -> None:
|
||||
"""Test that context providers work with run_stream method."""
|
||||
mock_provider = MockContextProvider(contents=[TextContent("Stream context instructions")])
|
||||
mock_provider = MockContextProvider(messages=[ChatMessage(role=Role.SYSTEM, text="Stream context instructions")])
|
||||
agent = ChatAgent(chat_client=chat_client, context_providers=mock_provider)
|
||||
|
||||
# Collect all stream updates
|
||||
@@ -332,47 +312,48 @@ async def test_chat_agent_run_stream_context_providers(chat_client: ChatClientPr
|
||||
updates.append(update)
|
||||
|
||||
# Verify context provider was called
|
||||
assert mock_provider.model_invoking_called
|
||||
assert mock_provider.thread_created_called
|
||||
assert mock_provider.messages_adding_called
|
||||
assert mock_provider.invoking_called
|
||||
# no conversation id is created, so no need to thread_create to be called.
|
||||
assert not mock_provider.thread_created_called
|
||||
assert mock_provider.invoked_called
|
||||
|
||||
|
||||
async def test_chat_agent_multiple_context_providers(chat_client: ChatClientProtocol) -> None:
|
||||
"""Test that multiple context providers work together."""
|
||||
provider1 = MockContextProvider(contents=[TextContent("First provider instructions")])
|
||||
provider2 = MockContextProvider(contents=[TextContent("Second provider instructions")])
|
||||
provider1 = MockContextProvider(messages=[ChatMessage(role=Role.SYSTEM, text="First provider instructions")])
|
||||
provider2 = MockContextProvider(messages=[ChatMessage(role=Role.SYSTEM, text="Second provider instructions")])
|
||||
|
||||
agent = ChatAgent(chat_client=chat_client, context_providers=[provider1, provider2])
|
||||
|
||||
await agent.run("Hello")
|
||||
|
||||
# Both providers should be called
|
||||
assert provider1.model_invoking_called
|
||||
assert provider1.thread_created_called
|
||||
assert provider1.messages_adding_called
|
||||
assert provider1.invoking_called
|
||||
assert not provider1.thread_created_called
|
||||
assert provider1.invoked_called
|
||||
|
||||
assert provider2.model_invoking_called
|
||||
assert provider2.thread_created_called
|
||||
assert provider2.messages_adding_called
|
||||
assert provider2.invoking_called
|
||||
assert not provider2.thread_created_called
|
||||
assert provider2.invoked_called
|
||||
|
||||
|
||||
async def test_chat_agent_aggregate_context_provider_combines_instructions() -> None:
|
||||
"""Test that AggregateContextProvider combines instructions from multiple providers."""
|
||||
provider1 = MockContextProvider(contents=[TextContent("First instruction")])
|
||||
provider2 = MockContextProvider(contents=[TextContent("Second instruction")])
|
||||
provider1 = MockContextProvider(messages=[ChatMessage(role=Role.SYSTEM, text="First instruction")])
|
||||
provider2 = MockContextProvider(messages=[ChatMessage(role=Role.SYSTEM, text="Second instruction")])
|
||||
|
||||
aggregate = AggregateContextProvider()
|
||||
aggregate.providers.append(provider1)
|
||||
aggregate.providers.append(provider2)
|
||||
|
||||
# Test model_invoking combines instructions
|
||||
result = await aggregate.model_invoking([ChatMessage(role=Role.USER, text="Test")])
|
||||
# Test invoking combines instructions
|
||||
result = await aggregate.invoking([ChatMessage(role=Role.USER, text="Test")])
|
||||
|
||||
assert result.contents
|
||||
assert isinstance(result.contents[0], TextContent)
|
||||
assert isinstance(result.contents[1], TextContent)
|
||||
assert result.contents[0].text == "First instruction"
|
||||
assert result.contents[1].text == "Second instruction"
|
||||
assert result.messages
|
||||
assert isinstance(result.messages[0], ChatMessage)
|
||||
assert isinstance(result.messages[1], ChatMessage)
|
||||
assert result.messages[0].text == "First instruction"
|
||||
assert result.messages[1].text == "Second instruction"
|
||||
|
||||
|
||||
async def test_chat_agent_context_providers_with_thread_service_id(chat_client_base: ChatClientProtocol) -> None:
|
||||
@@ -388,12 +369,11 @@ async def test_chat_agent_context_providers_with_thread_service_id(chat_client_b
|
||||
agent = ChatAgent(chat_client=chat_client_base, context_providers=mock_provider)
|
||||
|
||||
# Use existing service-managed thread
|
||||
thread = AgentThread(service_thread_id="existing-thread-id")
|
||||
thread = agent.get_new_thread(service_thread_id="existing-thread-id")
|
||||
await agent.run("Hello", thread=thread)
|
||||
|
||||
# messages_adding should be called with the service thread ID from response
|
||||
assert mock_provider.messages_adding_called
|
||||
assert mock_provider.messages_adding_thread_id == "service-thread-123" # Updated thread ID from response
|
||||
# invoked should be called with the service thread ID from response
|
||||
assert mock_provider.invoked_called
|
||||
|
||||
|
||||
# Tests for as_tool method
|
||||
|
||||
@@ -1,33 +1,23 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from collections.abc import MutableSequence, Sequence
|
||||
from collections.abc import MutableSequence
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
from agent_framework import ChatMessage, Contents, Role, TextContent
|
||||
from agent_framework import ChatMessage, Role, TextContent
|
||||
from agent_framework._memory import AggregateContextProvider, Context, ContextProvider
|
||||
|
||||
|
||||
class MockContextProvider(ContextProvider):
|
||||
"""Mock ContextProvider for testing."""
|
||||
|
||||
context_contents: list[Contents] | None = None
|
||||
thread_created_called: bool = False
|
||||
messages_adding_called: bool = False
|
||||
model_invoking_called: bool = False
|
||||
thread_created_thread_id: str | None = None
|
||||
messages_adding_thread_id: str | None = None
|
||||
messages_adding_new_messages: ChatMessage | Sequence[ChatMessage] | None = None
|
||||
model_invoking_messages: ChatMessage | MutableSequence[ChatMessage] | None = None
|
||||
|
||||
def __init__(self, context_contents: list[Contents] | None = None) -> None:
|
||||
super().__init__()
|
||||
self.context_contents = context_contents
|
||||
def __init__(self, messages: list[ChatMessage] | None = None) -> None:
|
||||
self.context_messages = messages
|
||||
self.thread_created_called = False
|
||||
self.messages_adding_called = False
|
||||
self.model_invoking_called = False
|
||||
self.invoked_called = False
|
||||
self.invoking_called = False
|
||||
self.thread_created_thread_id = None
|
||||
self.messages_adding_thread_id = None
|
||||
self.messages_adding_new_messages = None
|
||||
self.new_messages = None
|
||||
self.model_invoking_messages = None
|
||||
|
||||
async def thread_created(self, thread_id: str | None) -> None:
|
||||
@@ -35,18 +25,23 @@ class MockContextProvider(ContextProvider):
|
||||
self.thread_created_called = True
|
||||
self.thread_created_thread_id = thread_id
|
||||
|
||||
async def messages_adding(self, thread_id: str | None, new_messages: ChatMessage | Sequence[ChatMessage]) -> None:
|
||||
"""Track messages_adding calls."""
|
||||
self.messages_adding_called = True
|
||||
self.messages_adding_thread_id = thread_id
|
||||
self.messages_adding_new_messages = new_messages
|
||||
async def invoked(
|
||||
self,
|
||||
request_messages: Any,
|
||||
response_messages: Any | None = None,
|
||||
invoke_exception: Exception | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Track invoked calls."""
|
||||
self.invoked_called = True
|
||||
self.new_messages = request_messages
|
||||
|
||||
async def model_invoking(self, messages: ChatMessage | MutableSequence[ChatMessage]) -> Context:
|
||||
"""Track model_invoking calls and return context."""
|
||||
self.model_invoking_called = True
|
||||
async def invoking(self, messages: ChatMessage | MutableSequence[ChatMessage], **kwargs: Any) -> Context:
|
||||
"""Track invoking calls and return context."""
|
||||
self.invoking_called = True
|
||||
self.model_invoking_messages = messages
|
||||
context = Context()
|
||||
context.contents = self.context_contents
|
||||
context.messages = self.context_messages
|
||||
return context
|
||||
|
||||
|
||||
@@ -65,19 +60,21 @@ class TestAggregateContextProvider:
|
||||
|
||||
def test_init_with_providers(self) -> None:
|
||||
"""Test initialization with providers."""
|
||||
provider1 = MockContextProvider([TextContent("instructions1")])
|
||||
provider2 = MockContextProvider([TextContent("instructions1")])
|
||||
providers = [provider1, provider2]
|
||||
provider1 = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions 1")])
|
||||
provider2 = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions 2")])
|
||||
provider3 = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions 3")])
|
||||
providers = [provider1, provider2, provider3]
|
||||
|
||||
aggregate = AggregateContextProvider(providers)
|
||||
assert len(aggregate.providers) == 2
|
||||
assert len(aggregate.providers) == 3
|
||||
assert aggregate.providers[0] is provider1
|
||||
assert aggregate.providers[1] is provider2
|
||||
assert aggregate.providers[2] is provider3
|
||||
|
||||
def test_add_provider(self) -> None:
|
||||
"""Test adding a provider."""
|
||||
aggregate = AggregateContextProvider()
|
||||
provider = MockContextProvider([TextContent("instructions")])
|
||||
provider = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions")])
|
||||
|
||||
aggregate.add(provider)
|
||||
assert len(aggregate.providers) == 1
|
||||
@@ -86,8 +83,8 @@ class TestAggregateContextProvider:
|
||||
def test_add_multiple_providers(self) -> None:
|
||||
"""Test adding multiple providers."""
|
||||
aggregate = AggregateContextProvider()
|
||||
provider1 = MockContextProvider([TextContent("instructions1")])
|
||||
provider2 = MockContextProvider([TextContent("instructions2")])
|
||||
provider1 = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions 1")])
|
||||
provider2 = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions 2")])
|
||||
|
||||
aggregate.add(provider1)
|
||||
aggregate.add(provider2)
|
||||
@@ -105,8 +102,8 @@ class TestAggregateContextProvider:
|
||||
|
||||
async def test_thread_created_with_providers(self) -> None:
|
||||
"""Test thread_created calls all providers."""
|
||||
provider1 = MockContextProvider([TextContent("instructions1")])
|
||||
provider2 = MockContextProvider([TextContent("instructions2")])
|
||||
provider1 = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions 1")])
|
||||
provider2 = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions 2")])
|
||||
aggregate = AggregateContextProvider([provider1, provider2])
|
||||
|
||||
thread_id = "thread-123"
|
||||
@@ -119,7 +116,7 @@ class TestAggregateContextProvider:
|
||||
|
||||
async def test_thread_created_with_none_thread_id(self) -> None:
|
||||
"""Test thread_created with None thread_id."""
|
||||
provider = MockContextProvider([TextContent("instructions")])
|
||||
provider = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions")])
|
||||
aggregate = AggregateContextProvider([provider])
|
||||
|
||||
await aggregate.thread_created(None)
|
||||
@@ -128,155 +125,150 @@ class TestAggregateContextProvider:
|
||||
assert provider.thread_created_thread_id is None
|
||||
|
||||
async def test_messages_adding_with_no_providers(self) -> None:
|
||||
"""Test messages_adding with no providers."""
|
||||
"""Test invoked with no providers."""
|
||||
aggregate = AggregateContextProvider()
|
||||
message = ChatMessage(text="Hello", role=Role.USER)
|
||||
|
||||
# Should not raise an exception
|
||||
await aggregate.messages_adding("thread-123", message)
|
||||
await aggregate.invoked(message)
|
||||
|
||||
async def test_messages_adding_with_single_message(self) -> None:
|
||||
"""Test messages_adding with a single message."""
|
||||
provider1 = MockContextProvider([TextContent("instructions1")])
|
||||
provider2 = MockContextProvider([TextContent("instructions2")])
|
||||
"""Test invoked with a single message."""
|
||||
provider1 = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions 1")])
|
||||
provider2 = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions 2")])
|
||||
aggregate = AggregateContextProvider([provider1, provider2])
|
||||
|
||||
thread_id = "thread-123"
|
||||
message = ChatMessage(text="Hello", role=Role.USER)
|
||||
await aggregate.messages_adding(thread_id, message)
|
||||
await aggregate.invoked(message)
|
||||
|
||||
assert provider1.messages_adding_called
|
||||
assert provider1.messages_adding_thread_id == thread_id
|
||||
assert provider1.messages_adding_new_messages == message
|
||||
assert provider2.messages_adding_called
|
||||
assert provider2.messages_adding_thread_id == thread_id
|
||||
assert provider2.messages_adding_new_messages == message
|
||||
assert provider1.invoked_called
|
||||
assert provider1.new_messages == message
|
||||
assert provider2.invoked_called
|
||||
assert provider2.new_messages == message
|
||||
|
||||
async def test_messages_adding_with_message_sequence(self) -> None:
|
||||
"""Test messages_adding with a sequence of messages."""
|
||||
provider = MockContextProvider([TextContent("instructions")])
|
||||
"""Test invoked with a sequence of messages."""
|
||||
provider = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions")])
|
||||
aggregate = AggregateContextProvider([provider])
|
||||
|
||||
thread_id = "thread-123"
|
||||
messages = [
|
||||
ChatMessage(text="Hello", role=Role.USER),
|
||||
ChatMessage(text="Hi there", role=Role.ASSISTANT),
|
||||
]
|
||||
await aggregate.messages_adding(thread_id, messages)
|
||||
await aggregate.invoked(messages)
|
||||
|
||||
assert provider.messages_adding_called
|
||||
assert provider.messages_adding_thread_id == thread_id
|
||||
assert provider.messages_adding_new_messages == messages
|
||||
assert provider.invoked_called
|
||||
assert provider.new_messages == messages
|
||||
|
||||
async def test_model_invoking_with_no_providers(self) -> None:
|
||||
"""Test model_invoking with no providers."""
|
||||
"""Test invoking with no providers."""
|
||||
aggregate = AggregateContextProvider()
|
||||
message = ChatMessage(text="Hello", role=Role.USER)
|
||||
|
||||
context = await aggregate.model_invoking(message)
|
||||
context = await aggregate.invoking(message)
|
||||
|
||||
assert isinstance(context, Context)
|
||||
assert not context.contents
|
||||
assert not context.messages
|
||||
|
||||
async def test_model_invoking_with_single_provider(self) -> None:
|
||||
"""Test model_invoking with a single provider."""
|
||||
provider = MockContextProvider([TextContent("Test instructions")])
|
||||
"""Test invoking with a single provider."""
|
||||
provider = MockContextProvider(messages=[ChatMessage(role="user", text="Test instructions")])
|
||||
aggregate = AggregateContextProvider([provider])
|
||||
|
||||
message = ChatMessage(text="Hello", role=Role.USER)
|
||||
context = await aggregate.model_invoking(message)
|
||||
message = [ChatMessage(text="Hello", role=Role.USER)]
|
||||
context = await aggregate.invoking(message)
|
||||
|
||||
assert provider.model_invoking_called
|
||||
assert provider.invoking_called
|
||||
assert provider.model_invoking_messages == message
|
||||
assert isinstance(context, Context)
|
||||
|
||||
assert context.contents
|
||||
assert isinstance(context.contents[0], TextContent)
|
||||
assert context.contents[0].text == "Test instructions"
|
||||
assert context.messages
|
||||
assert isinstance(context.messages[0].contents[0], TextContent)
|
||||
assert context.messages[0].text == "Test instructions"
|
||||
|
||||
async def test_model_invoking_with_multiple_providers(self) -> None:
|
||||
"""Test model_invoking combines contexts from multiple providers."""
|
||||
provider1 = MockContextProvider([TextContent("Instructions 1")])
|
||||
provider2 = MockContextProvider([TextContent("Instructions 2")])
|
||||
provider3 = MockContextProvider([TextContent("Instructions 3")])
|
||||
"""Test invoking combines contexts from multiple providers."""
|
||||
provider1 = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions 1")])
|
||||
provider2 = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions 2")])
|
||||
provider3 = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions 3")])
|
||||
aggregate = AggregateContextProvider([provider1, provider2, provider3])
|
||||
|
||||
messages = [ChatMessage(text="Hello", role=Role.USER)]
|
||||
context = await aggregate.model_invoking(messages)
|
||||
context = await aggregate.invoking(messages)
|
||||
|
||||
assert provider1.model_invoking_called
|
||||
assert provider1.invoking_called
|
||||
assert provider1.model_invoking_messages == messages
|
||||
assert provider2.model_invoking_called
|
||||
assert provider2.invoking_called
|
||||
assert provider2.model_invoking_messages == messages
|
||||
assert provider3.model_invoking_called
|
||||
assert provider3.invoking_called
|
||||
assert provider3.model_invoking_messages == messages
|
||||
|
||||
assert isinstance(context, Context)
|
||||
|
||||
assert context.contents
|
||||
assert isinstance(context.contents[0], TextContent)
|
||||
assert isinstance(context.contents[1], TextContent)
|
||||
assert isinstance(context.contents[2], TextContent)
|
||||
assert context.contents[0].text == "Instructions 1"
|
||||
assert context.contents[1].text == "Instructions 2"
|
||||
assert context.contents[2].text == "Instructions 3"
|
||||
assert context.messages
|
||||
assert isinstance(context.messages[0].contents[0], TextContent)
|
||||
assert isinstance(context.messages[1].contents[0], TextContent)
|
||||
assert isinstance(context.messages[2].contents[0], TextContent)
|
||||
assert context.messages[0].text == "Instructions 1"
|
||||
assert context.messages[1].text == "Instructions 2"
|
||||
assert context.messages[2].text == "Instructions 3"
|
||||
|
||||
async def test_model_invoking_with_none_instructions(self) -> None:
|
||||
"""Test model_invoking filters out None instructions."""
|
||||
provider1 = MockContextProvider([TextContent("Instructions 1")])
|
||||
provider2 = MockContextProvider(None) # None instructions
|
||||
provider3 = MockContextProvider([TextContent("Instructions 3")])
|
||||
"""Test invoking filters out None instructions."""
|
||||
provider1 = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions 1")])
|
||||
provider2 = MockContextProvider(messages=None) # None instructions
|
||||
provider3 = MockContextProvider(messages=[ChatMessage(role="user", text="Instructions 3")])
|
||||
aggregate = AggregateContextProvider([provider1, provider2, provider3])
|
||||
|
||||
message = ChatMessage(text="Hello", role=Role.USER)
|
||||
context = await aggregate.model_invoking(message)
|
||||
context = await aggregate.invoking(message)
|
||||
|
||||
assert isinstance(context, Context)
|
||||
assert context.contents
|
||||
assert isinstance(context.contents[0], TextContent)
|
||||
assert isinstance(context.contents[1], TextContent)
|
||||
assert context.contents[0].text == "Instructions 1"
|
||||
assert context.contents[1].text == "Instructions 3"
|
||||
assert context.messages
|
||||
assert isinstance(context.messages[0].contents[0], TextContent)
|
||||
assert isinstance(context.messages[1].contents[0], TextContent)
|
||||
assert context.messages[0].text == "Instructions 1"
|
||||
assert context.messages[1].text == "Instructions 3"
|
||||
|
||||
async def test_model_invoking_with_all_none_instructions(self) -> None:
|
||||
"""Test model_invoking when all providers return None instructions."""
|
||||
"""Test invoking when all providers return None instructions."""
|
||||
provider1 = MockContextProvider(None)
|
||||
provider2 = MockContextProvider(None)
|
||||
aggregate = AggregateContextProvider([provider1, provider2])
|
||||
|
||||
message = ChatMessage(text="Hello", role=Role.USER)
|
||||
context = await aggregate.model_invoking(message)
|
||||
context = await aggregate.invoking(message)
|
||||
|
||||
assert isinstance(context, Context)
|
||||
assert not context.contents
|
||||
assert not context.messages
|
||||
|
||||
async def test_model_invoking_with_mutable_sequence(self) -> None:
|
||||
"""Test model_invoking with MutableSequence of messages."""
|
||||
provider = MockContextProvider([TextContent("Test instructions")])
|
||||
"""Test invoking with MutableSequence of messages."""
|
||||
provider = MockContextProvider(messages=[ChatMessage(role="user", text="Test instructions")])
|
||||
aggregate = AggregateContextProvider([provider])
|
||||
|
||||
messages = [ChatMessage(text="Hello", role=Role.USER)]
|
||||
context = await aggregate.model_invoking(messages)
|
||||
context = await aggregate.invoking(messages)
|
||||
|
||||
assert provider.model_invoking_called
|
||||
assert provider.invoking_called
|
||||
assert provider.model_invoking_messages == messages
|
||||
assert isinstance(context, Context)
|
||||
assert context.contents
|
||||
assert isinstance(context.contents[0], TextContent)
|
||||
assert context.contents[0].text == "Test instructions"
|
||||
assert context.messages
|
||||
assert isinstance(context.messages[0].contents[0], TextContent)
|
||||
assert context.messages[0].text == "Test instructions"
|
||||
|
||||
async def test_async_methods_concurrent_execution(self) -> None:
|
||||
"""Test that async methods execute providers concurrently."""
|
||||
# Use AsyncMock to verify concurrent execution
|
||||
provider1 = Mock(spec=ContextProvider)
|
||||
provider1.thread_created = AsyncMock()
|
||||
provider1.messages_adding = AsyncMock()
|
||||
provider1.model_invoking = AsyncMock(return_value=Context(contents=[TextContent("Test 1")]))
|
||||
provider1.invoked = AsyncMock()
|
||||
provider1.invoking = AsyncMock(return_value=Context(messages=[ChatMessage(role="user", text="Test 1")]))
|
||||
|
||||
provider2 = Mock(spec=ContextProvider)
|
||||
provider2.thread_created = AsyncMock()
|
||||
provider2.messages_adding = AsyncMock()
|
||||
provider2.model_invoking = AsyncMock(return_value=Context(contents=[TextContent("Test 2")]))
|
||||
provider2.invoked = AsyncMock()
|
||||
provider2.invoking = AsyncMock(return_value=Context(messages=[ChatMessage(role="user", text="Test 2")]))
|
||||
|
||||
aggregate = AggregateContextProvider([provider1, provider2])
|
||||
|
||||
@@ -285,18 +277,20 @@ class TestAggregateContextProvider:
|
||||
provider1.thread_created.assert_called_once_with("thread-123")
|
||||
provider2.thread_created.assert_called_once_with("thread-123")
|
||||
|
||||
# Test messages_adding
|
||||
# Test invoked
|
||||
message = ChatMessage(text="Hello", role=Role.USER)
|
||||
await aggregate.messages_adding("thread-123", message)
|
||||
provider1.messages_adding.assert_called_once_with("thread-123", message)
|
||||
provider2.messages_adding.assert_called_once_with("thread-123", message)
|
||||
await aggregate.invoked(message)
|
||||
provider1.invoked.assert_called_once_with(
|
||||
request_messages=message, response_messages=None, invoke_exception=None
|
||||
)
|
||||
provider2.invoked.assert_called_once_with(
|
||||
request_messages=message, response_messages=None, invoke_exception=None
|
||||
)
|
||||
|
||||
# Test model_invoking
|
||||
context = await aggregate.model_invoking(message)
|
||||
provider1.model_invoking.assert_called_once_with(message)
|
||||
provider2.model_invoking.assert_called_once_with(message)
|
||||
assert context.contents
|
||||
assert isinstance(context.contents[0], TextContent)
|
||||
assert isinstance(context.contents[1], TextContent)
|
||||
assert context.contents[0].text == "Test 1"
|
||||
assert context.contents[1].text == "Test 2"
|
||||
# Test invoking
|
||||
context = await aggregate.invoking(message)
|
||||
provider1.invoking.assert_called_once_with(message)
|
||||
provider2.invoking.assert_called_once_with(message)
|
||||
assert context.messages
|
||||
assert context.messages[0].text == "Test 1"
|
||||
assert context.messages[1].text == "Test 2"
|
||||
|
||||
@@ -1505,9 +1505,13 @@ class TestChatAgentChatMiddleware:
|
||||
context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]
|
||||
) -> None:
|
||||
# Modify the first message by adding a prefix
|
||||
if context.messages and len(context.messages) > 0:
|
||||
original_text = context.messages[0].text or ""
|
||||
context.messages[0] = ChatMessage(role=context.messages[0].role, text=f"MODIFIED: {original_text}")
|
||||
if context.messages:
|
||||
for idx, msg in enumerate(context.messages):
|
||||
if msg.role.value == "system":
|
||||
continue
|
||||
original_text = msg.text or ""
|
||||
context.messages[idx] = ChatMessage(role=msg.role, text=f"MODIFIED: {original_text}")
|
||||
break
|
||||
await next(context)
|
||||
|
||||
# Create ChatAgent with message-modifying middleware
|
||||
@@ -1519,8 +1523,7 @@ class TestChatAgentChatMiddleware:
|
||||
response = await agent.run(messages)
|
||||
|
||||
# Verify that the message was modified (MockBaseChatClient echoes back the input)
|
||||
assert response is not None
|
||||
assert len(response.messages) > 0
|
||||
assert response and response.messages
|
||||
assert "MODIFIED: test message" in response.messages[0].text
|
||||
|
||||
async def test_chat_middleware_can_override_response(self) -> None:
|
||||
|
||||
@@ -5,12 +5,13 @@ from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from agent_framework import AgentThread, ChatMessage, ChatMessageList, Role
|
||||
from agent_framework._threads import StoreState, ThreadState, deserialize_thread_state, thread_on_new_messages
|
||||
from agent_framework import AgentThread, ChatMessage, ChatMessageStore, Role
|
||||
from agent_framework._threads import AgentThreadState, ChatMessageStoreState
|
||||
from agent_framework.exceptions import AgentThreadException
|
||||
|
||||
|
||||
class MockChatMessageStore:
|
||||
"""Mock implementation of ChatMessageStore for testing."""
|
||||
"""Mock implementation of ChatMessageStoreProtocol for testing."""
|
||||
|
||||
def __init__(self, messages: list[ChatMessage] | None = None) -> None:
|
||||
self._messages = messages or []
|
||||
@@ -23,15 +24,21 @@ class MockChatMessageStore:
|
||||
async def add_messages(self, messages: Sequence[ChatMessage]) -> None:
|
||||
self._messages.extend(messages)
|
||||
|
||||
async def serialize_state(self, **kwargs: Any) -> Any:
|
||||
async def serialize(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:
|
||||
async def update_from_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"]
|
||||
|
||||
@classmethod
|
||||
async def deserialize(cls, serialized_store_state: Any, **kwargs: Any) -> "MockChatMessageStore":
|
||||
instance = cls()
|
||||
await instance.update_from_state(serialized_store_state, **kwargs)
|
||||
return instance
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_messages() -> list[ChatMessage]:
|
||||
@@ -67,7 +74,7 @@ class TestAgentThread:
|
||||
|
||||
def test_init_with_message_store(self) -> None:
|
||||
"""Test AgentThread initialization with message_store."""
|
||||
store = ChatMessageList()
|
||||
store = ChatMessageStore()
|
||||
thread = AgentThread(message_store=store)
|
||||
assert thread.service_thread_id is None
|
||||
assert thread.message_store is store
|
||||
@@ -81,11 +88,11 @@ class TestAgentThread:
|
||||
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()
|
||||
"""Test that setting service_thread_id when message_store exists raises AgentThreadException."""
|
||||
store = ChatMessageStore()
|
||||
thread = AgentThread(message_store=store)
|
||||
|
||||
with pytest.raises(ValueError, match="Only the service_thread_id or message_store may be set"):
|
||||
with pytest.raises(AgentThreadException, 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:
|
||||
@@ -97,18 +104,18 @@ class TestAgentThread:
|
||||
def test_message_store_property_setter(self) -> None:
|
||||
"""Test message_store property setter."""
|
||||
thread = AgentThread()
|
||||
store = ChatMessageList()
|
||||
store = ChatMessageStore()
|
||||
|
||||
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."""
|
||||
"""Test that setting message_store when service_thread_id exists raises AgentThreadException."""
|
||||
service_thread_id = "test-conversation-999"
|
||||
thread = AgentThread(service_thread_id=service_thread_id)
|
||||
store = ChatMessageList()
|
||||
store = ChatMessageStore()
|
||||
|
||||
with pytest.raises(ValueError, match="Only the service_thread_id or message_store may be set"):
|
||||
with pytest.raises(AgentThreadException, 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:
|
||||
@@ -119,7 +126,7 @@ class TestAgentThread:
|
||||
|
||||
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)
|
||||
store = ChatMessageStore(sample_messages)
|
||||
thread = AgentThread(message_store=store)
|
||||
|
||||
assert thread.message_store is not None
|
||||
@@ -142,19 +149,19 @@ class TestAgentThread:
|
||||
"""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)
|
||||
await thread.on_new_messages(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."""
|
||||
"""Test _on_new_messages with single message creates ChatMessageStore."""
|
||||
thread = AgentThread()
|
||||
|
||||
await thread_on_new_messages(thread, sample_message)
|
||||
await thread.on_new_messages(sample_message)
|
||||
|
||||
assert thread.message_store is not None
|
||||
assert isinstance(thread.message_store, ChatMessageList)
|
||||
assert isinstance(thread.message_store, ChatMessageStore)
|
||||
messages = await thread.message_store.list_messages()
|
||||
assert len(messages) == 1
|
||||
assert messages[0].text == "Test message"
|
||||
@@ -163,7 +170,7 @@ class TestAgentThread:
|
||||
"""Test _on_new_messages with multiple messages."""
|
||||
thread = AgentThread()
|
||||
|
||||
await thread_on_new_messages(thread, sample_messages)
|
||||
await thread.on_new_messages(sample_messages)
|
||||
|
||||
assert thread.message_store is not None
|
||||
messages = await thread.message_store.list_messages()
|
||||
@@ -172,10 +179,10 @@ class TestAgentThread:
|
||||
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=Role.USER, text="Initial", message_id="init1")]
|
||||
store = ChatMessageList(initial_messages)
|
||||
store = ChatMessageStore(initial_messages)
|
||||
thread = AgentThread(message_store=store)
|
||||
|
||||
await thread_on_new_messages(thread, sample_message)
|
||||
await thread.on_new_messages(sample_message)
|
||||
|
||||
assert thread.message_store is not None
|
||||
messages = await thread.message_store.list_messages()
|
||||
@@ -185,32 +192,30 @@ class TestAgentThread:
|
||||
|
||||
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)
|
||||
thread = await AgentThread.deserialize(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)
|
||||
thread = await AgentThread.deserialize(serialized_data)
|
||||
|
||||
assert thread.service_thread_id is None
|
||||
assert thread.message_store is not None
|
||||
assert isinstance(thread.message_store, ChatMessageList)
|
||||
assert isinstance(thread.message_store, ChatMessageStore)
|
||||
|
||||
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)
|
||||
await thread.deserialize(serialized_data)
|
||||
|
||||
assert thread.service_thread_id is None
|
||||
assert thread.message_store is None
|
||||
@@ -221,7 +226,7 @@ class TestAgentThread:
|
||||
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)
|
||||
await thread.update_from_thread_state(serialized_data)
|
||||
|
||||
assert store._deserialize_calls == 1 # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
@@ -265,31 +270,31 @@ class TestAgentThread:
|
||||
|
||||
|
||||
class TestChatMessageList:
|
||||
"""Test cases for ChatMessageList class."""
|
||||
"""Test cases for ChatMessageStore class."""
|
||||
|
||||
def test_init_empty(self) -> None:
|
||||
"""Test ChatMessageList initialization with no messages."""
|
||||
store = ChatMessageList()
|
||||
assert len(store) == 0
|
||||
"""Test ChatMessageStore initialization with no messages."""
|
||||
store = ChatMessageStore()
|
||||
assert len(store.messages) == 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
|
||||
"""Test ChatMessageStore initialization with messages."""
|
||||
store = ChatMessageStore(sample_messages)
|
||||
assert len(store.messages) == 3
|
||||
|
||||
async def test_add_messages(self, sample_messages: list[ChatMessage]) -> None:
|
||||
"""Test adding messages to the store."""
|
||||
store = ChatMessageList()
|
||||
store = ChatMessageStore()
|
||||
|
||||
await store.add_messages(sample_messages)
|
||||
|
||||
assert len(store) == 3
|
||||
assert len(store.messages) == 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)
|
||||
store = ChatMessageStore(sample_messages)
|
||||
|
||||
messages = await store.list_messages()
|
||||
|
||||
@@ -298,28 +303,28 @@ class TestChatMessageList:
|
||||
|
||||
async def test_serialize_state(self, sample_messages: list[ChatMessage]) -> None:
|
||||
"""Test serializing store state."""
|
||||
store = ChatMessageList(sample_messages)
|
||||
store = ChatMessageStore(sample_messages)
|
||||
|
||||
result = await store.serialize_state()
|
||||
result = await store.serialize()
|
||||
|
||||
assert "messages" in result
|
||||
assert len(result["messages"]) == 3
|
||||
|
||||
async def test_serialize_state_empty(self) -> None:
|
||||
"""Test serializing empty store state."""
|
||||
store = ChatMessageList()
|
||||
store = ChatMessageStore()
|
||||
|
||||
result = await store.serialize_state()
|
||||
result = await store.serialize()
|
||||
|
||||
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()
|
||||
store = ChatMessageStore()
|
||||
state_data = {"messages": sample_messages}
|
||||
|
||||
await store.deserialize_state(state_data)
|
||||
await store.update_from_state(state_data)
|
||||
|
||||
messages = await store.list_messages()
|
||||
assert len(messages) == 3
|
||||
@@ -327,156 +332,67 @@ class TestChatMessageList:
|
||||
|
||||
async def test_deserialize_state_none(self) -> None:
|
||||
"""Test deserializing None state."""
|
||||
store = ChatMessageList()
|
||||
store = ChatMessageStore()
|
||||
|
||||
await store.deserialize_state(None)
|
||||
await store.update_from_state(None)
|
||||
|
||||
assert len(store) == 0
|
||||
assert len(store.messages) == 0
|
||||
|
||||
async def test_deserialize_state_empty(self) -> None:
|
||||
"""Test deserializing empty state."""
|
||||
store = ChatMessageList()
|
||||
store = ChatMessageStore()
|
||||
|
||||
await store.deserialize_state({})
|
||||
await store.update_from_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?"
|
||||
assert len(store.messages) == 0
|
||||
|
||||
|
||||
class TestStoreState:
|
||||
"""Test cases for StoreState class."""
|
||||
"""Test cases for ChatMessageStoreState class."""
|
||||
|
||||
def test_init(self, sample_messages: list[ChatMessage]) -> None:
|
||||
"""Test StoreState initialization."""
|
||||
state = StoreState(messages=sample_messages)
|
||||
"""Test ChatMessageStoreState initialization."""
|
||||
state = ChatMessageStoreState(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=[])
|
||||
"""Test ChatMessageStoreState initialization with empty messages."""
|
||||
state = ChatMessageStoreState(messages=[])
|
||||
|
||||
assert len(state.messages) == 0
|
||||
|
||||
|
||||
class TestThreadState:
|
||||
"""Test cases for ThreadState class."""
|
||||
"""Test cases for AgentThreadState 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")
|
||||
"""Test AgentThreadState initialization with service_thread_id."""
|
||||
state = AgentThreadState(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."""
|
||||
"""Test AgentThreadState initialization with chat_message_store_state."""
|
||||
store_data: dict[str, Any] = {"messages": []}
|
||||
state = ThreadState(chat_message_store_state=store_data)
|
||||
state = AgentThreadState(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."""
|
||||
"""Test AgentThreadState 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
|
||||
with pytest.raises(
|
||||
AgentThreadException, match="Only one of service_thread_id or chat_message_store_state may be set"
|
||||
):
|
||||
AgentThreadState(service_thread_id="test-conv-123", chat_message_store_state=store_data)
|
||||
|
||||
def test_init_defaults(self) -> None:
|
||||
"""Test ThreadState initialization with defaults."""
|
||||
state = ThreadState()
|
||||
"""Test AgentThreadState initialization with defaults."""
|
||||
state = AgentThreadState()
|
||||
|
||||
assert state.service_thread_id is None
|
||||
assert state.chat_message_store_state is None
|
||||
|
||||
@@ -678,7 +678,6 @@ def test_chat_response_updates_to_chat_response_multiple_multiple():
|
||||
assert chat_response.text == "I'm doing well, thank you! More contextFinal part"
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_chat_response_from_async_generator():
|
||||
async def gen() -> AsyncIterable[ChatResponseUpdate]:
|
||||
yield ChatResponseUpdate(text="Hello", message_id="1")
|
||||
@@ -688,7 +687,6 @@ async def test_chat_response_from_async_generator():
|
||||
assert resp.text == "Hello world"
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_chat_response_from_async_generator_output_format():
|
||||
async def gen() -> AsyncIterable[ChatResponseUpdate]:
|
||||
yield ChatResponseUpdate(text='{ "respon', message_id="1")
|
||||
@@ -702,7 +700,6 @@ async def test_chat_response_from_async_generator_output_format():
|
||||
assert resp.value.response == "Hello"
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_chat_response_from_async_generator_output_format_in_method():
|
||||
async def gen() -> AsyncIterable[ChatResponseUpdate]:
|
||||
yield ChatResponseUpdate(text='{ "respon', message_id="1")
|
||||
@@ -1199,7 +1196,6 @@ def agent_run_response_async() -> AgentRunResponse:
|
||||
return AgentRunResponse(messages=[ChatMessage(role="user", text="Hello")])
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_agent_run_response_from_async_generator():
|
||||
async def gen():
|
||||
yield AgentRunResponseUpdate(contents=[TextContent("A")])
|
||||
|
||||
@@ -383,7 +383,6 @@ async def test_openai_assistants_client_prepare_thread_existing_no_run(mock_asyn
|
||||
mock_async_openai.beta.threads.runs.cancel.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_assistants_client_process_stream_events_thread_run_created(mock_async_openai: MagicMock) -> None:
|
||||
"""Test _process_stream_events with thread.run.created event."""
|
||||
chat_client = create_test_openai_assistants_client(mock_async_openai)
|
||||
@@ -417,7 +416,6 @@ async def test_openai_assistants_client_process_stream_events_thread_run_created
|
||||
assert update.raw_representation == mock_response.data
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_assistants_client_process_stream_events_message_delta_text(mock_async_openai: MagicMock) -> None:
|
||||
"""Test _process_stream_events with thread.message.delta event containing text."""
|
||||
chat_client = create_test_openai_assistants_client(mock_async_openai)
|
||||
@@ -462,7 +460,6 @@ async def test_openai_assistants_client_process_stream_events_message_delta_text
|
||||
assert update.raw_representation == mock_message_delta
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_assistants_client_process_stream_events_requires_action(mock_async_openai: MagicMock) -> None:
|
||||
"""Test _process_stream_events with thread.run.requires_action event."""
|
||||
chat_client = create_test_openai_assistants_client(mock_async_openai)
|
||||
@@ -506,7 +503,6 @@ async def test_openai_assistants_client_process_stream_events_requires_action(mo
|
||||
chat_client._create_function_call_contents.assert_called_once_with(mock_run, None) # type: ignore
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_assistants_client_process_stream_events_run_step_created(mock_async_openai: MagicMock) -> None:
|
||||
"""Test _process_stream_events with thread.run.step.created event."""
|
||||
|
||||
@@ -539,7 +535,6 @@ async def test_openai_assistants_client_process_stream_events_run_step_created(m
|
||||
assert len(updates) == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_assistants_client_process_stream_events_run_completed_with_usage(
|
||||
mock_async_openai: MagicMock,
|
||||
) -> None:
|
||||
|
||||
@@ -101,7 +101,6 @@ class TimedApproval(RequestInfoMessage):
|
||||
issued_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rehydrate_falls_back_when_request_type_missing() -> None:
|
||||
"""Rehydration should succeed even if the original request type cannot be imported.
|
||||
|
||||
@@ -144,7 +143,6 @@ async def test_rehydrate_falls_back_when_request_type_missing() -> None:
|
||||
assert getattr(event.data, "iteration", None) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_has_pending_request_detects_snapshot() -> None:
|
||||
request_id = "req-pending"
|
||||
snapshot = {
|
||||
@@ -172,7 +170,6 @@ async def test_has_pending_request_detects_snapshot() -> None:
|
||||
assert await executor.has_pending_request(request_id, ctx)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_has_pending_request_false_when_snapshot_absent() -> None:
|
||||
shared_state = SharedState()
|
||||
runner_ctx = _StubRunnerContext({"pending_requests": {}})
|
||||
|
||||
Reference in New Issue
Block a user