123

backend/scientist_team.py

@@ -7,7 +7,8 @@ from autogen_agentchat.messages import AgentEvent, ChatMessage, TextMessage, Too
 from autogen_agentchat.teams import SelectorGroupChat, RoundRobinGroupChat, Swarm
 from autogen_agentchat.ui import Console
 from autogen_ext.models.openai import OpenAIChatCompletionClient
-from backend.constant import MODEL, OPENAI_API_KEY, OPENAI_BASE_URL
+from backend.constant import MODEL, OPENAI_API_KEY, OPENAI_BASE_URL, GRAPH_RAG
+from autogen_ext.tools.graphrag import GlobalSearchTool
 from backend.tools import hybird_retrieval_from_knowledge_base, search_from_oqmd_by_composition
 from backend.custom import SocietyOfMindAgent
 
@@ -47,6 +48,7 @@ def create_scientist_team(user_input_func: Callable[[str, Optional[CancellationT
         handoffs=["Synthesis_Scientist", "Structure_Scientist", "Property_Scientist", "Application_Scientist"]
     )
 
+    graph_rag_search = GlobalSearchTool.from_settings(settings_path=GRAPH_RAG)
     synthesis_agent = AssistantAgent(
         "Synthesis_Scientist",
         description="An experienced materials scientist agent who is particularly good at coming up with detailed synthesis schemes, and should be called when the task around a material synthesis topic.",
@@ -86,7 +88,7 @@ def create_scientist_team(user_input_func: Callable[[str, Optional[CancellationT
         Always handoff back to Scientist_Admin when synthesis scheme is complete.
         Let's think step by step and answer with {lang}:
         """,
-        tools=[hybird_retrieval_from_knowledge_base],
+        tools=[hybird_retrieval_from_knowledge_base, search_from_oqmd_by_composition, graph_rag_search],
         reflect_on_tool_use=True,
         handoffs=["Scientist_Admin"]
     )
@@ -106,7 +108,7 @@ def create_scientist_team(user_input_func: Callable[[str, Optional[CancellationT
         Always handoff back to Scientist_Admin when response is complete.
         Answer with {lang}:
         """,
-        tools=[hybird_retrieval_from_knowledge_base],
+        tools=[hybird_retrieval_from_knowledge_base, search_from_oqmd_by_composition, graph_rag_search],
         reflect_on_tool_use=True,
         handoffs=["Scientist_Admin"]
     )
@@ -127,7 +129,7 @@ def create_scientist_team(user_input_func: Callable[[str, Optional[CancellationT
         Always handoff back to Scientist_Admin when response is complete.
         Answer with {lang}:
         """,
-        tools=[hybird_retrieval_from_knowledge_base],
+        tools=[hybird_retrieval_from_knowledge_base, search_from_oqmd_by_composition, graph_rag_search],
         reflect_on_tool_use=True,
         handoffs=["Scientist_Admin"]
     )
@@ -148,7 +150,7 @@ def create_scientist_team(user_input_func: Callable[[str, Optional[CancellationT
         Always handoff back to Scientist_Admin when response is complete.
         Answer with {lang}:
         """,
-        tools=[hybird_retrieval_from_knowledge_base],
+        tools=[hybird_retrieval_from_knowledge_base, search_from_oqmd_by_composition, graph_rag_search],
         reflect_on_tool_use=True,
         handoffs=["Scientist_Admin"]
     )