mirror of
https://github.com/microsoft/autogen.git
synced 2025-10-31 17:59:50 +00:00
56 lines
1.7 KiB
Python
56 lines
1.7 KiB
Python
import os
|
|
|
|
import pytest
|
|
from autogen_agentchat.agents import AssistantAgent
|
|
from autogen_agentchat.teams import SelectorGroupChat
|
|
from autogen_agentchat.ui import Console
|
|
from autogen_core.models import ChatCompletionClient
|
|
from autogen_ext.models.openai import OpenAIChatCompletionClient
|
|
|
|
|
|
async def _test_selector_group_chat(model_client: ChatCompletionClient) -> None:
|
|
assistant = AssistantAgent(
|
|
"assistant",
|
|
description="A helpful assistant agent.",
|
|
model_client=model_client,
|
|
system_message="You are a helpful assistant.",
|
|
)
|
|
|
|
critic = AssistantAgent(
|
|
"critic",
|
|
description="A critic agent to provide feedback.",
|
|
model_client=model_client,
|
|
system_message="Provide feedback.",
|
|
)
|
|
|
|
team = SelectorGroupChat([assistant, critic], model_client=model_client, max_turns=2)
|
|
await Console(team.run_stream(task="Draft a short email about organizing a holiday party for new year."))
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_selector_group_chat_gemini() -> None:
|
|
try:
|
|
api_key = os.environ["GEMINI_API_KEY"]
|
|
except KeyError:
|
|
pytest.skip("GEMINI_API_KEY not set in environment variables.")
|
|
|
|
model_client = OpenAIChatCompletionClient(
|
|
model="gemini-1.5-flash",
|
|
api_key=api_key,
|
|
)
|
|
await _test_selector_group_chat(model_client)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_selector_group_chat_openai() -> None:
|
|
try:
|
|
api_key = os.environ["OPENAI_API_KEY"]
|
|
except KeyError:
|
|
pytest.skip("OPENAI_API_KEY not set in environment variables.")
|
|
|
|
model_client = OpenAIChatCompletionClient(
|
|
model="gpt-4o-mini",
|
|
api_key=api_key,
|
|
)
|
|
await _test_selector_group_chat(model_client)
|