Merge pull request #175 from AnguseZhang/dev/zhouh

refactor: replace params_check_completed_agent with dynamic prompt in…

Author: Rasic2

agents/matmaster_agent/base_agents/job_agent.py

@@ -4,12 +4,11 @@
 from typing import AsyncGenerator, override, Optional
 
 import jsonpickle
+import litellm
 from google.adk.agents import LlmAgent, SequentialAgent
 from google.adk.agents.invocation_context import InvocationContext
 from google.adk.events import Event, EventActions
-from mcp.types import CallToolResult
 from pydantic import Field, BaseModel
-from yaml.scanner import ScannerError
 
 from agents.matmaster_agent.base_agents.callback import check_before_tool_callback_effect, default_before_tool_callback, \
     catch_before_tool_callback_error, check_job_create, inject_username_ticket, inject_ak_projectId, \
@@ -33,7 +32,6 @@
     get_BohriumStorage,
     get_DFlowExecutor, OpenAPIJobAPI,
 )
-from agents.matmaster_agent.llm_config import MatMasterLlmConfig
 from agents.matmaster_agent.model import BohrJobInfo, DFlowJobInfo
 from agents.matmaster_agent.prompt import (
     ResultCoreAgentDescription,
@@ -43,7 +41,7 @@
 )
 from agents.matmaster_agent.utils.event_utils import is_function_call, is_function_response, send_error_event, is_text, \
     context_function_event, all_text_event, context_text_event, frontend_text_event, is_text_and_not_bohrium, \
-    get_function_call_indexes, context_multipart2function_event
+    get_function_call_indexes, context_multipart2function_event, cherry_pick_events
 from agents.matmaster_agent.utils.frontend import get_frontend_job_result_data
 from agents.matmaster_agent.utils.helper_func import update_session_state, parse_result, get_session_state, \
     load_tool_response
@@ -528,7 +526,6 @@ async def _run_async_impl(self, ctx: InvocationContext) -> AsyncGenerator[Event,
 class BaseAsyncJobAgent(LlmAgent):
     submit_agent: SequentialAgent
     result_agent: SequentialAgent
-    params_check_completed_agent: LlmAgent
     params_check_info_agent: LlmAgent
     dflow_flag: bool = Field(False, description='Whether the agent is dflow related', exclude=True)
     supervisor_agent: str
@@ -594,16 +591,6 @@ def __init__(
             sub_agents=[result_core_agent]
         )
 
-        params_check_completed_agent = ParamsCheckCompletedAgent(
-            model=MatMasterLlmConfig.gpt_4o,
-            # model=model,
-            name=f"{agent_prefix}_params_check_completed_agent",
-            instruction=gen_params_check_completed_agent_instruction(),
-            output_schema=ParamsCheckComplete,
-            disallow_transfer_to_parent=True,
-            disallow_transfer_to_peers=True
-        )
-
         params_check_info_agent = ParamsCheckInfoAgent(
             model=model,
             name=f"{agent_prefix}_params_check_info_agent",
@@ -621,10 +608,9 @@ def __init__(
             description=agent_description,
             submit_agent=submit_agent,
             result_agent=result_agent,
-            params_check_completed_agent=params_check_completed_agent,
             params_check_info_agent=params_check_info_agent,
             dflow_flag=dflow_flag,
-            sub_agents=[submit_agent, result_agent, params_check_completed_agent, params_check_info_agent],
+            sub_agents=[submit_agent, result_agent, params_check_info_agent],
             supervisor_agent=supervisor_agent,
             sync_tools=sync_tools,
             enable_tgz_unpack=enable_tgz_unpack
@@ -651,10 +637,14 @@ async def _run_async_impl(self, ctx: InvocationContext) -> AsyncGenerator[Event,
         ):  # Only Query Job Result
             pass
         else:
-            last_params_check_completed_event = None
-            async for params_check_completed_event in self.params_check_completed_agent.run_async(ctx):
-                last_params_check_completed_event = params_check_completed_event
-            params_check_completed_json = json.loads(last_params_check_completed_event.content.parts[0].text)
+            cherry_pick_parts = cherry_pick_events(ctx)[:5]
+            context_messages = '\n'.join([f'<{item[0].title()}> said: \n{item[1]}\n' for item in cherry_pick_parts])
+            logger.info(f"[BaseAsyncJobAgent] context_messages = {context_messages}")
+
+            prompt = gen_params_check_completed_agent_instruction().format(context_messages=context_messages)
+            response = litellm.completion(model='azure/gpt-4o', messages=[{'role': 'user', 'content': prompt}],
+                                          response_format=ParamsCheckComplete)
+            params_check_completed_json: dict = json.loads(response.choices[0].message.content)
             params_check_completed = params_check_completed_json['flag']
             params_check_reason = params_check_completed_json['reason']
             params_check_msg = params_check_completed_json['analyzed_message']

agents/matmaster_agent/main.py

@@ -64,11 +64,11 @@ async def agent_main() -> None:
     # user_input = "高熵合金AlCoCr0.5FeNi2的可能结构是什么"
     # user_input = "调用 thermoelectric_agent 帮我生成10个具有Sn和Te元素的热电结构"
     # user_input = "plot perovstite 2021 to 2025"
-    # user_input = "帮我创建一个 FCC Bulk Cu 的结构"
+    user_input = '帮我创建一个 FCC Bulk Cu 的结构'
     # 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"
     # user_input = "帮我检索TiO2"
     # user_input = "请你为我搭建一个氯化钠的结构"
-    user_input = '我想要一个bandgap 小于0.5ev的结构，空间群225，生成数量1'
+    # user_input = '我想要一个bandgap 小于0.5ev的结构，空间群225，生成数量1'
     # user_input = '用openlam查找一个TiO2'
     print(f"🧑 用户：{user_input}")
 

agents/matmaster_agent/prompt.py

@@ -653,8 +653,10 @@ def gen_result_agent_description():
 
 def gen_params_check_completed_agent_instruction():
     return """
-Analyze the most recent message where the 'author' field ends with '_agent' and which contains parameter information. This message may not be the immediate preceding one.
 Your task is to determine if the parameters requiring user confirmation have been fully presented and a confirmation is being requested in that message.
+Analyze the messages from `user` and `model` listed below (only listed Latest 5 messages):
+
+{context_messages}
 
 Your output MUST be a valid JSON object with the following structure:
 {{

agents/matmaster_agent/utils/event_utils.py

@@ -150,7 +150,7 @@ def context_multipart2function_event(ctx: InvocationContext, author: str, event:
     for part in event.content.parts:
         if part.text:
             yield from context_function_event(ctx, author, function_call_name, {'msg': part.text},
-                                                         ModelRole)
+                                              ModelRole)
         elif part.function_call:
             yield context_function_call_event(ctx, author, function_call_id=part.function_call.id,
                                               function_call_name=part.function_call.name, role=ModelRole,
@@ -163,6 +163,18 @@ def all_text_event(ctx: InvocationContext, author: str, text: str, role: str):
     yield context_text_event(ctx, author, text, role)
 
 
+def cherry_pick_events(ctx: InvocationContext):
+    events = ctx.session.events
+    cherry_pick_parts = []
+    for event in events:
+        if event.content:
+            for part in event.content.parts:
+                if part.text:
+                    cherry_pick_parts.append((event.content.role, part.text))
+
+    return cherry_pick_parts
+
+
 async def send_error_event(err, ctx: InvocationContext, author):
     ctx.session.state['error_occurred'] = True
     await update_session_state(ctx, author)