langchain-agentx-python 0.1__py3-none-any.whl

langchain_agentx/memory/memdir/agent_memory.py

@@ -0,0 +1,380 @@
+"""
+memory/memdir/agent_memory.py — Layer 4 agent 记忆路径协作者。
+
+职责：
+    负责 agent_type 安全化与 Layer 4 路径边界判断，供加载/权限/抽取链路复用。
+
+链路位置：
+    subagent 启动加载与 Layer 4 写权限判定都会调用本模块。
+
+当前裁剪范围：
+    Phase 2 补充 prompt 注入能力（subagent 专用），抽取编排在后续 Phase。
+"""
+
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Awaitable, Callable
+
+from langchain_agentx.workspace import AgentWorkspaceConfig
+from langchain_agentx.workspace.config import _sanitize_agent_type
+
+from .loader import MemoryPromptLoader
+from .paths import is_memory_enabled
+
+if TYPE_CHECKING:
+    from langchain_agentx.loop.prompt.sections import SystemPromptSection
+
+
+_SCOPE_GUIDANCE: dict[str, str] = {
+    "user": "- Since this memory is user-scope, keep learnings general since they apply across all projects",
+    "project": "- Since this memory is project-scope and shared with your team via version control, tailor your memories to this project",
+    "local": "- Since this memory is local-scope (not checked into version control), tailor your memories to this project and machine",
+}
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class _AgentMemoryWorkspaceView:
+    memory_dir: Path
+    memory_entrypoint: Path
+
+
+@dataclass(frozen=True)
+class AgentMemoryExtractionContext:
+    session_id: str
+    container_type: str
+    workspace_cfg: AgentWorkspaceConfig
+    agent_type: str
+    messages: list[Any]
+    agent_id: str | None = None
+
+
+class AgentMemoryPathPolicy:
+    """Layer 4 路径策略协作者。"""
+
+    def __init__(self, cfg: AgentWorkspaceConfig) -> None:
+        self._cfg = cfg
+        self._agents_root = (cfg.memory_dir / "agents").resolve()
+
+    def is_agent_memory_path(self, absolute_path: str | Path) -> bool:
+        resolved = Path(absolute_path).resolve()
+        try:
+            resolved.relative_to(self._agents_root)
+            return True
+        except ValueError:
+            return False
+
+    def is_agent_memory_path_for_type(
+        self,
+        *,
+        agent_type: str,
+        absolute_path: str | Path,
+        allowed_scopes: tuple[str, ...] = ("project", "user", "local"),
+    ) -> bool:
+        resolved = Path(absolute_path).resolve()
+        sanitized = _sanitize_agent_type(agent_type)
+        for scope in allowed_scopes:
+            mem_dir = self._cfg.agent_memory_dir(sanitized, scope).resolve()
+            try:
+                resolved.relative_to(mem_dir)
+                return True
+            except ValueError:
+                continue
+        return False
+
+
+def sanitize_agent_type(agent_type: str) -> str:
+    return _sanitize_agent_type(agent_type)
+
+
+class AgentMemoryPromptService:
+    """Layer 4 agent memory prompt 构建协作者。"""
+
+    def __init__(self, cfg: AgentWorkspaceConfig) -> None:
+        self._cfg = cfg
+
+    def load_agent_memory_prompt(self, *, agent_type: str, scope: str = "project") -> str | None:
+        sanitized = _sanitize_agent_type(agent_type)
+        memory_dir = self._cfg.agent_memory_dir(sanitized, scope)
+        entrypoint = self._cfg.agent_memory_entrypoint(sanitized, scope)
+        prompt = MemoryPromptLoader(
+            _AgentMemoryWorkspaceView(memory_dir=memory_dir, memory_entrypoint=entrypoint)  # type: ignore[arg-type]
+        ).load_memory_prompt()
+        if prompt is None:
+            return None
+        guidance = _SCOPE_GUIDANCE.get(scope)
+        if guidance:
+            return prompt + "\n\n" + guidance
+        return prompt
+
+
+class AgentMemoryPromptBootstrap:
+    """Loop 侧 Layer 4 section 注入协作者。"""
+
+    def __init__(
+        self,
+        *,
+        workspace_cfg: AgentWorkspaceConfig,
+        is_subagent: bool,
+        subagent_type: str | None,
+        scope: str = "project",
+    ) -> None:
+        self._workspace_cfg = workspace_cfg
+        self._is_subagent = is_subagent
+        self._subagent_type = subagent_type
+        self._scope = scope
+
+    def build_sections(self) -> list["SystemPromptSection"]:
+        if not self._is_subagent or not self._subagent_type:
+            return []
+        from langchain_agentx.loop.prompt.sections import section
+
+        service = AgentMemoryPromptService(self._workspace_cfg)
+        return [
+            section(
+                "agent_memory",
+                lambda: service.load_agent_memory_prompt(
+                    agent_type=self._subagent_type or "",
+                    scope=self._scope,
+                ),
+            )
+        ]
+
+
+class AgentMemoryExtractionCoordinator:
+    """Layer 4 agent memory 后台抽取协作者（独立于 Layer 2）。"""
+
+    def __init__(
+        self,
+        *,
+        extraction_executor: Callable[[AgentMemoryExtractionContext, list[Any]], Awaitable[None]]
+        | None = None,
+    ) -> None:
+        self._executor = extraction_executor or _noop_agent_memory_extraction_executor
+        self._in_progress = False
+        self._pending_context: AgentMemoryExtractionContext | None = None
+        self._last_message_id: dict[str, str] = {}
+        self._tasks: set[asyncio.Task[None]] = set()
+
+    def trigger_from_state(
+        self,
+        *,
+        state: dict[str, Any],
+        workspace_cfg: AgentWorkspaceConfig,
+        session_id: str,
+        container_type: str,
+        agent_type: str | None,
+    ) -> None:
+        if not agent_type:
+            logger.warning(
+                "skip agent memory extraction trigger: empty agent_type (container=%s)",
+                container_type,
+            )
+            return
+        context = AgentMemoryExtractionContext(
+            session_id=session_id,
+            container_type=container_type,
+            workspace_cfg=workspace_cfg,
+            agent_type=agent_type,
+            messages=list(state.get("messages") or []),
+            agent_id=state.get("agent_id"),
+        )
+        self.trigger(context)
+
+    def trigger(self, context: AgentMemoryExtractionContext) -> None:
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            logger.warning("skip agent memory extraction trigger: no running event loop")
+            return
+        task = loop.create_task(self.execute_extract_memories(context))
+        self._tasks.add(task)
+        task.add_done_callback(self._tasks.discard)
+
+    async def execute_extract_memories(self, context: AgentMemoryExtractionContext) -> None:
+        if not self._can_extract(context):
+            return
+        if self._in_progress:
+            self._pending_context = context
+            return
+        self._in_progress = True
+        try:
+            current = context
+            while True:
+                await self._execute_once(current)
+                if self._pending_context is None:
+                    break
+                current = self._pending_context
+                self._pending_context = None
+        finally:
+            self._in_progress = False
+
+    async def drain_pending_extraction(self, timeout_ms: int = 60_000) -> None:
+        if not self._tasks:
+            return
+        try:
+            await asyncio.wait_for(
+                asyncio.gather(*list(self._tasks), return_exceptions=True),
+                timeout=timeout_ms / 1000.0,
+            )
+        except asyncio.TimeoutError:
+            logger.warning("agent memory extraction drain timed out after %sms", timeout_ms)
+
+    async def _execute_once(self, context: AgentMemoryExtractionContext) -> None:
+        new_messages = self._messages_since_cursor(context)
+        if not new_messages:
+            return
+        if self._has_agent_memory_writes_since(context, new_messages):
+            self._advance_cursor(context, new_messages)
+            return
+        await self._executor(context, new_messages)
+        self._advance_cursor(context, new_messages)
+
+    @staticmethod
+    def _can_extract(context: AgentMemoryExtractionContext) -> bool:
+        return (
+            context.container_type == "subagent"
+            and bool(context.agent_type)
+            and is_memory_enabled(context.workspace_cfg)
+        )
+
+    def _messages_since_cursor(self, context: AgentMemoryExtractionContext) -> list[Any]:
+        if not context.messages:
+            return []
+        cursor = self._last_message_id.get(self._session_key(context))
+        if cursor is None:
+            return context.messages
+        for idx, message in enumerate(context.messages):
+            if self._message_id(message) == cursor:
+                return context.messages[idx + 1 :]
+        return context.messages
+
+    def _advance_cursor(self, context: AgentMemoryExtractionContext, messages: list[Any]) -> None:
+        last_id = self._message_id(messages[-1])
+        if last_id:
+            self._last_message_id[self._session_key(context)] = last_id
+
+    def _has_agent_memory_writes_since(
+        self, context: AgentMemoryExtractionContext, messages: list[Any]
+    ) -> bool:
+        for message in messages:
+            calls = []
+            if isinstance(message, dict):
+                calls = list(message.get("tool_calls") or [])
+            else:
+                calls = list(getattr(message, "tool_calls", None) or [])
+            for call in calls:
+                if _is_agent_memory_write_call(call, context.workspace_cfg, context.agent_type):
+                    return True
+        return False
+
+    @staticmethod
+    def _session_key(context: AgentMemoryExtractionContext) -> str:
+        return f"{context.container_type}:{context.session_id}:{context.agent_type}"
+
+    @staticmethod
+    def _message_id(message: Any) -> str | None:
+        if isinstance(message, dict):
+            value = message.get("id")
+            return str(value) if value else None
+        value = getattr(message, "id", None)
+        return str(value) if value else None
+
+
+async def _noop_agent_memory_extraction_executor(
+    context: AgentMemoryExtractionContext, messages: list[Any]
+) -> None:
+    logger.info(
+        "agent memory extraction placeholder executor triggered session=%s agent_type=%s message_count=%s",
+        context.session_id,
+        context.agent_type,
+        len(messages),
+    )
+    return None
+
+
+def _is_agent_memory_write_call(
+    call: Any,
+    workspace_cfg: AgentWorkspaceConfig,
+    agent_type: str,
+) -> bool:
+    if isinstance(call, dict):
+        name = str(call.get("name") or call.get("tool_name") or "")
+        args = call.get("args") or call.get("tool_input") or {}
+    else:
+        name = str(getattr(call, "name", "") or getattr(call, "tool_name", ""))
+        args = getattr(call, "args", None) or getattr(call, "tool_input", None) or {}
+    if name not in {"Write", "Edit"}:
+        return False
+    if not isinstance(args, dict):
+        return False
+    candidate = args.get("path") or args.get("file_path")
+    if not candidate:
+        return False
+    return is_agent_memory_path_for_type(
+        workspace_cfg,
+        agent_type,
+        candidate,
+        allowed_scopes=("project",),
+    )
+
+
+_DEFAULT_AGENT_MEMORY_COORDINATOR = AgentMemoryExtractionCoordinator()
+
+
+def get_default_agent_memory_extraction_coordinator() -> AgentMemoryExtractionCoordinator:
+    return _DEFAULT_AGENT_MEMORY_COORDINATOR
+
+
+async def drain_pending_agent_memory_extraction(timeout_ms: int = 60_000) -> None:
+    await _DEFAULT_AGENT_MEMORY_COORDINATOR.drain_pending_extraction(timeout_ms=timeout_ms)
+
+
+def is_agent_memory_path(cfg: AgentWorkspaceConfig, absolute_path: str | Path) -> bool:
+    return AgentMemoryPathPolicy(cfg).is_agent_memory_path(absolute_path)
+
+
+def is_agent_memory_path_for_type(
+    cfg: AgentWorkspaceConfig,
+    agent_type: str,
+    absolute_path: str | Path,
+    *,
+    allowed_scopes: tuple[str, ...] = ("project", "user", "local"),
+) -> bool:
+    return AgentMemoryPathPolicy(cfg).is_agent_memory_path_for_type(
+        agent_type=agent_type,
+        absolute_path=absolute_path,
+        allowed_scopes=allowed_scopes,
+    )
+
+
+def load_agent_memory_prompt(
+    cfg: AgentWorkspaceConfig,
+    *,
+    agent_type: str,
+    scope: str = "project",
+) -> str | None:
+    return AgentMemoryPromptService(cfg).load_agent_memory_prompt(
+        agent_type=agent_type,
+        scope=scope,
+    )
+
+
+__all__ = [
+    "AgentMemoryExtractionContext",
+    "AgentMemoryExtractionCoordinator",
+    "AgentMemoryPromptBootstrap",
+    "AgentMemoryPromptService",
+    "AgentMemoryPathPolicy",
+    "drain_pending_agent_memory_extraction",
+    "get_default_agent_memory_extraction_coordinator",
+    "is_agent_memory_path",
+    "is_agent_memory_path_for_type",
+    "load_agent_memory_prompt",
+    "sanitize_agent_type",
+]

langchain_agentx/memory/memdir/extractor.py

@@ -0,0 +1,309 @@
+"""
+memory/memdir/extractor.py — Layer 2 后台抽取协作者。
+
+职责：
+    在 loop stop 时以 fire-and-forget 触发记忆抽取，提供并发保护与 session 退出 drain。
+
+链路位置：
+    HookGraphWiring.stop_hooks_node() 触发 -> MemoryExtractionCoordinator 执行。
+
+当前裁剪范围：
+    Phase 4 先实现门禁链、游标推进、并发与 drain；抽取写入细节使用可注入执行器。
+"""
+
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass
+import logging
+from pathlib import Path
+from typing import Any, Awaitable, Callable
+
+from langchain_agentx.workspace import AgentWorkspaceConfig
+
+from .agent_memory import is_agent_memory_path
+from .paths import is_memory_enabled
+from .scan import format_memory_manifest, scan_memory_files
+from .types import HOW_TO_SAVE_SECTION, TYPES_SECTION_INDIVIDUAL, WHAT_NOT_TO_SAVE_SECTION
+
+logger = logging.getLogger(__name__)
+
+_WRITE_TOOLS = {"Write", "Edit"}
+_READ_TOOLS = {"Read", "Grep", "Glob"}
+_READ_ONLY_BASH_PREFIX = ("ls", "find", "grep", "cat", "stat", "wc", "head", "tail")
+DEFAULT_EXTRACTION_WINDOW_MESSAGES = 30
+
+
+@dataclass(frozen=True)
+class MemoryExtractionContext:
+    session_id: str
+    container_type: str
+    workspace_cfg: AgentWorkspaceConfig
+    messages: list[Any]
+    agent_id: str | None = None
+
+
+class MemoryExtractionCoordinator:
+    def __init__(
+        self,
+        *,
+        extraction_executor: Callable[[MemoryExtractionContext, list[Any]], Awaitable[None]] | None = None,
+    ) -> None:
+        self._executor = extraction_executor or _noop_extraction_executor
+        self._in_progress = False
+        self._pending_context: MemoryExtractionContext | None = None
+        self._last_message_id: dict[str, str] = {}
+        self._tasks: set[asyncio.Task[None]] = set()
+
+    def trigger_from_state(
+        self,
+        *,
+        state: dict[str, Any],
+        workspace_cfg: AgentWorkspaceConfig,
+        session_id: str,
+        container_type: str,
+    ) -> None:
+        context = MemoryExtractionContext(
+            session_id=session_id,
+            container_type=container_type,
+            workspace_cfg=workspace_cfg,
+            messages=list(state.get("messages") or []),
+            agent_id=state.get("agent_id"),
+        )
+        self.trigger(context)
+
+    def trigger(self, context: MemoryExtractionContext) -> None:
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            logger.warning("skip memory extraction trigger: no running event loop")
+            return
+        task = loop.create_task(self.execute_extract_memories(context))
+        self._tasks.add(task)
+        task.add_done_callback(self._tasks.discard)
+
+    async def execute_extract_memories(self, context: MemoryExtractionContext) -> None:
+        if not self._can_extract(context):
+            return
+        if self._in_progress:
+            # 单槽覆盖语义：只保留最新上下文。
+            self._pending_context = context
+            return
+        self._in_progress = True
+        try:
+            current = context
+            while True:
+                await self._execute_once(current)
+                if self._pending_context is None:
+                    break
+                current = self._pending_context
+                self._pending_context = None
+        finally:
+            self._in_progress = False
+
+    async def drain_pending_extraction(self, timeout_ms: int = 60_000) -> None:
+        if not self._tasks:
+            return
+        try:
+            await asyncio.wait_for(
+                asyncio.gather(*list(self._tasks), return_exceptions=True),
+                timeout=timeout_ms / 1000.0,
+            )
+        except asyncio.TimeoutError:
+            logger.warning("memory extraction drain timed out after %sms", timeout_ms)
+
+    async def _execute_once(self, context: MemoryExtractionContext) -> None:
+        new_messages = self._messages_since_cursor(context)
+        if not new_messages:
+            return
+        if self._has_memory_writes_since(context, new_messages):
+            self._advance_cursor(context, new_messages)
+            return
+        await self._executor(context, new_messages)
+        self._advance_cursor(context, new_messages)
+
+    def _can_extract(self, context: MemoryExtractionContext) -> bool:
+        if context.container_type not in ("interactive", "conversation"):
+            return False
+        if context.agent_id:
+            return False
+        return is_memory_enabled(context.workspace_cfg)
+
+    def _messages_since_cursor(self, context: MemoryExtractionContext) -> list[Any]:
+        if not context.messages:
+            return []
+        session_key = self._session_key(context)
+        cursor = self._last_message_id.get(session_key) or self._load_cursor(context)
+        if cursor is None:
+            return context.messages
+        for idx, message in enumerate(context.messages):
+            if self._message_id(message) == cursor:
+                return context.messages[idx + 1 :]
+        return context.messages
+
+    def _advance_cursor(self, context: MemoryExtractionContext, messages: list[Any]) -> None:
+        last_id = self._message_id(messages[-1])
+        if not last_id:
+            return
+        session_key = self._session_key(context)
+        self._last_message_id[session_key] = last_id
+        if context.container_type == "conversation":
+            cursor_path = context.workspace_cfg.sessions_dir / f"{context.session_id}.cursor"
+            cursor_path.parent.mkdir(parents=True, exist_ok=True)
+            cursor_path.write_text(last_id, encoding="utf-8")
+
+    def _load_cursor(self, context: MemoryExtractionContext) -> str | None:
+        if context.container_type != "conversation":
+            return None
+        path = context.workspace_cfg.sessions_dir / f"{context.session_id}.cursor"
+        if not path.exists():
+            return None
+        return path.read_text(encoding="utf-8").strip() or None
+
+    @staticmethod
+    def _session_key(context: MemoryExtractionContext) -> str:
+        return f"{context.container_type}:{context.session_id}"
+
+    @staticmethod
+    def _message_id(message: Any) -> str | None:
+        if isinstance(message, dict):
+            value = message.get("id")
+            return str(value) if value else None
+        value = getattr(message, "id", None)
+        return str(value) if value else None
+
+    @staticmethod
+    def _has_memory_writes_since(context: MemoryExtractionContext, messages: list[Any]) -> bool:
+        for message in messages:
+            calls = []
+            if isinstance(message, dict):
+                calls = list(message.get("tool_calls") or [])
+            else:
+                calls = list(getattr(message, "tool_calls", None) or [])
+            for call in calls:
+                if _is_memory_write_call(call, context.workspace_cfg):
+                    return True
+        return False
+
+
+def create_auto_mem_can_use_tool(memory_dir: Path) -> Callable[[str, dict[str, Any]], bool]:
+    resolved_memory_dir = memory_dir.resolve()
+
+    def _can_use(tool_name: str, tool_input: dict[str, Any]) -> bool:
+        if tool_name in _READ_TOOLS:
+            return True
+        if tool_name == "Bash":
+            command = str(tool_input.get("command") or "").strip()
+            return command.startswith(_READ_ONLY_BASH_PREFIX)
+        if tool_name in _WRITE_TOOLS:
+            candidate = tool_input.get("path") or tool_input.get("file_path")
+            if candidate is None:
+                return False
+            try:
+                Path(candidate).resolve().relative_to(resolved_memory_dir)
+                return True
+            except (ValueError, OSError, RuntimeError, TypeError):
+                return False
+        return False
+
+    return _can_use
+
+
+async def _noop_extraction_executor(
+    context: MemoryExtractionContext, messages: list[Any]
+) -> None:
+    prompt = build_extract_prompt(
+        workspace_cfg=context.workspace_cfg,
+        window_messages=min(DEFAULT_EXTRACTION_WINDOW_MESSAGES, len(messages)),
+    )
+    logger.info(
+        "memory extraction placeholder executor triggered session=%s container=%s prompt_chars=%s message_count=%s",
+        context.session_id,
+        context.container_type,
+        len(prompt),
+        len(messages),
+    )
+    return None
+
+
+_DEFAULT_COORDINATOR = MemoryExtractionCoordinator()
+
+
+async def execute_extract_memories(context: MemoryExtractionContext) -> None:
+    await _DEFAULT_COORDINATOR.execute_extract_memories(context)
+
+
+async def drain_pending_extraction(timeout_ms: int = 60_000) -> None:
+    await _DEFAULT_COORDINATOR.drain_pending_extraction(timeout_ms=timeout_ms)
+
+
+def _is_memory_write_call(call: Any, workspace_cfg: AgentWorkspaceConfig) -> bool:
+    if isinstance(call, dict):
+        name = str(call.get("name") or call.get("tool_name") or "")
+        args = call.get("args") or call.get("tool_input") or {}
+    else:
+        name = str(getattr(call, "name", "") or getattr(call, "tool_name", ""))
+        args = getattr(call, "args", None) or getattr(call, "tool_input", None) or {}
+    if name not in _WRITE_TOOLS:
+        return False
+    if not isinstance(args, dict):
+        return False
+    candidate = args.get("path") or args.get("file_path")
+    if not candidate:
+        return False
+    try:
+        resolved = Path(candidate).resolve()
+        if is_agent_memory_path(workspace_cfg, resolved):
+            return True
+        resolved.relative_to(workspace_cfg.memory_dir.resolve())
+        return True
+    except (ValueError, OSError, RuntimeError, TypeError):
+        return False
+
+
+def build_extract_prompt(
+    *,
+    workspace_cfg: AgentWorkspaceConfig,
+    window_messages: int = DEFAULT_EXTRACTION_WINDOW_MESSAGES,
+) -> str:
+    headers = scan_memory_files(workspace_cfg.memory_dir)
+    manifest = format_memory_manifest(headers)
+    opener = (
+        "You are now acting as the memory extraction subagent. Analyze the most recent\n"
+        f"~{window_messages} messages above and use them to update your persistent memory systems.\n\n"
+        "Available tools: Read, Grep, Glob, read-only Bash (ls/find/cat/stat/wc/head/tail\n"
+        "and similar), and Edit/Write for paths inside the memory directory only.\n"
+        "Bash rm is not permitted. All other tools — MCP, Agent, write-capable Bash,\n"
+        "etc — will be denied.\n\n"
+        "You have a limited turn budget. Edit requires a prior Read of the same file,\n"
+        "so the efficient strategy is: turn 1 — issue all Read calls in parallel for\n"
+        "every file you might update; turn 2 — issue all Write/Edit calls in parallel.\n"
+        "Do not interleave reads and writes across multiple turns.\n"
+    )
+    existing = ""
+    if manifest:
+        existing = (
+            "\n\n## Existing memory files\n\n"
+            f"{manifest}\n\n"
+            "Check this list before writing — update an existing file rather than creating\n"
+            "a duplicate."
+        )
+    body = (
+        "If the user explicitly asks you to remember something, save it immediately as\n"
+        "whichever type fits best. If they ask you to forget something, find and remove\n"
+        "the relevant entry.\n\n"
+        f"{TYPES_SECTION_INDIVIDUAL}\n\n"
+        f"{WHAT_NOT_TO_SAVE_SECTION}\n\n"
+        f"{HOW_TO_SAVE_SECTION}"
+    )
+    return opener + existing + "\n\n" + body
+
+
+__all__ = [
+    "MemoryExtractionContext",
+    "MemoryExtractionCoordinator",
+    "build_extract_prompt",
+    "create_auto_mem_can_use_tool",
+    "drain_pending_extraction",
+    "execute_extract_memories",
+]