Merge pull request #241 from AnguseZhang/feat/ticket-auth

feat: Enhance job tracking and processing with URL functionality, improved error handling, and tool call management

Author: Rasic2

agents/matmaster_agent/agent.py

@@ -11,6 +11,7 @@
 from agents.matmaster_agent.base_agents.io_agent import HandleFileUploadLlmAgent
 from agents.matmaster_agent.callback import (
     matmaster_check_job_status,
+    matmaster_hallucination_retry,
     matmaster_prepare_state,
     matmaster_set_lang,
 )
@@ -110,6 +111,7 @@ def __init__(self, llm_config):
                     prompt=MatMasterCheckTransferPrompt,
                     target_agent_enum=MatMasterTargetAgentEnum,
                 ),
+                matmaster_hallucination_retry,
             ],
         )
 

agents/matmaster_agent/base_agents/job_agent.py

@@ -15,7 +15,7 @@
 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.style import get_job_complete_card, hallucination_card
 from agents.matmaster_agent.utils.job_utils import (
     get_job_status,
     get_running_jobs_detail,
@@ -77,6 +77,12 @@ async def matmaster_prepare_state(
     callback_context.state['new_query_job_status'] = callback_context.state.get(
         'new_query_job_status', {}
     )
+    callback_context.state['hallucination'] = callback_context.state.get(
+        'hallucination', False
+    )
+    callback_context.state['hallucination_agent'] = callback_context.state.get(
+        'hallucination_agent', None
+    )
 
 
 async def matmaster_set_lang(
@@ -93,6 +99,15 @@ async def matmaster_set_lang(
     logger.info(f"[{inspect.currentframe().f_code.co_name}] result = {result}")
     language = str(result.get('language', 'zh'))
     callback_context.state['target_language'] = language
+    if callback_context.state['target_language'] in [
+        'Chinese',
+        'zh-CN',
+        '简体中文',
+        'Chinese (Simplified)',
+    ]:
+        i18n.language = 'zh'
+    else:
+        i18n.language = 'en'
 
 
 # after_model_callback
@@ -113,16 +128,6 @@ async def matmaster_check_job_status(
         jobs_dict
     ):  # 确认当前有在运行中的任务
         running_job_ids = get_running_jobs_detail(jobs_dict)  # 从 state 里面拿
-        if callback_context.state['target_language'] in [
-            'Chinese',
-            'zh-CN',
-            '简体中文',
-            'Chinese (Simplified)',
-        ]:
-            i18n.language = 'zh'
-        else:
-            i18n.language = 'en'
-
         reset = False
         for origin_job_id, job_id, agent_name in running_job_ids:
             if not callback_context.state['last_llm_response_partial']:
@@ -149,9 +154,8 @@ async def matmaster_check_job_status(
                     llm_response.content = None
                     break
                 if not reset:
-                    callback_context.state['special_llm_response'] = (
-                        True  # 标记开始处理原来消息的非流式版本
-                    )
+                    # 标记开始处理原来消息的非流式版本
+                    callback_context.state['special_llm_response'] = True
                     llm_response.content.parts = []
                     reset = True
                 function_call_id = f"call_{str(uuid.uuid4()).replace('-', '')[:24]}"
@@ -169,7 +173,42 @@ async def matmaster_check_job_status(
                     )
                 )
         callback_context.state['last_llm_response_partial'] = llm_response.partial
-        return llm_response
+        # return llm_response
 
     callback_context.state['last_llm_response_partial'] = llm_response.partial
-    return None
+    return
+
+
+async def matmaster_hallucination_retry(
+    callback_context: CallbackContext, llm_response: LlmResponse
+) -> Optional[LlmResponse]:
+    hallucination_flag = callback_context.state['hallucination']
+    hallucination_agent = callback_context.state['hallucination_agent']
+    logger.info(
+        f'[matmaster_hallucination_retry] hallucination_flag={hallucination_flag}, hallucination_agent={hallucination_agent}, i18n.language = {i18n.language}'
+    )
+    if not callback_context.state['hallucination']:
+        return
+
+    if llm_response.partial:  # 原来消息的流式版本置空 None
+        llm_response.content = None
+        return
+
+    # 标记开始处理原来消息的非流式版本
+    callback_context.state['special_llm_response'] = True
+    llm_response.content.parts = []
+
+    llm_response.content.parts.append(Part(text=hallucination_card(i18n=i18n)))
+    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': callback_context.state['hallucination_agent']},
+            )
+        )
+    )
+    callback_context.state['hallucination'] = False
+
+    return llm_response

agents/matmaster_agent/callback.py

@@ -83,14 +83,13 @@ def _init_model(
             model_name: str,
             **kwargs,
         ):
-            llm_kwargs = {'stream_options': {'include_usage': True}}
-            if kwargs:
-                llm_kwargs.update(kwargs)
             return LiteLlm(
                 model=MODEL_MAPPING.get((provider_key, model_name), DEFAULT_MODEL),
-                **llm_kwargs,
+                **kwargs,
             )
 
+        llm_kwargs = {'stream_options': {'include_usage': True}}
+
         self.gpt_4o_mini = _init_model(azure_provider, gpt_4o_mini)
         self.gpt_4o = _init_model(azure_provider, gpt_4o)
         self.gemini_2_0_flash = _init_model(litellm_provider, gemini_2_0_flash)
@@ -103,7 +102,7 @@ def _init_model(
         self.gpt_5 = _init_model(litellm_provider, gpt_5)
         self.gpt_5_nano = _init_model(litellm_provider, gpt_5_nano)
         self.gpt_5_mini = _init_model(litellm_provider, gpt_5_mini)
-        self.gpt_5_chat = _init_model(litellm_provider, gpt_5_chat)
+        self.gpt_5_chat = _init_model(litellm_provider, gpt_5_chat, **llm_kwargs)
 
         # tracing
         self.opik_tracer = OpikTracer()

agents/matmaster_agent/llm_config.py

@@ -6,12 +6,16 @@
         'Job': 'Job',
         'JobCompleted': 'Job Completed',
         'ResultRetrieving': 'Result Retrieving',
+        'JobSubmitHallucination': 'Job Submit Hallucination',
+        'JobSubmitHallucinationAction': 'Detected a task submission hallucination, I will retry the submission with the same parameters.',
     },
     'zh': {
         'JobStatus': '任务状态',
         'Job': '任务',
         'JobCompleted': '任务已完成',
         'ResultRetrieving': '结果获取中',
+        'JobSubmitHallucination': '任务提交幻觉',
+        'JobSubmitHallucinationAction': '检测到任务提交幻觉，我将使用相同参数重试提交。',
     },
 }
 

agents/matmaster_agent/locales.py

@@ -64,8 +64,9 @@ async def agent_main() -> None:
     # 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'
+    user_input = '对 NbSe₂ 超导体进行声子谱计算并结合电子–声子耦合分析估算临界温度 Tc（从开源数据库获取初始结构并自行设定参数），并以 URL 导出声子谱与 e–ph 计算结果'
     print(f"🧑 用户：{user_input}")
 
     # Create the initial content with user input

agents/matmaster_agent/main.py

@@ -49,6 +49,7 @@
    - Upon user confirmation or applied parameters, execute the step using the sub-agent.
 6. **Result Handling**:
    - Present the execution result and a brief analysis.
+   - If the result contains images in markdown format, display them to the user using proper markdown syntax.
    - Await user instruction: either proceed to the next step in the plan, adjust parameters, or modify the plan.
 
 **Response Formatting:**
@@ -720,6 +721,44 @@ def gen_params_check_info_agent_instruction():
 """
 
 
+def gen_tool_call_info_instruction():
+    return """
+You are an AI agent that matches user requests to available tools. Your task is to analyze the user's query and return a JSON object with the following structure:
+{{
+  "tool_name": "string",
+  "tool_args": {{"param1_name": "value1", "param2_name": "value2"}},
+  "missing_tool_args": ["param3_name", "param4_name"]
+}}
+
+**Key Rules:**
+- The `tool_args` object should contain parameter names as keys and the actual values extracted from the user's request as values
+- For parameters where values cannot be extracted from the user's request, include the parameter name in the `missing_tool_args` list
+- If any parameter involves an input file, the parameter name should indicate it requires an HTTP URL (e.g., "file_url", "image_url")
+- For output file parameters, use appropriate names (e.g., "output_path", "result_file") - these will handle OSS URLs automatically
+- Only return the JSON object - do not execute any tools directly
+- Extract and include all available parameter values from the user's request in `tool_args`
+- List all missing required parameter names in the `missing_tool_args` array
+
+**Example Response:**
+{{
+  "tool_name": "image_processor",
+  "tool_args": {{
+    "image_url": "https://example.com/image.jpg",
+    "operation": "resize",
+    "width": 800
+  }},
+  "missing_tool_args": ["height", "output_format"]
+}}
+
+**Constraints:**
+- Return only valid JSON - no additional text or explanations
+- Include all available parameter values from the user's request in `tool_args`
+- List all missing required parameter names in `missing_tool_args`
+- Match the tool precisely based on the user's request
+- If no suitable tool is found, return an empty object: {{}}
+"""
+
+
 SubmitRenderAgentDescription = 'Sends specific messages to the frontend for rendering dedicated task list components'
 
 ResultCoreAgentDescription = (

agents/matmaster_agent/prompt.py

@@ -40,3 +40,77 @@ def get_job_complete_card(i18n: I18N, job_id):
     ">{i18n.t("ResultRetrieving")}...</div>
 </div>
 """
+
+
+def hallucination_card(i18n: I18N):
+    def _inner_css():
+        return """
+    <style>
+        @keyframes pulse {
+            0% { transform: scale(1); }
+            50% { transform: scale(1.1); }
+            100% { transform: scale(1); }
+        }
+        @keyframes progress {
+            0% { transform: translateX(-100%); }
+            100% { transform: translateX(250%); }
+        }
+    </style>
+</div>
+"""
+
+    return (
+        f"""
+<div style="
+    background: linear-gradient(135deg, #6a11cb 0%, #2575fc 100%);
+    color: white;
+    padding: 24px 32px;
+    border-radius: 12px;
+    box-shadow: 0 10px 30px rgba(0, 0, 0, 0.15);
+    max-width: 500px;
+    width: 100%;
+    position: relative;
+    overflow: hidden;
+    transition: transform 0.3s ease, box-shadow 0.3s ease;
+    margin: 20px auto;
+">
+    <div style="
+        position: relative;
+        z-index: 1;
+        display: flex;
+        align-items: center;
+    ">
+        <div style="
+            font-size: 28px;
+            margin-right: 16px;
+            animation: pulse 2s infinite;
+        ">🔄</div>
+        <div style="flex: 1;">
+            <div style="
+                font-size: 18px;
+                font-weight: 600;
+                margin-bottom: 4px;
+            ">{i18n.t("JobSubmitHallucination")}</div>
+            <div style="
+                font-size: 14px;
+                opacity: 0.9;
+            ">{i18n.t("JobSubmitHallucinationAction")}</div>
+        </div>
+    </div>
+    <div style="
+        height: 4px;
+        background: rgba(255, 255, 255, 0.3);
+        border-radius: 2px;
+        margin-top: 16px;
+        overflow: hidden;
+    ">
+        <div style="
+            height: 100%;
+            width: 60%;
+            background: white;
+            border-radius: 2px;
+            animation: progress 2s ease-in-out infinite;
+        "></div>
+    </div>"""
+        + _inner_css()
+    )

agents/matmaster_agent/style.py

@@ -42,12 +42,6 @@ def is_text(event: Event):
     return has_part(event) and event.content.parts[0].text
 
 
-def is_text_and_not_bohrium(event: Event):
-    return is_text(event) and not event.content.parts[0].text.startswith(
-        '<bohrium-chat-msg>'
-    )
-
-
 def is_function_call(event: Event) -> bool:
     """检查事件是否包含函数调用"""
     return has_part(event) and any(part.function_call for part in event.content.parts)