refactor: streamline matmaster_agent transfer logic by abstracting callback functions and modifying the agent model structure

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

@@ -12,10 +12,9 @@
 
 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.model import UserContent
+from agents.matmaster_agent.prompt import get_user_content_lang
 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__)
 
@@ -109,40 +108,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/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

@@ -715,10 +715,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.
 

agents/matmaster_agent/public/__init__.py

@@ -0,0 +1 @@
+

agents/matmaster_agent/public/callback.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/utils/model_utils.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
+    )