feat: add memory modules

agents/matmaster_agent/constant.py

@@ -49,6 +49,7 @@
 BOHRIUM_COM = f'https://www{URL_PART}.bohrium.com'
 BOHRIUM_HOST = f'https://bohrium{URL_PART}.dp.tech'
 ICL_SERVICE_URL = '101.126.90.82:8001'
+MEMORY_SERVICE_URL = '101.126.90.82:8002'
 MATMASTER_TOOLS_SERVER = f'https://matmaster-tools-server{URL_PART}.bohrium.com'
 if CURRENT_ENV == 'test':
     DFLOW_HOST = 'https://lbg-workflow-mlops.test.dp.tech'

agents/matmaster_agent/core_agents/base_agents/mcp_agent.py

@@ -13,6 +13,7 @@
     check_user_phonon_balance,
     default_after_model_callback,
     default_after_tool_callback,
+    default_before_model_callback,
     default_before_tool_callback,
     default_cost_func,
     filter_function_calls,
@@ -42,6 +43,12 @@
 from agents.matmaster_agent.flow_agents.model import PlanStepStatusEnum
 from agents.matmaster_agent.locales import i18n
 from agents.matmaster_agent.logger import PrefixFilter
+from agents.matmaster_agent.memory.inject_memory_callback import (
+    inject_memory_before_model,
+)
+from agents.matmaster_agent.memory.store_tool_result_callback import (
+    store_tool_result_in_memory,
+)
 from agents.matmaster_agent.model import CostFuncType
 from agents.matmaster_agent.state import PLAN
 from agents.matmaster_agent.style import tool_response_failed_card
@@ -84,6 +91,12 @@ def mcp_callback_model_validator(data: Any):
     if data.get('after_model_callback') is None:
         data['after_model_callback'] = default_after_model_callback
 
+    if data.get('before_model_callback') is None:
+        data['before_model_callback'] = default_before_model_callback
+    data['before_model_callback'] = inject_memory_before_model(
+        data['before_model_callback']
+    )
+
     if data.get('before_tool_callback') is None:
         data['before_tool_callback'] = default_before_tool_callback
 
@@ -118,11 +131,13 @@ def mcp_callback_model_validator(data: Any):
 
     data['before_tool_callback'] = catch_before_tool_callback_error(pipeline)
 
-    data['after_tool_callback'] = check_before_tool_callback_effect(
-        catch_after_tool_callback_error(
-            remove_job_link(
-                tgz_oss_to_oss_list(
-                    data['after_tool_callback'], data['enable_tgz_unpack']
+    data['after_tool_callback'] = store_tool_result_in_memory(
+        check_before_tool_callback_effect(
+            catch_after_tool_callback_error(
+                remove_job_link(
+                    tgz_oss_to_oss_list(
+                        data['after_tool_callback'], data['enable_tgz_unpack']
+                    )
                 )
             )
         )

agents/matmaster_agent/flow_agents/agent.py

@@ -48,7 +48,10 @@
 )
 from agents.matmaster_agent.flow_agents.expand_agent.agent import ExpandAgent
 from agents.matmaster_agent.flow_agents.expand_agent.constant import EXPAND_AGENT
-from agents.matmaster_agent.flow_agents.expand_agent.prompt import EXPAND_INSTRUCTION
+from agents.matmaster_agent.flow_agents.expand_agent.prompt import (
+    EXPAND_INSTRUCTION,
+    build_expand_context,
+)
 from agents.matmaster_agent.flow_agents.expand_agent.schema import ExpandSchema
 from agents.matmaster_agent.flow_agents.handle_upload_agent.agent import (
     HandleUploadAgent,
@@ -99,6 +102,11 @@
 from agents.matmaster_agent.llm_config import MatMasterLlmConfig
 from agents.matmaster_agent.locales import i18n
 from agents.matmaster_agent.logger import PrefixFilter
+from agents.matmaster_agent.memory.agent import MemoryWriterAgent
+from agents.matmaster_agent.memory.prompt import (
+    LONG_CONTEXT_THRESHOLD,
+    get_memory_writer_instruction,
+)
 from agents.matmaster_agent.prompt import (
     GLOBAL_INSTRUCTION,
     HUMAN_FRIENDLY_FORMAT_REQUIREMENT,
@@ -110,6 +118,10 @@
     select_update_examples,
     toolchain_from_examples,
 )
+from agents.matmaster_agent.services.memory import (
+    format_short_term_memory,
+    memory_write,
+)
 from agents.matmaster_agent.services.questions import get_random_questions
 from agents.matmaster_agent.services.session_files import get_session_files
 from agents.matmaster_agent.state import (
@@ -193,6 +205,8 @@ def after_init(self):
             before_model_callback=filter_plan_make_llm_contents,
         )
 
+        self._memory_writer_agent = MemoryWriterAgent(MatMasterLlmConfig)
+
         self._thinking_agent = ThinkingAgent(
             name=THINKING_AGENT,
             model=MatMasterLlmConfig.default_litellm_model,
@@ -261,6 +275,11 @@ def scene_agent(self) -> LlmAgent:
     def plan_make_agent(self) -> LlmAgent:
         return self._plan_make_agent
 
+    @computed_field
+    @property
+    def memory_writer_agent(self) -> LlmAgent:
+        return self._memory_writer_agent
+
     @computed_field
     @property
     def execution_agent(self) -> LlmAgent:
@@ -326,20 +345,25 @@ def _build_execution_agent_for_plan(
         return execution_agent
 
     async def _run_expand_agent(
-        self, ctx: InvocationContext
+        self,
+        ctx: InvocationContext,
+        short_term_memory_block: str = '',
+        session_file_summary: str = '',
     ) -> AsyncGenerator[Event, None]:
         # 1. 检索 ICL 示例
+        raw_user_text = ctx.user_content.parts[0].text if ctx.user_content.parts else ''
         icl_examples = select_examples(
-            ctx.user_content.parts[0].text,
+            raw_user_text,
             ctx.session.id,
             CURRENT_ENV,
             logger,
         )
         EXPAND_INPUT_EXAMPLES_PROMPT = expand_input_examples(icl_examples)
         logger.info(f'{ctx.session.id} {EXPAND_INPUT_EXAMPLES_PROMPT}')
         # 2. 动态构造 instruction
+        context = build_expand_context(short_term_memory_block, session_file_summary)
         self.expand_agent.instruction = (
-            EXPAND_INSTRUCTION + EXPAND_INPUT_EXAMPLES_PROMPT
+            context + EXPAND_INSTRUCTION + EXPAND_INPUT_EXAMPLES_PROMPT
         )
         # 3. 运行 Agent
         async for expand_event in self.expand_agent.run_async(ctx):
@@ -407,6 +431,17 @@ async def _run_plan_make_agent(
                 for key, value in available_tools_with_info.items()
             ]
         )
+        query_for_memory = ctx.session.state.get('expand', {}).get(
+            'update_user_content', ''
+        ) or (
+            ctx.user_content.parts[0].text
+            if ctx.user_content and ctx.user_content.parts
+            else ''
+        )
+        short_term_memory_block = await format_short_term_memory(
+            query_text=query_for_memory,
+            session_id=ctx.session.id,
+        )
 
         # Get session files (after full tool list is available)
         try:
@@ -443,6 +478,7 @@ async def _run_plan_make_agent(
                 session_file_summary,
                 original_query,
                 expanded_query,
+                short_term_memory=short_term_memory_block,
             )
             last_full_text = ''
             async for thinking_event in self._thinking_agent.run_async(ctx):
@@ -475,6 +511,7 @@ async def _run_plan_make_agent(
             available_tools_with_info_str
             + UPDATE_USER_CONTENT
             + TOOLCHAIN_EXAMPLES_PROMPT,
+            short_term_memory=short_term_memory_block,
             thinking_context=thinking_text,
             session_file_summary=session_file_summary,
         )
@@ -484,6 +521,46 @@ async def _run_plan_make_agent(
         async for plan_event in self.plan_make_agent.run_async(ctx):
             yield plan_event
 
+        # 记忆写入：用 memory_writer_agent 从当前请求和计划提炼 insights，写入 kernel（不向用户展示）
+        plan_info = ctx.session.state.get(MULTI_PLANS) or {}
+        intro = plan_info.get('intro', '')
+        plans = plan_info.get('plans', [])
+        plan_intro = intro
+        if plans:
+            parts = [intro]
+            for i, plan in enumerate(plans):
+                desc = plan.get('plan_description', '')
+                steps_brief = '; '.join(
+                    s.get('step_description', '')[:60]
+                    for s in plan.get('steps', [])[:5]
+                )
+                parts.append(f"方案{i + 1}摘要: {desc}\n步骤: {steps_brief}")
+            plan_intro = '\n\n'.join(parts)
+        is_long_context = len(UPDATE_USER_CONTENT) >= LONG_CONTEXT_THRESHOLD
+        self.memory_writer_agent.instruction = get_memory_writer_instruction(
+            UPDATE_USER_CONTENT, plan_intro, is_long_context=is_long_context
+        )
+        async for _ in self.memory_writer_agent.run_async(ctx):
+            pass
+        output = ctx.session.state.get('memory_writer_output') or {}
+        insights = output.get('insights', []) if isinstance(output, dict) else []
+        if insights:
+            session_id = ctx.session.id
+            written = 0
+            for text in insights:
+                if isinstance(text, str) and text.strip():
+                    await memory_write(
+                        session_id=session_id, text=text.strip(), metadata={}
+                    )
+                    written += 1
+            logger.info(
+                '%s memory_writer wrote %d insight(s) to memory',
+                ctx.session.id,
+                written,
+            )
+        else:
+            logger.debug('%s memory_writer output 0 insights', ctx.session.id)
+
         # 总结计划
         yield update_state_event(
             ctx,
@@ -659,8 +736,22 @@ async def _run_plan_execute_and_summary_agent(
                 self._analysis_agent.instruction = get_analysis_instruction(
                     ctx.session.state['plan']
                 )
+                analysis_text = ''
                 async for analysis_event in self.analysis_agent.run_async(ctx):
+                    if (cur := is_text(analysis_event)) and not analysis_event.partial:
+                        analysis_text += cur
                     yield analysis_event
+                if analysis_text.strip():
+                    await memory_write(
+                        session_id=ctx.session.id,
+                        text=f"Plan execution summary: {analysis_text.strip()}",
+                        metadata={'source': 'execution_summary'},
+                    )
+                    logger.info(
+                        '%s wrote execution summary to memory (%d chars)',
+                        ctx.session.id,
+                        len(analysis_text),
+                    )
                 self._report_agent.instruction = get_report_instruction(
                     ctx.session.state.get('plan', {})
                 )
@@ -671,6 +762,19 @@ async def _run_plan_execute_and_summary_agent(
                     if (cur_text := is_text(report_event)) and not report_event.partial:
                         report_markdown += cur_text
 
+                if report_markdown.strip():
+                    excerpt = report_markdown.strip()[:5000]
+                    await memory_write(
+                        session_id=ctx.session.id,
+                        text=f"Plan execution report (excerpt): {excerpt}",
+                        metadata={'source': 'execution_report'},
+                    )
+                    logger.info(
+                        '%s wrote report excerpt to memory (%d chars)',
+                        ctx.session.id,
+                        len(excerpt),
+                    )
+
                 # matmaster_report_md.md
                 upload_result = await upload_report_md_to_oss(
                     ReportUploadParams(
@@ -737,8 +841,19 @@ async def _run_plan_execute_and_summary_agent(
     async def _run_research_flow(
         self, ctx: InvocationContext
     ) -> AsyncGenerator[Event, None]:
-        # 扩写用户问题
-        async for _expand_event in self._run_expand_agent(ctx):
+        # 先取短期记忆和会话已有文件，再扩写，避免第二步仍从头 expand（如“第一步的Fe，扩胞到20A”只做扩胞）
+        raw_user_text = ctx.user_content.parts[0].text if ctx.user_content.parts else ''
+        short_term_memory_block = await format_short_term_memory(
+            raw_user_text, ctx.session.id
+        )
+        session_files = await get_session_files(ctx.session.id)
+        session_file_summary = '\n'.join(session_files) if session_files else ''
+        # 扩写用户问题（带记忆 + 会话文件，延续上一步时只 expand 新步骤）
+        async for _expand_event in self._run_expand_agent(
+            ctx,
+            short_term_memory_block=short_term_memory_block,
+            session_file_summary=session_file_summary,
+        ):
             yield _expand_event
 
         # 构造 UPDATE_USER_CONTENT, SCENE_EXAMPLES_PROMPT, TOOLCHAIN_EXAMPLES_PROMPT

agents/matmaster_agent/flow_agents/expand_agent/prompt.py

@@ -1,3 +1,33 @@
+# --- Context section headers (injected before EXPAND_INSTRUCTION when present) ---
+SECTION_SHORT_TERM_MEMORY = 'SHORT-TERM WORKING MEMORY'
+SECTION_SESSION_FILES = 'SESSION FILES'
+
+MEMORY_SECTION_HEADER = (
+    f'# {SECTION_SHORT_TERM_MEMORY}\n'
+    'Use the following when expanding the user request.\n\n'
+)
+SESSION_FILES_SECTION_HEADER = (
+    f'# {SECTION_SESSION_FILES}\n'
+    'Files already produced in this session. If the user refers to a previous step (e.g. "第一步", "上一步", "刚才") '
+    'and these files exist, expand only the new step; do not re-add structure-building steps.\n\n'
+)
+
+
+def build_expand_context(
+    short_term_memory_block: str = '',
+    session_file_summary: str = '',
+) -> str:
+    """Build the optional context block (memory + session files) to prepend to expand instruction."""
+    parts = []
+    if short_term_memory_block:
+        parts.append(MEMORY_SECTION_HEADER + short_term_memory_block.strip() + '\n\n')
+    if session_file_summary:
+        parts.append(
+            SESSION_FILES_SECTION_HEADER + session_file_summary.strip() + '\n\n'
+        )
+    return ''.join(parts)
+
+
 EXPAND_INSTRUCTION = """
 You are a computational materials science assistant specializing in structure generation. Follow this structured protocol for all user requests:
 

agents/matmaster_agent/flow_agents/plan_make_agent/prompt.py

@@ -103,16 +103,25 @@ def get_static_plan_system_block(available_tools_with_info: str) -> str:
 def get_dynamic_plan_user_block(
     thinking_context: str = '',
     session_file_summary: str = '',
+    short_term_memory: str = '',
 ) -> str:
     """
-    Mutable content: <Prior Thinking> and <Session File Info>. Changes every turn.
+    Mutable content: <Prior Thinking>, <Session File Info>, and optional <Session Memory>.
     """
     parts = []
     if session_file_summary:
         parts.append(
             f"""
 <Session File Info>
 {session_file_summary}
+"""
+        )
+    if short_term_memory:
+        parts.append(
+            f"""
+<Session Memory>
+{short_term_memory.strip()}
+</Session Memory>
 """
         )
     if thinking_context:
@@ -134,12 +143,15 @@ def get_plan_make_instruction(
     available_tools_with_info: str,
     thinking_context: str = '',
     session_file_summary: str = '',
+    short_term_memory: str = '',
 ) -> str:
     """
-    Returns a single prompt: static content (tools + rules) then dynamic (session + thinking).
+    Returns a single prompt: static content (tools + rules) then dynamic (session + memory + thinking).
     """
     static = get_static_plan_system_block(available_tools_with_info)
-    dynamic = get_dynamic_plan_user_block(thinking_context, session_file_summary)
+    dynamic = get_dynamic_plan_user_block(
+        thinking_context, session_file_summary, short_term_memory
+    )
     if not dynamic:
         return static
     return static + '\n\n' + dynamic