优化结构新增scientist和engineer
This commit is contained in:
@@ -1,13 +1,14 @@
|
||||
import asyncio
|
||||
from typing import Sequence
|
||||
from autogen_agentchat.agents import AssistantAgent
|
||||
from autogen_agentchat.agents import AssistantAgent, SocietyOfMindAgent
|
||||
from autogen_agentchat.conditions import MaxMessageTermination, TextMentionTermination
|
||||
from autogen_agentchat.messages import AgentEvent, ChatMessage, TextMessage, ToolCallExecutionEvent
|
||||
from autogen_agentchat.teams import SelectorGroupChat, RoundRobinGroupChat
|
||||
from autogen_agentchat.ui import Console
|
||||
from autogen_ext.models.openai import OpenAIChatCompletionClient
|
||||
from constant import MODEL, OPENAI_API_KEY, OPENAI_BASE_URL
|
||||
from tools import retrieval_from_knowledge_base, search_from_oqmd_by_composition
|
||||
from scientist_team import create_scientist_team
|
||||
from engineer_team import create_engineer_team
|
||||
|
||||
model_client = OpenAIChatCompletionClient(
|
||||
model=MODEL,
|
||||
@@ -22,54 +23,34 @@ model_client = OpenAIChatCompletionClient(
|
||||
)
|
||||
|
||||
|
||||
def create_team() -> SelectorGroupChat:
|
||||
async def main(task: str = "") -> dict:
|
||||
scientist_team = create_scientist_team()
|
||||
engineer_team = create_engineer_team()
|
||||
|
||||
result = {}
|
||||
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 search tasks into smaller, manageable subtasks.
|
||||
Assign these subtasks to the appropriate team members; not all team members are required to participate in every task.
|
||||
Your team members are:
|
||||
Vector search agent: Searches for paper information in Vector database of knowledge base.
|
||||
OQMD search agent: Searches for crystal structure and property information in OQMD database by composition.
|
||||
|
||||
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:
|
||||
Scientist team: A professional team of material scientists who are mainly responsible for consulting on material synthesis, structure, application and properties.
|
||||
Engineer team: A team of professional engineers who are responsible for writing code, visualizing experimental schemes, converting experimental schemes to machine code, and more.
|
||||
You only plan and delegate tasks - you do not execute them yourself.
|
||||
|
||||
When assigning tasks, use this format:
|
||||
When assigning subtasks, use this format:
|
||||
1. <agent> : <task>
|
||||
|
||||
After all search tasks are complete, summarize the findings and end with "TERMINATE".
|
||||
""",
|
||||
)
|
||||
When assigning subtasks, give a flow chart with following format or mermaid to visualize the collaboration between the various teams, such as:
|
||||
<agent 1> -> <agent 2> -> <agent 3>
|
||||
|
||||
vector_search_agent = AssistantAgent(
|
||||
"VectorSearcher",
|
||||
description="A vector search agent.",
|
||||
tools=[retrieval_from_knowledge_base],
|
||||
model_client=model_client,
|
||||
system_message="""
|
||||
You are a vector search agent.
|
||||
Your only tool is retrieval_from_knowledge_base - use it to find information.
|
||||
You make only one search call at a time.
|
||||
Once you have the results, you never do calculations based on them.
|
||||
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, # Set to True to have the model reflect on the tool use, set to False to return the tool call result directly.
|
||||
)
|
||||
|
||||
oqmd_database_search_agent = AssistantAgent(
|
||||
"OQMDDatabaseSearcher",
|
||||
description="A database search agent.",
|
||||
tools=[search_from_oqmd_by_composition],
|
||||
model_client=model_client,
|
||||
system_message="""
|
||||
You are a database search agent of OQMD.
|
||||
Your only tool is search_from_oqmd_by_composition - use it to find information.
|
||||
You make only one search call at a time.
|
||||
Once you have the results, you never do calculations based on them.
|
||||
""",
|
||||
reflect_on_tool_use=False, # Set to True to have the model reflect on the tool use, set to False to return the tool call result directly.
|
||||
reflect_on_tool_use=False
|
||||
)
|
||||
|
||||
# The termination condition is a combination of text mention termination and max message termination.
|
||||
@@ -85,31 +66,26 @@ def create_team() -> SelectorGroupChat:
|
||||
return None
|
||||
|
||||
team = SelectorGroupChat(
|
||||
[planning_agent, vector_search_agent, oqmd_database_search_agent],
|
||||
[planning_agent, scientist_team, engineer_team],
|
||||
model_client=model_client, # Use a smaller model for the selector.
|
||||
termination_condition=termination,
|
||||
selector_func=selector_func,
|
||||
)
|
||||
return team
|
||||
|
||||
async def main(task: str = "") -> dict:
|
||||
team = create_team()
|
||||
|
||||
result = {}
|
||||
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]
|
||||
|
||||
await Console(team.run_stream(task=task))
|
||||
# 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("How to synthesis CsPbBr3 nanocubes at room temperature?")
|
||||
# Now you can use result in main_1
|
||||
result = await main("Let the robot synthesize CsPbBr3 nanocubes at room temperature")
|
||||
# result = await main("查一下CsPbBr3的晶体结构")
|
||||
|
||||
print(result)
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
Reference in New Issue
Block a user