134 lines
7.1 KiB
Python
Executable File
134 lines
7.1 KiB
Python
Executable File
import asyncio
|
||
from typing import Sequence
|
||
from autogen_agentchat.agents import AssistantAgent, SocietyOfMindAgent, UserProxyAgent
|
||
from autogen_agentchat.conditions import MaxMessageTermination, TextMentionTermination, HandoffTermination
|
||
from autogen_agentchat.messages import AgentEvent, ChatMessage, TextMessage, ToolCallExecutionEvent
|
||
from autogen_agentchat.teams import SelectorGroupChat, RoundRobinGroupChat
|
||
from autogen_agentchat.ui import Console
|
||
from autogen_agentchat.base import Handoff
|
||
from autogen_ext.models.openai import OpenAIChatCompletionClient
|
||
from constant import MODEL, OPENAI_API_KEY, OPENAI_BASE_URL, code_executor
|
||
from scientist_team import create_scientist_team
|
||
from engineer_team import create_engineer_team
|
||
from robot_platform import create_robot_team
|
||
|
||
model_client = OpenAIChatCompletionClient(
|
||
model=MODEL,
|
||
base_url=OPENAI_BASE_URL,
|
||
api_key=OPENAI_API_KEY,
|
||
model_info={
|
||
"vision": True,
|
||
"function_calling": True,
|
||
"json_output": True,
|
||
"family": "unknown",
|
||
},
|
||
)
|
||
|
||
async def main(task: str = "") -> dict:
|
||
scientist_team = create_scientist_team()
|
||
engineer_team = create_engineer_team()
|
||
await code_executor.start()
|
||
robot_platform = create_robot_team(code_executor)
|
||
|
||
result = {}
|
||
user = UserProxyAgent("user", input_func=input)
|
||
|
||
planning_agent = AssistantAgent(
|
||
"PlanningAgent",
|
||
description="An agent for planning tasks, this agent should be the first to engage when given a new task.",
|
||
model_client=model_client,
|
||
system_message="""
|
||
You are a planning agent.
|
||
Your job is to break down complex Materials science research tasks into smaller, manageable subtasks.
|
||
Assign these subtasks to the appropriate sub-teams; not all sub-teams are required to participate in every task.
|
||
Your sub-teams are:
|
||
1. User: A human agent to whom you transfer information whenever you need to confirm your execution steps to a human.
|
||
2. Engineer team: A team of professional engineers who are responsible for writing code, visualizing experimental schemes, converting experimental schemes to JSON, and more.
|
||
- The engineer team has the following members:
|
||
2.1 Structural engineer: A professional structural engineer who focus on converting natural language synthesis schemes to JSON or XML formated scheme, and then upload this JSON to S3 Storage.
|
||
2.2 Software_Engineer: A professional Python software engineer will use Python to implement tasks.
|
||
2.3 SandBox_Env: A computer terminal that performs no other action than running Python scripts (provided to it quoted in ```python code blocks), or sh shell scripts (provided to it quoted in ```sh code blocks).
|
||
3. Scientist team: A professional team of material scientists who are mainly responsible for consulting on material synthesis, structure, application and properties.
|
||
- The scientist team has the following members:
|
||
3.1 Synthesis Scientist: who is good at giving perfect and correct synthesis solutions.
|
||
3.2 Structure Scientist: focusing on agents of structural topics in materials science.
|
||
3.3 Property Scientist: focuses on physical and chemistry property topics in materials science.
|
||
3.4 Application Scientist: Focus on practical applications of materials, such as devices, chips, etc.
|
||
4. Robot Platform: A robotic platform is responsible for performing automated synthesis experiments, automated characterization experiments, and collecting experimental logs.
|
||
- The Robot Platform has the following members:
|
||
4.1 RobotIO_Agent: The agent responsible for the input and output of the robot platform. Pass a structured JSON schema to the robot; Get the experiment log of the robot.
|
||
|
||
You only plan and delegate tasks - you do not execute them yourself.
|
||
|
||
回答时你需要初始化/更新如下任务分配表和Mermaid流程图，并按顺序执行，使用如下格式并利用：
|
||
| Team_name | Member_name | sub-task |
|
||
| ----------- | ------------- | ------------------------------------ |
|
||
| <team_name> | <member_name> | <status: brief sub-task description> |
|
||
|
||
```mermaid
|
||
graph TD
|
||
User[User]
|
||
subgraph <team_name>
|
||
A1[<member_name>]
|
||
end
|
||
style xxx # 推荐多样的风格
|
||
...
|
||
User --> A1
|
||
...
|
||
```
|
||
|
||
每次回答时，你需要清晰明确的指出已经完成的子任务下一步子任务，使用如下格式：
|
||
**已完成子任务：**
|
||
1. <team> : <subtask>
|
||
**Next sub-task:**
|
||
n. <team> : <subtask>
|
||
|
||
You can end with "HUMAN" if you need to, which means you need human approval or other advice or instructions;
|
||
After plan and delegate tasks are complete, end with "START";
|
||
Determine if all sub-teams have completed their tasks, and if so, summarize the findings and end with "TERMINATE".
|
||
""",
|
||
reflect_on_tool_use=False
|
||
)
|
||
|
||
# The termination condition is a combination of text mention termination and max message termination.
|
||
text_mention_termination = TextMentionTermination("TERMINATE")
|
||
max_messages_termination = MaxMessageTermination(max_messages=100)
|
||
termination = text_mention_termination | max_messages_termination
|
||
|
||
# The selector function is a function that takes the current message thread of the group chat
|
||
# and returns the next speaker's name. If None is returned, the LLM-based selection method will be used.
|
||
def selector_func(messages: Sequence[AgentEvent | ChatMessage]) -> str | None:
|
||
if messages[-1].source != planning_agent.name:
|
||
return planning_agent.name # Always return to the planning agent after the other agents have spoken.
|
||
elif "HUMAN" in messages[-1].content:
|
||
return user.name
|
||
return None
|
||
|
||
team = SelectorGroupChat(
|
||
[planning_agent, user, scientist_team, engineer_team, robot_platform],
|
||
model_client=model_client, # Use a smaller model for the selector.
|
||
termination_condition=termination,
|
||
selector_func=selector_func,
|
||
)
|
||
await Console(team.run_stream(task=task))
|
||
await code_executor.stop()
|
||
# async for message in team.run_stream(task=task):
|
||
# if isinstance(message, TextMessage):
|
||
# print(f"----------------{message.source}----------------\n {message.content}")
|
||
# result[message.source] = message.content
|
||
# elif isinstance(message, ToolCallExecutionEvent):
|
||
# print(f"----------------{message.source}----------------\n {message.content}")
|
||
# result[message.source] = [content.content for content in message.content]
|
||
return result
|
||
|
||
# Example usage in another function
|
||
async def main_1():
|
||
# result = await main(input("Enter your instructions below: \n"))
|
||
result = await main("Let the robot synthesize CsPbBr3 nanocubes at room temperature")
|
||
# result = await main("查一下CsPbBr3的晶体结构")
|
||
|
||
print(result)
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main_1())
|