Skip to content

Commit 96acca8

Browse files
Rasic2Rasic2
committed
refactor: rename MATMASTER_AGENT_NAME to MATMASTER_CORE_AGENT_NAME to improve clarity and consistency
1 parent ac28416 commit 96acca8

19 files changed

Lines changed: 137 additions & 39 deletions

File tree

‎agents/matmaster_agent/ABACUS_agent/agent.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
ABACUS_AGENT_NAME,
1414
)
1515
from agents.matmaster_agent.base_agents.job_agent import BaseAsyncJobAgent
16-
from agents.matmaster_agent.constant import MATMASTER_AGENT_NAME
16+
from agents.matmaster_agent.constant import MATMASTER_CORE_AGENT_NAME
1717
from agents.matmaster_agent.logger import matmodeler_logging_handler
1818

1919
mcp_tools_abacus = CalculationMCPToolset(
@@ -38,7 +38,7 @@ def __init__(self, llm_config):
3838
agent_description=ABACUS_AGENT_DESCRIPTION,
3939
agent_instruction=ABACUS_AGENT_INSTRUCTION,
4040
dflow_flag=False,
41-
supervisor_agent=MATMASTER_AGENT_NAME,
41+
supervisor_agent=MATMASTER_CORE_AGENT_NAME,
4242
sync_tools=[
4343
'abacus_prepare',
4444
'abacus_modify_input',

‎agents/matmaster_agent/DPACalculator_agent/agent.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
from google.adk.tools.mcp_tool.mcp_session_manager import SseServerParams
44

55
from agents.matmaster_agent.base_agents.job_agent import BaseAsyncJobAgent
6-
from agents.matmaster_agent.constant import MATMASTER_AGENT_NAME
6+
from agents.matmaster_agent.constant import MATMASTER_CORE_AGENT_NAME
77
from agents.matmaster_agent.DPACalculator_agent.constant import (
88
DPACalulator_AGENT_NAME,
99
DPACalulator_BOHRIUM_EXECUTOR,
@@ -35,7 +35,7 @@ def __init__(self, llm_config):
3535
agent_description=DPAAgentDescription,
3636
agent_instruction=DPAAgentInstruction,
3737
dflow_flag=False,
38-
supervisor_agent=MATMASTER_AGENT_NAME,
38+
supervisor_agent=MATMASTER_CORE_AGENT_NAME,
3939
)
4040

4141

‎agents/matmaster_agent/INVAR_agent/agent.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
from google.adk.tools.mcp_tool.mcp_session_manager import SseServerParams
44

55
from agents.matmaster_agent.base_agents.job_agent import BaseAsyncJobAgent
6-
from agents.matmaster_agent.constant import MATMASTER_AGENT_NAME
6+
from agents.matmaster_agent.constant import MATMASTER_CORE_AGENT_NAME
77
from agents.matmaster_agent.INVAR_agent.constant import (
88
INVAR_AGENT_NAME,
99
INVAR_BOHRIUM_EXECUTOR,
@@ -35,7 +35,7 @@ def __init__(self, llm_config):
3535
agent_description=INVARAgentDescription,
3636
agent_instruction=INVARAgentInstruction,
3737
dflow_flag=False,
38-
supervisor_agent=MATMASTER_AGENT_NAME,
38+
supervisor_agent=MATMASTER_CORE_AGENT_NAME,
3939
)
4040

4141

‎agents/matmaster_agent/MrDice_agent/agent.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
from google.adk.agents import BaseAgent
22

33
from agents.matmaster_agent.base_agents.llm_wrap_agent import LlmWrapAgent
4-
from agents.matmaster_agent.constant import MATMASTER_AGENT_NAME
4+
from agents.matmaster_agent.constant import MATMASTER_CORE_AGENT_NAME
55
from agents.matmaster_agent.MrDice_agent.bohriumpublic_agent.agent import (
66
init_bohriumpublic_database_agent,
77
)
@@ -38,7 +38,7 @@ def __init__(self, llm_config):
3838
bohriumpublic_agent,
3939
mofdb_agent,
4040
],
41-
supervisor_agent=MATMASTER_AGENT_NAME,
41+
supervisor_agent=MATMASTER_CORE_AGENT_NAME,
4242
)
4343

4444

‎agents/matmaster_agent/agent.py‎

Lines changed: 28 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import logging
22
from typing import AsyncGenerator, override
33

4-
from google.adk.agents import InvocationContext, LlmAgent
4+
from google.adk.agents import InvocationContext, LlmAgent, LoopAgent
55
from google.adk.events import Event
66
from google.adk.models.lite_llm import LiteLlm
77
from opik.integrations.adk import track_adk_agent_recursive
@@ -13,9 +13,10 @@
1313
matmaster_check_job_status,
1414
matmaster_prepare_state,
1515
matmaster_set_lang,
16+
matmaster_should_continue,
1617
)
1718
from agents.matmaster_agent.chembrain_agent.agent import init_chembrain_agent
18-
from agents.matmaster_agent.constant import MATMASTER_AGENT_NAME, ModelRole
19+
from agents.matmaster_agent.constant import MATMASTER_CORE_AGENT_NAME, ModelRole
1920
from agents.matmaster_agent.document_parser_agent.agent import (
2021
init_document_parser_agent,
2122
)
@@ -58,7 +59,7 @@
5859
)
5960

6061

61-
class MatMasterAgent(HandleFileUploadLlmAgent):
62+
class MatMasterCoreAgent(HandleFileUploadLlmAgent):
6263
def __init__(self, llm_config):
6364
piloteye_electro_agent = init_piloteye_electro_agent(llm_config)
6465
traj_analysis_agent = init_traj_analysis_agent(llm_config)
@@ -79,7 +80,7 @@ def __init__(self, llm_config):
7980
document_parser_agent = init_document_parser_agent(llm_config)
8081

8182
super().__init__(
82-
name=MATMASTER_AGENT_NAME,
83+
name=MATMASTER_CORE_AGENT_NAME,
8384
model=llm_config.gpt_5_chat,
8485
sub_agents=[
8586
piloteye_electro_agent,
@@ -111,6 +112,7 @@ def __init__(self, llm_config):
111112
target_agent_enum=MatMasterTargetAgentEnum,
112113
),
113114
],
115+
after_agent_callback=matmaster_should_continue,
114116
)
115117

116118
@override
@@ -141,7 +143,28 @@ async def _run_async_impl(
141143
yield error_handel_event
142144

143145

144-
def init_matmaster_agent() -> LlmAgent:
146+
class MatMasterAgent(LoopAgent):
147+
def __init__(self, llm_config):
148+
matmaster_core_agent = MatMasterCoreAgent(llm_config)
149+
150+
super().__init__(
151+
name='matmaster_agent',
152+
sub_agents=[matmaster_core_agent],
153+
before_agent_callback=matmaster_prepare_state,
154+
description=AgentDescription,
155+
)
156+
157+
@override
158+
async def _run_async_impl(
159+
self, ctx: InvocationContext
160+
) -> AsyncGenerator[Event, None]:
161+
async for event in super()._run_async_impl(ctx):
162+
if not ctx.session.state['should_continue']:
163+
return
164+
yield event
165+
166+
167+
def init_matmaster_agent() -> MatMasterAgent:
145168
matmaster_agent = MatMasterAgent(MatMasterLlmConfig)
146169
track_adk_agent_recursive(matmaster_agent, MatMasterLlmConfig.opik_tracer)
147170

‎agents/matmaster_agent/apex_agent/agent.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
from google.adk.tools.mcp_tool.mcp_session_manager import SseServerParams
44

55
from agents.matmaster_agent.base_agents.job_agent import BaseAsyncJobAgent
6-
from agents.matmaster_agent.constant import MATMASTER_AGENT_NAME
6+
from agents.matmaster_agent.constant import MATMASTER_CORE_AGENT_NAME
77
from agents.matmaster_agent.logger import matmodeler_logging_handler
88

99
from .constant import (
@@ -49,7 +49,7 @@ def __init__(self, llm_config):
4949
agent_instruction=ApexAgentInstruction,
5050
mcp_tools=[toolset],
5151
dflow_flag=False,
52-
supervisor_agent=MATMASTER_AGENT_NAME,
52+
supervisor_agent=MATMASTER_CORE_AGENT_NAME,
5353
)
5454

5555

‎agents/matmaster_agent/callback.py‎

Lines changed: 33 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,9 +13,10 @@
1313

1414
from agents.matmaster_agent.constant import FRONTEND_STATE_KEY, MATERIALS_ACCESS_KEY
1515
from agents.matmaster_agent.locales import i18n
16-
from agents.matmaster_agent.model import UserContent
17-
from agents.matmaster_agent.prompt import get_user_content_lang
16+
from agents.matmaster_agent.model import AgentLoop, UserContent
17+
from agents.matmaster_agent.prompt import check_should_continue, get_user_content_lang
1818
from agents.matmaster_agent.style import get_job_complete_card
19+
from agents.matmaster_agent.utils.event_utils import cherry_pick_events
1920
from agents.matmaster_agent.utils.job_utils import (
2021
get_job_status,
2122
get_running_jobs_detail,
@@ -77,6 +78,9 @@ async def matmaster_prepare_state(
7778
callback_context.state['new_query_job_status'] = callback_context.state.get(
7879
'new_query_job_status', {}
7980
)
81+
callback_context.state['should_continue'] = callback_context.state.get(
82+
'should_continue', True
83+
)
8084

8185

8286
async def matmaster_set_lang(
@@ -173,3 +177,30 @@ async def matmaster_check_job_status(
173177

174178
callback_context.state['last_llm_response_partial'] = llm_response.partial
175179
return None
180+
181+
182+
# after_agent_callback
183+
async def matmaster_should_continue(
184+
callback_context: CallbackContext,
185+
) -> Optional[types.Content]:
186+
cherry_pick_parts = cherry_pick_events(callback_context)[-1]
187+
context_messages = cherry_pick_parts[1]
188+
logger.info(f"[matmaster_should_continue] context_messages = {context_messages}")
189+
190+
prompt = check_should_continue().format(agent_response=context_messages)
191+
response = litellm.completion(
192+
model='azure/gpt-4o',
193+
messages=[{'role': 'user', 'content': prompt}],
194+
response_format=AgentLoop,
195+
)
196+
should_continue_json: dict = json.loads(response.choices[0].message.content)
197+
logger.info(
198+
f"[matmaster_should_continue] should_continue_json = {should_continue_json}"
199+
)
200+
201+
callback_context.state['should_continue'] = bool(
202+
should_continue_json['should_continue']
203+
)
204+
logger.info(
205+
f"[matmaster_should_continue] {callback_context.state['should_continue']}"
206+
)

‎agents/matmaster_agent/constant.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
TMP_FRONTEND_STATE_KEY = 'tmp_frontend_state'
99

1010
# Matmaster Constants
11-
MATMASTER_AGENT_NAME = 'matmaster_agent'
11+
MATMASTER_CORE_AGENT_NAME = 'matmaster_core_agent'
1212

1313
# Agent Constant
1414
ModelRole = 'model'

‎agents/matmaster_agent/document_parser_agent/agent.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
from google.adk.tools.mcp_tool.mcp_session_manager import SseServerParams
44

55
from agents.matmaster_agent.base_agents.job_agent import CalculationMCPLlmAgent
6-
from agents.matmaster_agent.constant import MATMASTER_AGENT_NAME, BohriumStorge
6+
from agents.matmaster_agent.constant import MATMASTER_CORE_AGENT_NAME, BohriumStorge
77

88
from .callback import validate_document_url
99
from .constant import DocumentParserAgentName, DocumentParserServerUrl
@@ -21,7 +21,7 @@ def __init__(self, llm_config):
2121
description=DocumentParserAgentDescription,
2222
instruction=DocumentParserAgentInstruction,
2323
tools=[toolset],
24-
supervisor_agent=MATMASTER_AGENT_NAME,
24+
supervisor_agent=MATMASTER_CORE_AGENT_NAME,
2525
after_model_callback=validate_document_url,
2626
)
2727

‎agents/matmaster_agent/main.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -64,8 +64,9 @@ async def agent_main() -> None:
6464
# user_input = "帮我用DPA优化这个结构：https://dp-storage-test2.oss-cn-zhangjiakou.aliyuncs.com/bohrium-test/110663/12791/store/7ba41529-5af4-4e38-a6fb-c569cd769dd9/outputs/structure_paths/structure_bulk.cif"
6565
# user_input = "帮我检索TiO2"
6666
# user_input = "请你为我搭建一个氯化钠的结构"
67-
user_input = '我想要一个bandgap 小于0.5ev的结构，空间群225，生成数量1'
67+
# user_input = '我想要一个bandgap 小于0.5ev的结构，空间群225，生成数量1'
6868
# user_input = '用openlam查找一个TiO2'
69+
user_input = '你是谁'
6970
print(f"🧑 用户：{user_input}")
7071

7172
# Create the initial content with user input

0 commit comments

Comments
 (0)