mirror of
https://github.com/microsoft/autogen.git
synced 2026-01-26 04:18:14 -05:00
* chat completion agent * use response format enum; use Message type for chat history; remove name from state
269 lines
8.8 KiB
Python
269 lines
8.8 KiB
Python
import argparse
|
|
import asyncio
|
|
import logging
|
|
|
|
import openai
|
|
from agnext.agent_components.types import SystemMessage
|
|
from agnext.application_components import (
|
|
SingleThreadedAgentRuntime,
|
|
)
|
|
from agnext.chat.agents.chat_completion_agent import ChatCompletionAgent
|
|
from agnext.chat.agents.oai_assistant import OpenAIAssistantAgent
|
|
from agnext.chat.patterns.group_chat import GroupChat, GroupChatOutput
|
|
from agnext.chat.patterns.orchestrator_chat import OrchestratorChat
|
|
from agnext.chat.types import TextMessage
|
|
from agnext.core._agent import Agent
|
|
from agnext.agent_components.model_client import OpenAI
|
|
from agnext.core.intervention import DefaultInterventionHandler, DropMessage
|
|
from typing_extensions import Any, override
|
|
|
|
logging.basicConfig(level=logging.WARNING)
|
|
logging.getLogger("agnext").setLevel(logging.DEBUG)
|
|
|
|
|
|
class ConcatOutput(GroupChatOutput):
|
|
def __init__(self) -> None:
|
|
self._output = ""
|
|
|
|
def on_message_received(self, message: Any) -> None:
|
|
match message:
|
|
case TextMessage(content=content):
|
|
self._output += content
|
|
case _:
|
|
...
|
|
|
|
def get_output(self) -> Any:
|
|
return self._output
|
|
|
|
def reset(self) -> None:
|
|
self._output = ""
|
|
|
|
|
|
class LoggingHandler(DefaultInterventionHandler):
|
|
send_color = "\033[31m"
|
|
response_color = "\033[34m"
|
|
reset_color = "\033[0m"
|
|
|
|
@override
|
|
async def on_send(self, message: Any, *, sender: Agent | None, recipient: Agent) -> Any | type[DropMessage]:
|
|
if sender is None:
|
|
print(f"{self.send_color}Sending message to {recipient.name}:{self.reset_color} {message}")
|
|
else:
|
|
print(
|
|
f"{self.send_color}Sending message from {sender.name} to {recipient.name}:{self.reset_color} {message}"
|
|
)
|
|
return message
|
|
|
|
@override
|
|
async def on_response(self, message: Any, *, sender: Agent, recipient: Agent | None) -> Any | type[DropMessage]:
|
|
if recipient is None:
|
|
print(f"{self.response_color}Received response from {sender.name}:{self.reset_color} {message}")
|
|
else:
|
|
print(
|
|
f"{self.response_color}Received response from {sender.name} to {recipient.name}:{self.reset_color} {message}"
|
|
)
|
|
return message
|
|
|
|
|
|
async def group_chat(message: str) -> None:
|
|
runtime = SingleThreadedAgentRuntime(before_send=LoggingHandler())
|
|
|
|
joe_oai_assistant = openai.beta.assistants.create(
|
|
model="gpt-3.5-turbo",
|
|
name="Joe",
|
|
instructions="You are a commedian named Joe. Make the audience laugh.",
|
|
)
|
|
joe_oai_thread = openai.beta.threads.create()
|
|
joe = OpenAIAssistantAgent(
|
|
name="Joe",
|
|
description="Joe the commedian.",
|
|
runtime=runtime,
|
|
client=openai.AsyncClient(),
|
|
assistant_id=joe_oai_assistant.id,
|
|
thread_id=joe_oai_thread.id,
|
|
)
|
|
|
|
cathy_oai_assistant = openai.beta.assistants.create(
|
|
model="gpt-3.5-turbo",
|
|
name="Cathy",
|
|
instructions="You are a poet named Cathy. Write beautiful poems.",
|
|
)
|
|
cathy_oai_thread = openai.beta.threads.create()
|
|
cathy = OpenAIAssistantAgent(
|
|
name="Cathy",
|
|
description="Cathy the poet.",
|
|
runtime=runtime,
|
|
client=openai.AsyncClient(),
|
|
assistant_id=cathy_oai_assistant.id,
|
|
thread_id=cathy_oai_thread.id,
|
|
)
|
|
|
|
chat = GroupChat(
|
|
"Host",
|
|
"A round-robin chat room.",
|
|
runtime,
|
|
[joe, cathy],
|
|
num_rounds=5,
|
|
output=ConcatOutput(),
|
|
)
|
|
|
|
response = runtime.send_message(TextMessage(content=message, source="host"), chat)
|
|
|
|
while not response.done():
|
|
await runtime.process_next()
|
|
|
|
await response
|
|
|
|
|
|
async def orchestrator_oai_assistant(message: str) -> None:
|
|
runtime = SingleThreadedAgentRuntime(before_send=LoggingHandler())
|
|
|
|
developer_oai_assistant = openai.beta.assistants.create(
|
|
model="gpt-3.5-turbo",
|
|
name="Developer",
|
|
instructions="You are a Python developer.",
|
|
)
|
|
developer_oai_thread = openai.beta.threads.create()
|
|
developer = OpenAIAssistantAgent(
|
|
name="Developer",
|
|
description="A developer that writes code.",
|
|
runtime=runtime,
|
|
client=openai.AsyncClient(),
|
|
assistant_id=developer_oai_assistant.id,
|
|
thread_id=developer_oai_thread.id,
|
|
)
|
|
|
|
product_manager_oai_assistant = openai.beta.assistants.create(
|
|
model="gpt-3.5-turbo",
|
|
name="ProductManager",
|
|
instructions="You are a product manager good at translating customer needs into software specifications.",
|
|
)
|
|
product_manager_oai_thread = openai.beta.threads.create()
|
|
product_manager = OpenAIAssistantAgent(
|
|
name="ProductManager",
|
|
description="A product manager that plans and comes up with specs.",
|
|
runtime=runtime,
|
|
client=openai.AsyncClient(),
|
|
assistant_id=product_manager_oai_assistant.id,
|
|
thread_id=product_manager_oai_thread.id,
|
|
)
|
|
|
|
planner_oai_assistant = openai.beta.assistants.create(
|
|
model="gpt-4-turbo",
|
|
name="Planner",
|
|
instructions="You are a planner of complex tasks.",
|
|
)
|
|
planner_oai_thread = openai.beta.threads.create()
|
|
planner = OpenAIAssistantAgent(
|
|
name="Planner",
|
|
description="A planner that organizes and schedules tasks.",
|
|
runtime=runtime,
|
|
client=openai.AsyncClient(),
|
|
assistant_id=planner_oai_assistant.id,
|
|
thread_id=planner_oai_thread.id,
|
|
)
|
|
|
|
orchestrator_oai_assistant = openai.beta.assistants.create(
|
|
model="gpt-4-turbo",
|
|
name="Orchestrator",
|
|
instructions="You are an orchestrator that coordinates the team to complete a complex task.",
|
|
)
|
|
orchestrator_oai_thread = openai.beta.threads.create()
|
|
orchestrator = OpenAIAssistantAgent(
|
|
name="Orchestrator",
|
|
description="An orchestrator that coordinates the team.",
|
|
runtime=runtime,
|
|
client=openai.AsyncClient(),
|
|
assistant_id=orchestrator_oai_assistant.id,
|
|
thread_id=orchestrator_oai_thread.id,
|
|
)
|
|
|
|
chat = OrchestratorChat(
|
|
"OrchestratorChat",
|
|
"A software development team.",
|
|
runtime,
|
|
orchestrator=orchestrator,
|
|
planner=planner,
|
|
specialists=[developer, product_manager],
|
|
)
|
|
|
|
response = runtime.send_message(TextMessage(content=message, source="Customer"), chat)
|
|
|
|
while not response.done():
|
|
await runtime.process_next()
|
|
|
|
print((await response).content) # type: ignore
|
|
|
|
|
|
async def orchestrator_chat_completion(message: str) -> None:
|
|
runtime = SingleThreadedAgentRuntime(before_send=LoggingHandler())
|
|
|
|
developer = ChatCompletionAgent(
|
|
name="Developer",
|
|
description="A developer that writes code.",
|
|
runtime=runtime,
|
|
system_messages=[SystemMessage("You are a Python developer.")],
|
|
model_client=OpenAI(model="gpt-3.5-turbo"),
|
|
)
|
|
|
|
product_manager = ChatCompletionAgent(
|
|
name="ProductManager",
|
|
description="A product manager that plans and comes up with specs.",
|
|
runtime=runtime,
|
|
system_messages=[
|
|
SystemMessage("You are a product manager good at translating customer needs into software specifications.")
|
|
],
|
|
model_client=OpenAI(model="gpt-3.5-turbo"),
|
|
)
|
|
|
|
planner = ChatCompletionAgent(
|
|
name="Planner",
|
|
description="A planner that organizes and schedules tasks.",
|
|
runtime=runtime,
|
|
system_messages=[SystemMessage("You are a planner of complex tasks.")],
|
|
model_client=OpenAI(model="gpt-4-turbo"),
|
|
)
|
|
|
|
orchestrator = ChatCompletionAgent(
|
|
name="Orchestrator",
|
|
description="An orchestrator that coordinates the team.",
|
|
runtime=runtime,
|
|
system_messages=[
|
|
SystemMessage("You are an orchestrator that coordinates the team to complete a complex task.")
|
|
],
|
|
model_client=OpenAI(model="gpt-4-turbo"),
|
|
)
|
|
|
|
chat = OrchestratorChat(
|
|
"OrchestratorChat",
|
|
"A software development team.",
|
|
runtime,
|
|
orchestrator=orchestrator,
|
|
planner=planner,
|
|
specialists=[developer, product_manager],
|
|
)
|
|
|
|
response = runtime.send_message(TextMessage(content=message, source="Customer"), chat)
|
|
|
|
while not response.done():
|
|
await runtime.process_next()
|
|
|
|
print((await response).content) # type: ignore
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Run a pattern demo.")
|
|
choices = {
|
|
"group_chat": group_chat,
|
|
"orchestrator_oai_assistant": orchestrator_oai_assistant,
|
|
"orchestrator_chat_completion": orchestrator_chat_completion,
|
|
}
|
|
parser.add_argument(
|
|
"--pattern",
|
|
choices=list(choices.keys()),
|
|
help="The pattern to demo.",
|
|
)
|
|
parser.add_argument("--message", help="The message to send.")
|
|
args = parser.parse_args()
|
|
asyncio.run(choices[args.pattern](args.message))
|