refactor: streamline job handling by removing URLs and adding validation functions

Author: Rasic2

agents/matmaster_agent/base_agents/callback.py

@@ -14,7 +14,7 @@
 from google.adk.agents.llm_agent import AfterToolCallback, BeforeToolCallback
 from google.adk.models import LlmResponse
 from google.adk.tools import BaseTool, ToolContext
-from mcp.types import CallToolResult
+from mcp.types import CallToolResult, TextContent
 
 from agents.matmaster_agent.constant import (
     CURRENT_ENV,
@@ -32,10 +32,10 @@
     function_calls_to_str,
     get_session_state,
     get_unique_function_call,
-    is_json,
     update_llm_response,
 )
 from agents.matmaster_agent.utils.io_oss import update_tgz_dict
+from agents.matmaster_agent.utils.tool_response_utils import check_valid_tool_response
 
 logger = logging.getLogger(__name__)
 
@@ -457,12 +457,7 @@ async def wrapper(
             raise TypeError('Not CalculationMCPTool type')
 
         # 检查是否为有效的 json 字典
-        if not (
-            tool_response
-            and tool_response.content
-            and tool_response.content[0].text
-            and is_json(tool_response.content[0].text)
-        ):
+        if not check_valid_tool_response(tool_response):
             return None
 
         tool_result = json.loads(tool_response.content[0].text)
@@ -473,6 +468,44 @@ async def wrapper(
     return wrapper
 
 
+def remove_job_link(func: AfterToolCallback) -> AfterToolCallback:
+    @wraps(func)
+    async def wrapper(
+        tool: BaseTool,
+        args: dict,
+        tool_context: ToolContext,
+        tool_response: Union[dict, CallToolResult],
+    ) -> Optional[dict]:
+        # 两步操作：
+        # 1. 调用被装饰的 after_tool_callback；
+        # 2. 如果调用的 after_tool_callback 有返回值，以这个为准
+        # 如果 tool 为 Transfer2Agent，直接 return
+        if tool.name == Transfer2Agent:
+            return None
+
+        if (
+            after_tool_result := await func(tool, args, tool_context, tool_response)
+        ) is not None:
+            return after_tool_result
+
+        # 检查是否为有效的 json 字典
+        if not check_valid_tool_response(tool_response):
+            return None
+
+        # 移除 job_link
+        tool_result: dict = json.loads(tool_response.content[0].text)
+        if tool_result.get('extra_info', None) is not None:
+            del tool_result['extra_info']['job_link']
+            tool_response.content[0] = TextContent(
+                type='text', text=json.dumps(tool_result)
+            )
+
+            logger.info(f"[remove_job_link] final_tool_result = {tool_response}")
+            return tool_response
+
+    return wrapper
+
+
 def catch_after_tool_callback_error(func: AfterToolCallback) -> AfterToolCallback:
     @wraps(func)
     async def wrapper(

agents/matmaster_agent/base_agents/job_agent.py

@@ -23,6 +23,7 @@
     inject_current_env,
     inject_username_ticket,
     remove_function_call,
+    remove_job_link,
     tgz_oss_to_oss_list,
 )
 from agents.matmaster_agent.base_agents.io_agent import HandleFileUploadLlmAgent
@@ -39,7 +40,6 @@
     MATERIALS_PROJECT_ID,
     TMP_FRONTEND_STATE_KEY,
     ModelRole,
-    OpenAPIJobAPI,
     get_BohriumExecutor,
     get_BohriumStorage,
     get_DFlowExecutor,
@@ -166,7 +166,9 @@ def __init__(
         )
         after_tool_callback = check_before_tool_callback_effect(
             catch_after_tool_callback_error(
-                tgz_oss_to_oss_list(after_tool_callback, enable_tgz_unpack)
+                remove_job_link(
+                    tgz_oss_to_oss_list(after_tool_callback, enable_tgz_unpack)
+                )
             )
         )
 
@@ -412,12 +414,10 @@ async def _run_async_impl(
                             job_status = results['status']
                             if not ctx.session.state['dflow']:
                                 bohr_job_id = results['extra_info']['bohr_job_id']
-                                job_query_url = f'{OpenAPIJobAPI}/{bohr_job_id}'
                                 frontend_result = BohrJobInfo(
                                     origin_job_id=origin_job_id,
                                     job_name=job_name,
                                     job_status=job_status,
-                                    job_query_url=job_query_url,
                                     job_id=bohr_job_id,
                                     agent_name=ctx.agent.parent_agent.parent_agent.name,
                                 ).model_dump(mode='json')

agents/matmaster_agent/callback.py

@@ -124,13 +124,13 @@ async def matmaster_check_job_status(
             i18n.language = 'en'
 
         reset = False
-        for origin_job_id, job_id, job_query_url, agent_name in running_job_ids:
+        for origin_job_id, job_id, agent_name in running_job_ids:
             if not callback_context.state['last_llm_response_partial']:
                 logger.info(
                     '[matmaster_check_job_status] new LlmResponse, prepare call API'
                 )
                 job_status = await get_job_status(
-                    job_query_url, access_key=MATERIALS_ACCESS_KEY
+                    job_id, access_key=MATERIALS_ACCESS_KEY
                 )  # 查询任务的最新状态
                 callback_context.state['new_query_job_status'][
                     'origin_job_id'

agents/matmaster_agent/main.py

@@ -60,11 +60,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/model.py

@@ -52,7 +52,6 @@ class JobResult(BaseModel):
 class BohrJobInfo(BaseModel):
     origin_job_id: str
     job_id: Union[int, str]
-    job_query_url: str
     job_status: JobStatus
     job_name: str
     job_result: Optional[List[JobResult]] = None

agents/matmaster_agent/utils/job_utils.py

@@ -1,5 +1,6 @@
 import aiohttp
 
+from agents.matmaster_agent.constant import OpenAPIJobAPI
 from agents.matmaster_agent.model import BohrJobInfo, JobStatus
 
 
@@ -19,10 +20,10 @@ def mapping_status(status):
     }.get(status, 'Unknown')
 
 
-async def get_job_status(job_query_url, access_key):
+async def get_job_status(job_id, access_key):
     async with aiohttp.ClientSession() as session:
         async with session.get(
-            job_query_url, headers={'accessKey': access_key}
+            f'{OpenAPIJobAPI}/{job_id}', headers={'accessKey': access_key}
         ) as response:
             data = await response.json()
             return mapping_status(data['data']['status'])
@@ -35,7 +36,7 @@ def has_job_running(jobs_dict: BohrJobInfo) -> bool:
 
 def get_running_jobs_detail(jobs_dict: BohrJobInfo):
     return [
-        (job['origin_job_id'], job['job_id'], job['job_query_url'], job['agent_name'])
+        (job['origin_job_id'], job['job_id'], job['agent_name'])
         for job in jobs_dict.values()
         if job['job_status'] == JobStatus.Running
     ]

agents/matmaster_agent/utils/tool_response_utils.py

@@ -0,0 +1,10 @@
+from agents.matmaster_agent.utils.helper_func import is_json
+
+
+def check_valid_tool_response(tool_response):
+    return (
+        tool_response
+        and tool_response.content
+        and tool_response.content[0].text
+        and is_json(tool_response.content[0].text)
+    )