feat(memory_collection): add MemoryConversation class for dialogue history management

Change-Id: I099853b5465234a34db594c9dea1210743174932
Co-developed-by: Cursor <noreply@cursor.com>
Signed-off-by: 久氢 <mapenghui.mph@alibaba-inc.com>

Author: 久氢

agentrun/memory_collection/memory_conversation.py

@@ -2,25 +2,19 @@
 
 提供与 TableStore Memory 的集成能力，自动存储用户和 Agent 的对话历史。
 
-Example (基本使用):
-    >>> from agentrun.server import AgentRunServer, AgentRequest
-    >>> from agentrun.memory_collection import MemoryConversation
-    >>>
-    >>> # 初始化 Memory Conversation
-    >>> memory = MemoryConversation(memory_collection_name="my-memory")
-    >>>
-    >>> # 包装 invoke_agent 函数
-    >>> async def invoke_agent(req: AgentRequest):
-    ...     async for event in memory.wrap_invoke_agent(req, my_agent_handler):
-    ...         yield event
-    >>>
-    >>> server = AgentRunServer(invoke_agent=invoke_agent)
-    >>> server.start()
 """
 
 import json
 import os
-from typing import Any, AsyncIterator, Callable, Dict, Optional, TYPE_CHECKING
+from typing import (
+    Any,
+    AsyncIterator,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    TYPE_CHECKING,
+)
 import uuid
 
 import tablestore
@@ -88,13 +82,17 @@ def _default_user_id_extractor(req: Any) -> str:
         """默认的 user_id 提取器
 
         优先级：
-        1. X-User-ID 请求头
+        1. X-User-ID 请求头（支持多种格式）
         2. user_id 查询参数
         3. 默认值 "default_user"
         """
         if req.raw_request:
             # 从请求头获取
-            user_id = req.raw_request.headers.get("X-User-ID")
+            user_id = (
+                req.raw_request.headers.get("X-AgentRun-User-ID")
+                or req.raw_request.headers.get("x-agentrun-user-id")
+                or req.raw_request.headers.get("X-Agentrun-User-Id")
+            )
             if user_id:
                 return user_id
 
@@ -110,26 +108,20 @@ def _default_session_id_extractor(req: Any) -> str:
         """默认的 session_id 提取器
 
         优先级：
-        1. X-Session-ID 请求头
-        2. sessionId 查询参数
-        3. 从最后一条消息的 id 生成
-        4. 生成新的 UUID
+        1. X-Session-ID 请求头（支持多种格式）
+        2. 生成新的 UUID
         """
         if req.raw_request:
-            # 从请求头获取
-            session_id = req.raw_request.headers.get("X-Conversation-ID")
-            if session_id:
-                return session_id
-
-            # 从查询参数获取
-            session_id = req.raw_request.query_params.get("sessionId")
+            # 从请求头获取（兼容多种格式）
+            # 支持：X-AgentRun-Session-ID, x-agentrun-session-id, X-Agentrun-Session-Id
+            session_id = (
+                req.raw_request.headers.get("X-AgentRun-Session-ID")
+                or req.raw_request.headers.get("x-agentrun-session-id")
+                or req.raw_request.headers.get("X-Agentrun-Session-Id")
+            )
             if session_id:
                 return session_id
 
-        # 从消息 ID 生成（如果有）
-        if req.messages and req.messages[-1].id:
-            return f"session_{req.messages[-1].id}"
-
         # 生成新的 session_id
         return f"session_{uuid.uuid4().hex[:16]}"
 
@@ -138,13 +130,18 @@ def _default_agent_id_extractor(req: Any) -> str:
         """默认的 agent_id 提取器
 
         优先级：
-        1. X-Agent-ID 请求头
+        1. X-Agent-ID 请求头（支持多种格式）
         2. 从 URL 路径中提取 /agent-runtimes/{agent_id}/... 格式
         3. 默认值 "default_agent"
         """
         if req.raw_request:
-            # 从请求头获取
-            agent_id = req.raw_request.headers.get("X-Agent-ID")
+            # 从请求头获取（兼容多种格式）
+            # 支持：X-AgentRun-Agent-ID, x-agentrun-agent-id, X-Agentrun-Agent-Id, X-Agent-ID, x-agent-id
+            agent_id = (
+                req.raw_request.headers.get("X-AgentRun-Agent-ID")
+                or req.raw_request.headers.get("x-agentrun-agent-id")
+                or req.raw_request.headers.get("X-Agentrun-Agent-Id")
+            )
             if agent_id:
                 return agent_id
 
@@ -407,8 +404,12 @@ async def wrap_invoke_agent(
                 "content": self._extract_message_content(msg.content),
             })
 
-        # 收集 Agent 响应
+        # 收集 Agent 响应（包括文本和工具调用）
         agent_response_content = ""
+        tool_calls: Dict[str, Dict[str, Any]] = (
+            {}
+        )  # tool_call_id -> tool_call_info
+        tool_results: List[Dict[str, Any]] = []  # 工具执行结果列表
 
         try:
             # 流式处理 Agent 响应
@@ -420,17 +421,80 @@ async def wrap_invoke_agent(
                     if event.event == EventType.TEXT and "delta" in event.data:
                         agent_response_content += event.data["delta"]
 
+                    # 收集工具调用信息
+                    elif event.event == EventType.TOOL_CALL:
+                        # 完整的工具调用
+                        tool_id = event.data.get("id", "")
+                        if tool_id:
+                            tool_calls[tool_id] = {
+                                "id": tool_id,
+                                "type": "function",
+                                "function": {
+                                    "name": event.data.get("name", ""),
+                                    "arguments": event.data.get("args", ""),
+                                },
+                            }
+
+                    elif event.event == EventType.TOOL_CALL_CHUNK:
+                        # 工具调用片段（流式场景）
+                        tool_id = event.data.get("id", "")
+                        if tool_id:
+                            if tool_id not in tool_calls:
+                                tool_calls[tool_id] = {
+                                    "id": tool_id,
+                                    "type": "function",
+                                    "function": {
+                                        "name": event.data.get("name", ""),
+                                        "arguments": "",
+                                    },
+                                }
+                            # 累积参数片段
+                            if "args_delta" in event.data:
+                                tool_calls[tool_id]["function"][
+                                    "arguments"
+                                ] += event.data["args_delta"]
+
+                    # 收集工具执行结果
+                    elif event.event == EventType.TOOL_RESULT:
+                        tool_id = event.data.get("id", "")
+                        if tool_id:
+                            tool_results.append({
+                                "role": "tool",
+                                "tool_call_id": tool_id,
+                                "content": str(event.data.get("result", "")),
+                            })
+
                 # 透传事件
                 yield event
 
             # 保存完整的对话轮次（输入 + 输出）
-            if agent_response_content:
+            # 只有当有文本内容或工具调用时才保存
+            if agent_response_content or tool_calls or tool_results:
                 try:
-                    # 将助手响应添加到消息列表
-                    output_messages = input_messages + [{
+                    # 构建助手响应消息
+                    assistant_message: Dict[str, Any] = {
                         "role": "assistant",
-                        "content": agent_response_content,
-                    }]
+                    }
+
+                    # 添加文本内容（如果有）
+                    if agent_response_content:
+                        assistant_message["content"] = agent_response_content
+                    else:
+                        # OpenAI 格式要求：如果有 tool_calls，content 可以为 null
+                        assistant_message["content"] = None
+
+                    # 添加工具调用（如果有）
+                    if tool_calls:
+                        assistant_message["tool_calls"] = list(
+                            tool_calls.values()
+                        )
+
+                    # 构建完整的消息列表
+                    output_messages = input_messages + [assistant_message]
+
+                    # 添加工具执行结果（如果有）
+                    if tool_results:
+                        output_messages.extend(tool_results)
 
                     # 将完整的对话历史存储为一条消息
                     # content 字段存储 JSON 格式的消息列表
@@ -446,8 +510,10 @@ async def wrap_invoke_agent(
                     await memory_store.update_session(session)
 
                     logger.debug(
-                        f"Saved conversation: {len(output_messages)} messages, "
-                        f"response length: {len(agent_response_content)} chars"
+                        f"Saved conversation: {len(output_messages)} messages,"
+                        f" text length: {len(agent_response_content)} chars,"
+                        f" tool_calls: {len(tool_calls)}, tool_results:"
+                        f" {len(tool_results)}"
                     )
                 except Exception as e:
                     logger.error(

agentrun/server/server.py

@@ -83,13 +83,22 @@ class AgentRunServer:
         ...     invoke_agent=invoke_agent,
         ...     config=ServerConfig(cors_origins=["http://localhost:3000"])
         ... )
+
+    Example (启用会话历史记录):
+        >>> server = AgentRunServer(
+        ...     invoke_agent=invoke_agent,
+        ...     memory_collection_name="my-memory-collection"
+        ... )
+        >>> server.start(port=8000)
+        # 会话历史将自动保存到 TableStore
     """
 
     def __init__(
         self,
         invoke_agent: InvokeAgentHandler,
         protocols: Optional[List[ProtocolHandler]] = None,
         config: Optional[ServerConfig] = None,
+        memory_collection_name: Optional[str] = None,
     ):
         """初始化 AgentRun Server
 
@@ -107,8 +116,20 @@ def __init__(
                 - cors_origins: CORS 允许的源列表
                 - openai: OpenAI 协议配置
                 - agui: AG-UI 协议配置
+
+            memory_collection_name: MemoryCollection 名称（可选）
+                - 如果提供，将自动启用会话历史记录功能
+                - 会话历史将保存到指定的 MemoryCollection 中
         """
         self.app = FastAPI(title="AgentRun Server")
+
+        # 如果启用了 memory，包装 invoke_agent
+        if memory_collection_name:
+            invoke_agent = self._wrap_with_memory(
+                invoke_agent,
+                memory_collection_name,
+            )
+
         self.agent_invoker = AgentInvoker(invoke_agent)
 
         # 配置 CORS
@@ -124,6 +145,39 @@ def __init__(
         # 挂载所有协议的 Router
         self._mount_protocols(protocols)
 
+    def _wrap_with_memory(
+        self,
+        invoke_agent: InvokeAgentHandler,
+        memory_collection_name: str,
+    ) -> InvokeAgentHandler:
+        """使用 MemoryConversation 包装 invoke_agent
+
+        Args:
+            invoke_agent: 原始的 invoke_agent 函数
+            memory_collection_name: MemoryCollection 名称
+
+        Returns:
+            包装后的 invoke_agent 函数
+        """
+        from agentrun.memory_collection import MemoryConversation
+
+        # 创建 MemoryConversation 实例
+        memory = MemoryConversation(
+            memory_collection_name=memory_collection_name,
+        )
+
+        logger.info(
+            "Memory integration enabled for collection:"
+            f" {memory_collection_name}"
+        )
+
+        # 包装 invoke_agent
+        async def wrapped_invoke_agent(request: Any):
+            async for event in memory.wrap_invoke_agent(request, invoke_agent):
+                yield event
+
+        return wrapped_invoke_agent
+
     def _setup_cors(self, cors_origins: Optional[Sequence[str]] = None):
         """配置 CORS 中间件
 

examples/server_with_memory.py

@@ -32,56 +32,60 @@
 # 配置参数
 MODEL_SERVICE = os.getenv("MODEL_SERVICE", "qwen3-max")
 MODEL_NAME = os.getenv("MODEL_NAME", "qwen3-max")
-SANDBOX_NAME = os.getenv("SANDBOX_NAME", "")
+SANDBOX_NAME = os.getenv("SANDBOX_NAME", "sandbox-browser-BmUyyD")
 MEMORY_COLLECTION_NAME = os.getenv("MEMORY_COLLECTION_NAME", "mem-ots0129")
 
 # 创建 Agent
 agent = create_agent(
     # 使用 AgentRun 注册的模型
     model=model(MODEL_SERVICE, model=MODEL_NAME),
     system_prompt="""
-你是一个诗人，根据用户输入内容写一个20字以内的诗文
+你是 AgentRun 的 AI 助手，可以通过网络搜索帮助用户解决问题
+
+
+你的工作流程如下
+- 当用户向你提问概念性问题时，不要直接回答，而是先进行网络搜索
+- 使用 Browser 工具打开百度搜索。如果要搜索 AgentRun，对应的搜索链接为: `https://www.baidu.com/s?ie=utf-8&wd=agentrun`。为了节省 token 使用，不要使用 `snapshot` 获取完整页面内容，而是通过 `evaluate` 获取你需要的部分
+- 获取百度搜索的结果，根据相关性分别打开子页面获取内容
+    - 如果子页面的相关度较低，则可以直接忽略
+    - 如果子页面的相关度较高，则将其记录为可参考的资料，记录页面标题和实时的 url
+- 当你获得至少 3 条网络信息后，可以结束搜索，并根据搜索到的结果回答用户的问题。
+- 如果某一部分回答引用了网络的信息，需要进行标注，并在回答的最后给出跳转链接
 """,
     # 使用 AgentRun 的 Sandbox 工具
-    # tools=[*sandbox_toolset(SANDBOX_NAME, template_type=TemplateType.BROWSER)],
+    tools=[*sandbox_toolset(SANDBOX_NAME, template_type=TemplateType.BROWSER)],
 )
 
 # 初始化 Memory Integration
 memory = MemoryConversation(memory_collection_name=MEMORY_COLLECTION_NAME)
 
 
 async def invoke_agent(req: AgentRequest):
-    """Agent 调用函数，集成了记忆存储功能"""
+    """Agent 调用函数"""
     try:
         converter = AgentRunConverter()
-
-        # 定义原始的 agent 处理函数
-        async def agent_handler(request: AgentRequest):
-            result = agent.astream_events(
-                {
-                    "messages": [
-                        {"role": msg.role, "content": msg.content}
-                        for msg in request.messages
-                    ]
-                },
-                config={"recursion_limit": 1000},
-            )
-            async for event in result:
-                for agentrun_event in converter.convert(event):
-                    yield agentrun_event
-
-        # 使用 MemoryIntegration 包装，自动存储对话历史
-        async for event in memory.wrap_invoke_agent(req, agent_handler):
-            yield event
-
+        result = agent.astream_events(
+            {
+                "messages": [
+                    {"role": msg.role, "content": msg.content}
+                    for msg in req.messages
+                ]
+            },
+            config={"recursion_limit": 1000},
+        )
+        async for event in result:
+            for agentrun_event in converter.convert(event):
+                yield agentrun_event
     except Exception as e:
-        print(f"Error in invoke_agent: {e}")
+        print(e)
         raise Exception("Internal Error")
 
 
 # 创建并启动 Server
 if __name__ == "__main__":
-    server = AgentRunServer(invoke_agent=invoke_agent)
+    server = AgentRunServer(
+        invoke_agent=invoke_agent, memory_collection_name=MEMORY_COLLECTION_NAME
+    )
     print(f"Server starting with memory collection: {MEMORY_COLLECTION_NAME}")
     print("Memory will be automatically saved to TableStore")
     server.start(port=9000)