Merge pull request #181 from AnguseZhang/dev/zhouh

feat: implement internationalized job status notifications with styled JobCompleteCard and refactor MatMaster agent for improved localization

Author: Rasic2

agents/matmaster_agent/agent.py

@@ -14,16 +14,18 @@
 from agents.matmaster_agent.MrDice_agent.agent import init_MrDice_agent
 from agents.matmaster_agent.apex_agent.agent import init_apex_agent
 from agents.matmaster_agent.base_agents.io_agent import HandleFileUploadLlmAgent
-from agents.matmaster_agent.callback import matmaster_prepare_state, matmaster_check_transfer, matmaster_set_lang, \
-    matmaster_check_job_status
+from agents.matmaster_agent.callback import matmaster_prepare_state, matmaster_set_lang, matmaster_check_job_status
 from agents.matmaster_agent.chembrain_agent.agent import init_chembrain_agent
 from agents.matmaster_agent.constant import MATMASTER_AGENT_NAME, ModelRole
 from agents.matmaster_agent.document_parser_agent.agent import init_document_parser_agent
 from agents.matmaster_agent.llm_config import MatMasterLlmConfig
+from agents.matmaster_agent.model import MatMasterTargetAgentEnum
 from agents.matmaster_agent.organic_reaction_agent.agent import init_organic_reaction_agent
 from agents.matmaster_agent.perovskite_agent.agent import init_perovskite_agent
 from agents.matmaster_agent.piloteye_electro_agent.agent import init_piloteye_electro_agent
-from agents.matmaster_agent.prompt import AgentDescription, AgentInstruction, GlobalInstruction
+from agents.matmaster_agent.prompt import AgentDescription, AgentInstruction, GlobalInstruction, \
+    MatMasterCheckTransferPrompt
+from agents.matmaster_agent.public.callback import check_transfer
 from agents.matmaster_agent.ssebrain_agent.agent import init_ssebrain_agent
 from agents.matmaster_agent.structure_generate_agent.agent import init_structure_generate_agent
 from agents.matmaster_agent.superconductor_agent.agent import init_superconductor_agent
@@ -82,8 +84,9 @@ def __init__(self, llm_config):
             instruction=AgentInstruction,
             description=AgentDescription,
             before_agent_callback=[matmaster_prepare_state, matmaster_set_lang],
-            after_model_callback=[matmaster_check_job_status, matmaster_check_transfer],
-        )
+            after_model_callback=[matmaster_check_job_status,
+                                  check_transfer(prompt=MatMasterCheckTransferPrompt,
+                                                 target_agent_enum=MatMasterTargetAgentEnum)])
 
     @override
     async def _run_async_impl(self, ctx: InvocationContext) -> AsyncGenerator[Event, None]:

agents/matmaster_agent/callback.py

@@ -1,3 +1,4 @@
+import inspect
 import json
 import logging
 import uuid
@@ -12,10 +13,11 @@
 
 from agents.matmaster_agent.base_agents.callback import _get_ak
 from agents.matmaster_agent.constant import FRONTEND_STATE_KEY
-from agents.matmaster_agent.model import TransferCheck, UserContent
-from agents.matmaster_agent.prompt import get_transfer_check_prompt, get_user_content_lang
+from agents.matmaster_agent.locales import i18n
+from agents.matmaster_agent.model import UserContent
+from agents.matmaster_agent.prompt import get_user_content_lang
+from agents.matmaster_agent.style import get_job_complete_card
 from agents.matmaster_agent.utils.job_utils import get_job_status, has_job_running, get_running_jobs_detail
-from agents.matmaster_agent.utils.llm_response_utils import has_function_call
 
 logger = logging.getLogger(__name__)
 
@@ -51,7 +53,7 @@ async def matmaster_set_lang(callback_context: CallbackContext) -> Optional[type
     response = litellm.completion(model='azure/gpt-4o', messages=[{'role': 'user', 'content': prompt}],
                                   response_format=UserContent)
     result: dict = json.loads(response.choices[0].message.content)
-    logger.info(f"[matmaster_prepare_state] user_content = {result}")
+    logger.info(f"[{inspect.currentframe().f_code.co_name}] result = {result}")
     language = str(result.get('language', 'zh'))
     callback_context.state['target_language'] = language
 
@@ -74,10 +76,9 @@ async def matmaster_check_job_status(callback_context: CallbackContext, llm_resp
         running_job_ids = get_running_jobs_detail(jobs_dict)  # 从 state 里面拿
         access_key = _get_ak(callback_context)  # 从 state 或环境变量里面拿
         if callback_context.state['target_language'] in ['Chinese', 'zh-CN', '简体中文', 'Chinese (Simplified)']:
-            job_complete_intro = '检测到任务 <{job_id}> 已完成，我将立刻转移至对应的 Agent 去获取任务结果。'
+            i18n.language = 'zh'
         else:
-            job_complete_intro = ('Job <{job_id}> has been detected as completed. '
-                                  'I will immediately transfer to the corresponding agent to retrieve the job results.')
+            i18n.language = 'en'
 
         reset = False
         for origin_job_id, job_id, job_query_url, agent_name in running_job_ids:
@@ -100,7 +101,7 @@ async def matmaster_check_job_status(callback_context: CallbackContext, llm_resp
                     reset = True
                 function_call_id = f"call_{str(uuid.uuid4()).replace('-', '')[:24]}"
                 callback_context.state['origin_job_id'] = origin_job_id
-                llm_response.content.parts.append(Part(text=job_complete_intro.format(job_id=job_id)))
+                llm_response.content.parts.append(Part(text=get_job_complete_card(i18n=i18n, job_id=job_id)))
                 llm_response.content.parts.append(Part(function_call=FunctionCall(id=function_call_id,
                                                                                   name='transfer_to_agent',
                                                                                   args={'agent_name': agent_name})))
@@ -109,40 +110,3 @@ async def matmaster_check_job_status(callback_context: CallbackContext, llm_resp
 
     callback_context.state['last_llm_response_partial'] = llm_response.partial
     return None
-
-
-async def matmaster_check_transfer(callback_context: CallbackContext, llm_response: LlmResponse) -> Optional[
-    LlmResponse]:
-    # 检查响应是否有效
-    if not (
-            llm_response and
-            not llm_response.partial and
-            llm_response.content and
-            llm_response.content.parts and
-            len(llm_response.content.parts) and
-            llm_response.content.parts[0].text
-    ):
-        return None
-
-    prompt = get_transfer_check_prompt().format(response_text=llm_response.content.parts[0].text)
-    response = litellm.completion(model='azure/gpt-4o', messages=[{'role': 'user', 'content': prompt}],
-                                  response_format=TransferCheck)
-
-    result: dict = json.loads(response.choices[0].message.content)
-    is_transfer = bool(result.get('is_transfer', False))
-    target_agent = str(result.get('target_agent', ''))
-    reason = str(result.get('reason', ''))
-    logger.info(f"[matmaster_check_transfer] target_agent = {target_agent}, is_transfer = {is_transfer}"
-                f"response_text = {llm_response.content.parts[0].text}, reason = {reason}")
-    if (
-            is_transfer and
-            not has_function_call(llm_response)
-    ):
-        logger.warning(f"[matmaster_check_transfer] add `transfer_to_agent`")
-        function_call_id = f"added_{str(uuid.uuid4()).replace('-', '')[:24]}"
-        llm_response.content.parts.append(Part(function_call=FunctionCall(id=function_call_id, name='transfer_to_agent',
-                                                                          args={'agent_name': target_agent})))
-
-        return llm_response
-
-    return None

agents/matmaster_agent/locales.py

@@ -0,0 +1,18 @@
+from toolsy.i8n import I18N
+
+translations = {
+    'en': {
+        'JobStatus': 'Job Status',
+        'Job': 'Job',
+        'JobCompleted': 'Job Completed',
+        'ResultRetrieving': 'Result Retrieving'
+    },
+    'zh': {
+        'JobStatus': '任务状态',
+        'Job': '任务',
+        'JobCompleted': '任务已完成',
+        'ResultRetrieving': '结果获取中'
+    }
+}
+
+i18n = I18N(translations=translations)

agents/matmaster_agent/model.py

@@ -71,7 +71,13 @@ class DFlowJobInfo(BaseModel):
     job_in_ctx: bool = False
 
 
-class TargetAgentEnum(str, Enum):
+class ParamsCheckComplete(BaseModel):
+    flag: bool
+    reason: str
+    analyzed_messages: List[str]
+
+
+class MatMasterTargetAgentEnum(str, Enum):
     ABACUSAgent = ABACUS_AGENT_NAME
     APEXAgent = ApexAgentName
     ChemBrainAgent = CHEMBRAIN_AGENT_NAME
@@ -91,17 +97,5 @@ class TargetAgentEnum(str, Enum):
     TrajAnalysisAgent = TrajAnalysisAgentName
 
 
-class ParamsCheckComplete(BaseModel):
-    flag: bool
-    reason: str
-    analyzed_messages: List[str]
-
-
-class TransferCheck(BaseModel):
-    is_transfer: bool
-    target_agent: TargetAgentEnum
-    reason: str
-
-
 class UserContent(BaseModel):
     language: str

agents/matmaster_agent/prompt.py

@@ -710,10 +710,8 @@ def gen_params_check_info_agent_instruction():
 ResultCoreAgentDescription = 'Provides real-time task status updates and result forwarding to UI'
 TransferAgentDescription = 'Transfer to proper agent to answer user query'
 
-
 # LLM-Helper Prompt
-def get_transfer_check_prompt():
-    return """
+MatMasterCheckTransferPrompt = """
 You are an expert judge tasked with evaluating whether the previous LLM's response contains a clear and explicit request or instruction to transfer the conversation to a specific agent (e.g., 'xxx agent').
 Analyze the provided RESPONSE TEXT to determine if it explicitly indicates a transfer action.
 
@@ -801,13 +799,30 @@ def get_params_check_info_prompt():
 
 def get_user_content_lang():
     return """
-You are a professional assistant responsible for analysing language of user_content.
+You are a professional linguistic analyst. Your task is to identify the primary language used in the user content provided.
 
 User Content:
 {user_content}
 
-Provide your analysis in the following JSON format:
+Analyze the text and determine the most likely language from the following predefined options:
+- English
+- Chinese
+- Spanish
+- French
+- German
+- Japanese
+- Korean
+- Russian
+- Arabic
+- Portuguese
+- Italian
+- Dutch
+- Other
+
+If the language does not clearly match any of the above options or is a mix of multiple languages, classify it as "Other".
+
+Provide your analysis in the following strict JSON format:
 {{
-    "language": <string>
+    "language": "<string>"
 }}
 """

agents/matmaster_agent/public/__init__.py

@@ -0,0 +1,58 @@
+import json
+import logging
+import uuid
+from enum import Enum
+from typing import Optional, Callable, Type
+
+import litellm
+from google.adk.agents.callback_context import CallbackContext
+from google.adk.agents.llm_agent import AfterModelCallback
+from google.adk.models import LlmResponse
+from google.genai.types import Part, FunctionCall
+
+from agents.matmaster_agent.utils.llm_response_utils import has_function_call
+from agents.matmaster_agent.utils.model_utils import create_transfer_check_model
+
+logger = logging.getLogger(__name__)
+
+
+def check_transfer(prompt: str, target_agent_enum: Type[Enum]) -> AfterModelCallback:
+    async def wrapper(callback_context: CallbackContext, llm_response: LlmResponse) -> Optional[
+        LlmResponse]:
+        # 检查响应是否有效
+        if not (
+                llm_response and
+                not llm_response.partial and
+                llm_response.content and
+                llm_response.content.parts and
+                len(llm_response.content.parts) and
+                llm_response.content.parts[0].text
+        ):
+            return None
+
+        llm_prompt = prompt.format(response_text=llm_response.content.parts[0].text)
+        response = litellm.completion(model='azure/gpt-4o', messages=[{'role': 'user', 'content': llm_prompt}],
+                                      response_format=create_transfer_check_model(target_agent_enum))
+
+        result: dict = json.loads(response.choices[0].message.content)
+        is_transfer = bool(result.get('is_transfer', False))
+        target_agent = str(result.get('target_agent', ''))
+        reason = str(result.get('reason', ''))
+        symbol_name = f"[{callback_context.agent_name.replace('_agent', '')}_check_transfer]"
+        logger.info(f"{symbol_name} target_agent = {target_agent}, is_transfer = {is_transfer}, "
+                    f"response_text = {llm_response.content.parts[0].text}, reason = {reason}")
+        if (
+                is_transfer and
+                not has_function_call(llm_response)
+        ):
+            logger.warning(f"{symbol_name} add `transfer_to_agent`")
+            function_call_id = f"added_{str(uuid.uuid4()).replace('-', '')[:24]}"
+            llm_response.content.parts.append(
+                Part(function_call=FunctionCall(id=function_call_id, name='transfer_to_agent',
+                                                args={'agent_name': target_agent})))
+
+            return llm_response
+
+        return None
+
+    return wrapper

agents/matmaster_agent/public/callback.py

@@ -53,7 +53,7 @@ def __init__(self, llm_config):
             dflow_flag=False,
             supervisor_agent=MATMASTER_AGENT_NAME,
             sync_tools=[
-                'build_bulk_structure_by_template',
+                # 'build_bulk_structure_by_template',
                 'build_bulk_structure_by_wyckoff',
                 'make_supercell_structure',
                 'build_molecule_structure_from_g2database',

agents/matmaster_agent/structure_generate_agent/agent.py

@@ -0,0 +1,42 @@
+from toolsy.i8n import I18N
+
+
+def get_job_complete_card(i18n: I18N, job_id):
+    return f"""
+<div style="
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+    border-radius: 12px;
+    padding: 20px;
+    margin: 16px 0;
+    color: white;
+    box-shadow: 0 8px 20px rgba(102, 126, 234, 0.3);
+">
+    <div style="display: flex; align-items: center; gap: 12px; margin-bottom: 8px;">
+        <div style="
+            background: rgba(255, 255, 255, 0.2);
+            padding: 6px 12px;
+            border-radius: 20px;
+            backdrop-filter: blur(10px);
+            font-size: 12px;
+            font-weight: bold;
+        ">🚀 {i18n.t("JobStatus")}</div>
+        <div style="
+            background: rgba(255, 255, 255, 0.1);
+            padding: 4px 10px;
+            border-radius: 6px;
+            font-family: monospace;
+            font-size: 12px;
+        ">{i18n.t("Job")}ID：{job_id}</div>
+    </div>
+    <div style="
+        font-size: 16px;
+        font-weight: 500;
+        margin-top: 8px;
+    ">✅ {i18n.t("JobCompleted")}</div>
+    <div style="
+        font-size: 14px;
+        opacity: 0.9;
+        margin-top: 6px;
+    ">{i18n.t("ResultRetrieving")}...</div>
+</div>
+"""

agents/matmaster_agent/style.py

@@ -0,0 +1,12 @@
+from pydantic import create_model, BaseModel
+
+
+def create_transfer_check_model(agent_type):
+    """动态创建具有特定 agent 类型的 TransferCheck 模型"""
+    return create_model(
+        'DynamicTransferCheck',
+        is_transfer=(bool, ...),
+        target_agent=(agent_type, ...),
+        reason=(str, ...),
+        __base__=BaseModel
+    )