Merge pull request #168 from AnguseZhang/dev/zhouh

refactor: streamline agent transition handling and enhance session tr…

Author: Rasic2

agents/matmaster_agent/agent.py

@@ -90,9 +90,8 @@ async def _run_async_impl(self, ctx: InvocationContext) -> AsyncGenerator[Event,
                 # 对于 [matmaster_check_job_status] 生成的消息， 手动拼一个流式消息
                 if ctx.session.state['special_llm_response']:
                     yield frontend_text_event(ctx, self.name, event.content.parts[0].text, ModelRole)
-
-                ctx.session.state['special_llm_response'] = False
-                await update_session_state(ctx, self.name)
+                    ctx.session.state['special_llm_response'] = False
+                    await update_session_state(ctx, self.name)
                 yield event
         except BaseException as err:
             async for error_event in send_error_event(err, ctx, self.name):

agents/matmaster_agent/callback.py

@@ -64,10 +64,10 @@ async def matmaster_check_job_status(callback_context: CallbackContext, llm_resp
         running_job_ids = get_running_jobs_detail(jobs_dict)
         access_key = _get_ak(callback_context)
         if callback_context.state['target_language'] in ['Chinese', 'zh-CN', '简体中文', 'Chinese (Simplified)']:
-            job_complete_intro = '检测到任务 <{job_id}> 已完成，我将立刻转移至对应的 Agent 去获取任务结果，请稍等...'
+            job_complete_intro = '检测到任务 <{job_id}> 已完成，我将立刻转移至对应的 Agent 去获取任务结果。'
         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. Please wait...')
+                                  'I will immediately transfer to the corresponding agent to retrieve the job results.')
 
         reset = False
         for origin_job_id, job_id, job_query_url, agent_name in running_job_ids:
@@ -116,7 +116,7 @@ async def matmaster_check_transfer(callback_context: CallbackContext, llm_respon
             is_transfer and
             not has_function_call(llm_response)
     ):
-        logger.warning(f"Detected Agent Transfer Hallucination, add `transfer_to_agent`")
+        logger.warning(f"[matmaster_check_transfer] target_agent = {target_agent}")
         function_call_id = f"call_{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})))

agents/matmaster_agent/model.py

@@ -3,6 +3,23 @@
 
 from pydantic import BaseModel, field_validator
 
+from agents.matmaster_agent.ABACUS_agent.constant import ABACUS_AGENT_NAME
+from agents.matmaster_agent.DPACalculator_agent.constant import DPACalulator_AGENT_NAME
+from agents.matmaster_agent.HEACalculator_agent.constant import HEACALCULATOR_AGENT_NAME
+from agents.matmaster_agent.HEA_assistant_agent.constant import HEA_assistant_AgentName
+from agents.matmaster_agent.INVAR_agent.constant import INVAR_AGENT_NAME
+from agents.matmaster_agent.MrDice_agent.constant import MrDice_Agent_Name
+from agents.matmaster_agent.apex_agent.constant import ApexAgentName
+from agents.matmaster_agent.chembrain_agent.constant import CHEMBRAIN_AGENT_NAME
+from agents.matmaster_agent.organic_reaction_agent.constant import ORGANIC_REACTION_AGENT_NAME
+from agents.matmaster_agent.perovskite_agent.constant import PerovskiteAgentName
+from agents.matmaster_agent.piloteye_electro_agent.constant import PILOTEYE_ELECTRO_AGENT_NAME
+from agents.matmaster_agent.ssebrain_agent.constant import SSEBRAIN_AGENT_NAME
+from agents.matmaster_agent.structure_generate_agent.constant import StructureGenerateAgentName
+from agents.matmaster_agent.superconductor_agent.constant import SuperconductorAgentName
+from agents.matmaster_agent.thermoelectric_agent.constant import ThermoelectricAgentName
+from agents.matmaster_agent.traj_analysis_agent.constant import TrajAnalysisAgentName
+
 
 class JobStatus(str, Enum):
     Running = 'Running'
@@ -53,9 +70,28 @@ class DFlowJobInfo(BaseModel):
     job_in_ctx: bool = False
 
 
+class TargetAgentEnum(str, Enum):
+    ABACUSAgent = ABACUS_AGENT_NAME
+    APEXAgent = ApexAgentName
+    ChemBrainAgent = CHEMBRAIN_AGENT_NAME
+    DPACalculatorAgent = DPACalulator_AGENT_NAME
+    HEAAssistantAgent = HEA_assistant_AgentName
+    HEACalculatorAgent = HEACALCULATOR_AGENT_NAME
+    INVARAgent = INVAR_AGENT_NAME
+    MrDiceAgent = MrDice_Agent_Name
+    OrganicReactionAgent = ORGANIC_REACTION_AGENT_NAME
+    PerovskiteAgent = PerovskiteAgentName
+    PiloteyeElectroAgent = PILOTEYE_ELECTRO_AGENT_NAME
+    SSEBrainAgent = SSEBRAIN_AGENT_NAME
+    StructureGenerateAgent = StructureGenerateAgentName
+    SuperConductorAgent = SuperconductorAgentName
+    ThermoElectricAgent = ThermoelectricAgentName
+    TrajAnalysisAgent = TrajAnalysisAgentName
+
+
 class TransferCheck(BaseModel):
     is_transfer: bool
-    target_agent: str
+    target_agent: TargetAgentEnum
 
 
 class UserContent(BaseModel):

agents/matmaster_agent/utils/helper_func.py

@@ -1,6 +1,8 @@
 import copy
+import inspect
 import json
 import logging
+import os
 from typing import List, Union
 
 from google.adk.agents.invocation_context import InvocationContext
@@ -19,9 +21,14 @@ def get_session_state(ctx: Union[InvocationContext, ToolContext]):
 
 
 async def update_session_state(ctx: InvocationContext, author: str):
+    stack = inspect.stack()
+    frame = stack[1]  # stack[1] 表示调用当前函数的上一层调用
+    filename = os.path.basename(frame.filename)
+    lineno = frame.lineno
     actions_with_update = EventActions(state_delta=ctx.session.state)
-    system_event = Event(invocation_id=ctx.invocation_id, author=author, actions=actions_with_update)
-    await ctx.session_service.append_event(ctx.session, system_event)
+    system_event = Event(invocation_id=ctx.invocation_id, author=f"{filename}:{lineno}",
+                         actions=actions_with_update)
+    await ctx.session_service.append_event(ctx.session, system_event)  # 会引入一个空消息
 
 
 def update_llm_response(llm_response: LlmResponse, current_function_calls: List[dict],