mirror of
https://github.com/microsoft/autogen.git
synced 2025-07-09 01:51:16 +00:00
7c0bbf674f
* feat: add support for list of messages as team task input * Update society of mind agent to use the list input task --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: Ryan Sweet <rysweet@microsoft.com> Co-authored-by: Eric Zhu <ekzhu@users.noreply.github.com>
92 lines
4.1 KiB
Python
92 lines
4.1 KiB
Python
import asyncio
|
|
from typing import Any, AsyncGenerator, List
|
|
|
|
import pytest
|
|
from autogen_agentchat.agents import AssistantAgent, SocietyOfMindAgent
|
|
from autogen_agentchat.conditions import MaxMessageTermination
|
|
from autogen_agentchat.teams import RoundRobinGroupChat
|
|
from autogen_ext.models.openai import OpenAIChatCompletionClient
|
|
from openai.resources.chat.completions import AsyncCompletions
|
|
from openai.types.chat.chat_completion import ChatCompletion, Choice
|
|
from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
|
|
from openai.types.chat.chat_completion_message import ChatCompletionMessage
|
|
from openai.types.completion_usage import CompletionUsage
|
|
|
|
|
|
class _MockChatCompletion:
|
|
def __init__(self, chat_completions: List[ChatCompletion]) -> None:
|
|
self._saved_chat_completions = chat_completions
|
|
self._curr_index = 0
|
|
|
|
async def mock_create(
|
|
self, *args: Any, **kwargs: Any
|
|
) -> ChatCompletion | AsyncGenerator[ChatCompletionChunk, None]:
|
|
await asyncio.sleep(0.1)
|
|
completion = self._saved_chat_completions[self._curr_index]
|
|
self._curr_index += 1
|
|
return completion
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_society_of_mind_agent(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
model = "gpt-4o-2024-05-13"
|
|
chat_completions = [
|
|
ChatCompletion(
|
|
id="id2",
|
|
choices=[
|
|
Choice(finish_reason="stop", index=0, message=ChatCompletionMessage(content="1", role="assistant"))
|
|
],
|
|
created=0,
|
|
model=model,
|
|
object="chat.completion",
|
|
usage=CompletionUsage(prompt_tokens=10, completion_tokens=5, total_tokens=0),
|
|
),
|
|
ChatCompletion(
|
|
id="id2",
|
|
choices=[
|
|
Choice(finish_reason="stop", index=0, message=ChatCompletionMessage(content="2", role="assistant"))
|
|
],
|
|
created=0,
|
|
model=model,
|
|
object="chat.completion",
|
|
usage=CompletionUsage(prompt_tokens=10, completion_tokens=5, total_tokens=0),
|
|
),
|
|
ChatCompletion(
|
|
id="id2",
|
|
choices=[
|
|
Choice(finish_reason="stop", index=0, message=ChatCompletionMessage(content="3", role="assistant"))
|
|
],
|
|
created=0,
|
|
model=model,
|
|
object="chat.completion",
|
|
usage=CompletionUsage(prompt_tokens=10, completion_tokens=5, total_tokens=0),
|
|
),
|
|
]
|
|
mock = _MockChatCompletion(chat_completions)
|
|
monkeypatch.setattr(AsyncCompletions, "create", mock.mock_create)
|
|
model_client = OpenAIChatCompletionClient(model="gpt-4o", api_key="")
|
|
|
|
agent1 = AssistantAgent("assistant1", model_client=model_client, system_message="You are a helpful assistant.")
|
|
agent2 = AssistantAgent("assistant2", model_client=model_client, system_message="You are a helpful assistant.")
|
|
inner_termination = MaxMessageTermination(3)
|
|
inner_team = RoundRobinGroupChat([agent1, agent2], termination_condition=inner_termination)
|
|
society_of_mind_agent = SocietyOfMindAgent("society_of_mind", team=inner_team, model_client=model_client)
|
|
response = await society_of_mind_agent.run(task="Count to 10.")
|
|
assert len(response.messages) == 4
|
|
assert response.messages[0].source == "user"
|
|
assert response.messages[1].source == "assistant1"
|
|
assert response.messages[2].source == "assistant2"
|
|
assert response.messages[3].source == "society_of_mind"
|
|
|
|
# Test save and load state.
|
|
state = await society_of_mind_agent.save_state()
|
|
assert state is not None
|
|
agent1 = AssistantAgent("assistant1", model_client=model_client, system_message="You are a helpful assistant.")
|
|
agent2 = AssistantAgent("assistant2", model_client=model_client, system_message="You are a helpful assistant.")
|
|
inner_termination = MaxMessageTermination(3)
|
|
inner_team = RoundRobinGroupChat([agent1, agent2], termination_condition=inner_termination)
|
|
society_of_mind_agent2 = SocietyOfMindAgent("society_of_mind", team=inner_team, model_client=model_client)
|
|
await society_of_mind_agent2.load_state(state)
|
|
state2 = await society_of_mind_agent2.save_state()
|
|
assert state == state2
|