mirror of
https://github.com/microsoft/autogen.git
synced 2025-06-26 22:30:10 +00:00
46 lines
1.4 KiB
Python
46 lines
1.4 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 ModelFamily
|
||
|
|
from autogen_ext.models.openai import OpenAIChatCompletionClient
|
||
|
|
|
||
|
|
|
||
|
|
@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",
|
||
|
|
base_url="https://generativelanguage.googleapis.com/v1beta/openai/",
|
||
|
|
api_key=api_key,
|
||
|
|
model_info={
|
||
|
|
"vision": True,
|
||
|
|
"function_calling": True,
|
||
|
|
"json_output": True,
|
||
|
|
"family": ModelFamily.GEMINI_1_5_FLASH,
|
||
|
|
},
|
||
|
|
)
|
||
|
|
|
||
|
|
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."))
|