Merge pull request #167 from AnguseZhang/dev/zhouh

feat: integrate session state update and frontend text event handling…

Author: Rasic2

agents/matmaster_agent/agent.py

@@ -17,7 +17,7 @@
 from agents.matmaster_agent.callback import matmaster_prepare_state, matmaster_check_transfer, 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
+from agents.matmaster_agent.constant import MATMASTER_AGENT_NAME, ModelRole
 from agents.matmaster_agent.llm_config import MatMasterLlmConfig
 from agents.matmaster_agent.organic_reaction_agent.agent import init_organic_reaction_agent
 from agents.matmaster_agent.perovskite_agent.agent import init_perovskite_agent
@@ -28,7 +28,8 @@
 from agents.matmaster_agent.superconductor_agent.agent import init_superconductor_agent
 from agents.matmaster_agent.thermoelectric_agent.agent import init_thermoelectric_agent
 from agents.matmaster_agent.traj_analysis_agent.agent import init_traj_analysis_agent
-from agents.matmaster_agent.utils.event_utils import send_error_event
+from agents.matmaster_agent.utils.event_utils import send_error_event, frontend_text_event
+from agents.matmaster_agent.utils.helper_func import update_session_state
 
 logging.getLogger('google_adk.google.adk.tools.base_authenticated_tool').setLevel(logging.ERROR)
 
@@ -86,6 +87,12 @@ async def _run_async_impl(self, ctx: InvocationContext) -> AsyncGenerator[Event,
         try:
             # Delegate to parent implementation for the actual processing
             async for event in super()._run_async_impl(ctx):
+                # 对于 [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)
                 yield event
         except BaseException as err:
             async for error_event in send_error_event(err, ctx, self.name):

agents/matmaster_agent/callback.py

@@ -8,7 +8,7 @@
 from google.adk.agents.callback_context import CallbackContext
 from google.adk.models import LlmResponse
 from google.genai import types
-from google.genai.types import FunctionCall, Part, FunctionResponse
+from google.genai.types import FunctionCall, Part
 
 from agents.matmaster_agent.base_agents.callback import _get_ak
 from agents.matmaster_agent.constant import FRONTEND_STATE_KEY
@@ -40,6 +40,7 @@ async def matmaster_prepare_state(callback_context: CallbackContext) -> Optional
     callback_context.state['sync_tools'] = callback_context.state.get('sync_tools', None)
     callback_context.state['invocation_id_with_tool_call'] = callback_context.state.get('invocation_id_with_tool_call',
                                                                                         None)
+    callback_context.state['special_llm_response'] = False
 
 
 async def matmaster_set_lang(callback_context: CallbackContext) -> Optional[types.Content]:
@@ -72,7 +73,11 @@ async def matmaster_check_job_status(callback_context: CallbackContext, llm_resp
         for origin_job_id, job_id, job_query_url, agent_name in running_job_ids:
             job_status = get_job_status(job_query_url, access_key=access_key)
             if job_status in ['Failed', 'Finished']:
+                if llm_response.partial:  # 原来消息的流式版本置空 None
+                    llm_response.content = None
+                    break
                 if not reset:
+                    callback_context.state['special_llm_response'] = True  # 标记开始处理原来消息的非流式版本
                     llm_response.content.parts = []
                     reset = True
                 logger.info(f"[matmaster_check_job_status] job_id = {job_id}, job_status = {job_status}")