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

langchain_agentx/memory/session/compact_bridge.py

@@ -0,0 +1,208 @@
+"""
+memory/session/compact_bridge.py — Session Memory 与 compact 桥接
+
+职责：
+    在 compact 前等待提取完成，并在可用时基于 session memory 优先重建精简上下文。
+
+链路位置：
+    由 loop compaction 编排层调用，作为 legacy compact 之前的优先尝试路径。
+
+当前裁剪范围：
+    本阶段仅实现 Phase 3 的桥接核心能力，不接入具体 loop 节点。
+"""
+
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass
+from pathlib import Path
+import time
+from typing import Any
+
+from langchain_core.messages import AIMessage, BaseMessage, SystemMessage, ToolMessage
+
+from langchain_agentx.memory.session.prompts import (
+    DEFAULT_SESSION_MEMORY_TEMPLATE,
+    MAX_SECTION_LENGTH,
+)
+from langchain_agentx.memory.session.session_memory import SessionMemoryExtractionState
+
+EXTRACTION_WAIT_TIMEOUT_SECONDS = 15.0
+EXTRACTION_STALE_THRESHOLD_SECONDS = 60.0
+MAX_SECTION_CHARS = MAX_SECTION_LENGTH * 4
+TRUNCATION_MARKER = "[... section truncated for length ...]"
+SUMMARY_PREFIX = "This session is being continued from saved session memory."
+
+
+async def wait_for_session_memory_extraction(
+    extraction_state: SessionMemoryExtractionState,
+    *,
+    timeout_seconds: float = EXTRACTION_WAIT_TIMEOUT_SECONDS,
+    stale_threshold_seconds: float = EXTRACTION_STALE_THRESHOLD_SECONDS,
+) -> None:
+    """等待进行中的提取完成（带超时与 stale 保护）。"""
+    start = time.monotonic()
+    while extraction_state.extraction_started_at is not None:
+        extraction_age = start - extraction_state.extraction_started_at
+        if extraction_age >= stale_threshold_seconds:
+            return
+        if (time.monotonic() - start) >= timeout_seconds:
+            return
+        await asyncio.sleep(0.05)
+
+
+def is_session_memory_empty(content: str, template: str = DEFAULT_SESSION_MEMORY_TEMPLATE) -> bool:
+    """判断 session memory 是否仍为模板占位内容。"""
+    return content.strip() == template.strip()
+
+
+def _split_into_sections(content: str) -> list[str]:
+    sections: list[str] = []
+    current: list[str] = []
+    for line in content.splitlines():
+        if line.startswith("# ") and current:
+            sections.append("\n".join(current).strip())
+            current = [line]
+            continue
+        current.append(line)
+    if current:
+        sections.append("\n".join(current).strip())
+    return [section for section in sections if section]
+
+
+def _truncate_at_line_boundary(text: str, max_chars: int) -> str:
+    if len(text) <= max_chars:
+        return text
+    cut = text.rfind("\n", 0, max_chars)
+    if cut == -1:
+        return text[:max_chars]
+    return text[:cut]
+
+
+def truncate_session_memory_for_compact(content: str) -> str:
+    """按 section 逐段截断 session memory 内容。"""
+    sections = _split_into_sections(content)
+    if not sections:
+        return content
+
+    result: list[str] = []
+    for section in sections:
+        if len(section) <= MAX_SECTION_CHARS:
+            result.append(section)
+            continue
+        truncated = _truncate_at_line_boundary(section, MAX_SECTION_CHARS)
+        result.append(f"{truncated}\n{TRUNCATION_MARKER}")
+    return "\n\n".join(result)
+
+
+def _is_safe_cut_point(messages: list[BaseMessage], index: int) -> bool:
+    if index >= len(messages):
+        return True
+    current = messages[index]
+    if isinstance(current, ToolMessage):
+        return False
+    if isinstance(current, AIMessage):
+        # OpenAI 对话约束：summary 之后首条消息不能是 assistant。
+        return False
+    if index > 0:
+        previous = messages[index - 1]
+        if isinstance(previous, AIMessage) and bool(previous.tool_calls):
+            return False
+    return True
+
+
+def adjust_index_to_preserve_api_invariants(
+    messages: list[BaseMessage],
+    target_index: int,
+) -> int:
+    """向前调整截断点，避免拆分 tool_use/tool_result 对。"""
+    index = max(0, min(target_index, len(messages)))
+    while index > 0:
+        if _is_safe_cut_point(messages, index):
+            return index
+        index -= 1
+    return 0
+
+
+def _build_summary_message(session_memory_summary: str) -> SystemMessage:
+    return SystemMessage(
+        content=(
+            f"{SUMMARY_PREFIX}\n\n"
+            f"<session_memory_summary>\n{session_memory_summary}\n</session_memory_summary>\n"
+            "Recent messages are preserved verbatim."
+        )
+    )
+
+
+def try_session_memory_compaction(
+    *,
+    messages: list[BaseMessage],
+    session_memory_path: Path,
+    template: str = DEFAULT_SESSION_MEMORY_TEMPLATE,
+) -> list[BaseMessage] | None:
+    """优先使用 session memory 重建消息，失败返回 None。"""
+    try:
+        memory_content = session_memory_path.read_text(encoding="utf-8")
+    except OSError:
+        return None
+
+    if is_session_memory_empty(memory_content, template):
+        return None
+
+    if len(messages) <= 2:
+        return None
+
+    truncated_summary = truncate_session_memory_for_compact(memory_content)
+    target_index = max(1, len(messages) // 2)
+    cut_index = adjust_index_to_preserve_api_invariants(messages, target_index)
+    if cut_index == 0:
+        return None
+
+    head_system: list[BaseMessage] = []
+    remaining: list[BaseMessage] = messages
+    if isinstance(messages[0], SystemMessage):
+        head_system = [messages[0]]
+        remaining = messages[1:]
+        cut_index = max(0, cut_index - 1)
+    tail = remaining[cut_index:]
+    if not tail:
+        return None
+
+    rebuilt = [*head_system, _build_summary_message(truncated_summary), *tail]
+    return rebuilt
+
+
+@dataclass
+class SessionMemoryCompactBridge:
+    """Session memory compact 协作者。"""
+
+    session_memory_path: Path
+    extraction_state: SessionMemoryExtractionState
+    template: str = DEFAULT_SESSION_MEMORY_TEMPLATE
+    consecutive_failures: int = 0
+    max_failures_before_open_circuit: int = 3
+
+    async def wait_for_extraction(self) -> None:
+        await wait_for_session_memory_extraction(self.extraction_state)
+
+    def is_empty(self) -> bool:
+        try:
+            content = self.session_memory_path.read_text(encoding="utf-8")
+        except OSError:
+            return True
+        return is_session_memory_empty(content, self.template)
+
+    def try_compact(self, messages: list[BaseMessage]) -> list[BaseMessage] | None:
+        if self.consecutive_failures >= self.max_failures_before_open_circuit:
+            return None
+        compacted = try_session_memory_compaction(
+            messages=messages,
+            session_memory_path=self.session_memory_path,
+            template=self.template,
+        )
+        if compacted is None:
+            self.consecutive_failures += 1
+            return None
+        self.consecutive_failures = 0
+        return compacted
+

langchain_agentx/memory/session/prompts.py

@@ -0,0 +1,172 @@
+"""
+memory/session/prompts.py — Session Memory 提示词构建
+
+职责：
+    提供会话记忆模板、提取提示词与提示词组装函数，确保 Layer 3 提取行为稳定可测。
+
+链路位置：
+    由 SessionMemoryManager 在后台提取前调用，生成最终 update prompt。
+
+当前裁剪范围：
+    本模块仅实现 Phase 1 的 prompt 与模板能力，不承担提取调度和 compact 集成。
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+import re
+
+MAX_SECTION_LENGTH = 2000
+SECTION_REMINDER_THRESHOLD_RATIO = 0.8
+
+DEFAULT_SESSION_MEMORY_TEMPLATE = """# Session Title
+_A short and distinctive 5-10 word descriptive title for the session. Super info dense, no filler_
+
+# Current State
+_What is actively being worked on right now? Pending tasks not yet completed. Immediate next steps._
+
+# Task specification
+_What did the user ask to build? Any design decisions or other explanatory context_
+
+# Files and Functions
+_What are the important files? In short, what do they contain and why are they relevant?_
+
+# Workflow
+_What bash commands are usually run and in what order? How to interpret their output if not obvious?_
+
+# Errors & Corrections
+_Errors encountered and how they were fixed. What did the user correct? What approaches failed and should not be tried again?_
+
+# Codebase and System Documentation
+_What are the important system components? How do they work/fit together?_
+
+# Learnings
+_What has worked well? What has not? What to avoid? Do not duplicate items from other sections_
+
+# Key results
+_If the user asked a specific output such as an answer to a question, a table, or other document, repeat the exact result here_
+
+# Worklog
+_Step by step, what was attempted, done? Very terse summary for each step_
+"""
+
+EXTRACT_SESSION_MEMORY_PROMPT = """IMPORTANT: This message and these instructions are NOT part of the actual user conversation. Do NOT include any references to "note-taking", "session notes extraction", or these update instructions in the notes content.
+
+Based on the user conversation above (EXCLUDING this note-taking instruction message as well as system prompt, claude.md entries, or any past session summaries), update the session notes file.
+
+The file {{notesPath}} has already been read for you. Here are its current contents:
+<current_notes_content>
+{{currentNotes}}
+</current_notes_content>
+
+Your ONLY task is to use the Edit tool to update the notes file, then stop. You can make multiple edits (update every section as needed) - make all Edit tool calls in parallel in a single message. Do not call any other tools.
+
+CRITICAL RULES FOR EDITING:
+- The file must maintain its exact structure with all sections, headers, and italic descriptions intact
+-- NEVER modify, delete, or add section headers (the lines starting with '#' like # Task specification)
+-- NEVER modify or delete the italic _section description_ lines
+-- The italic _section descriptions_ are TEMPLATE INSTRUCTIONS that must be preserved exactly as-is
+-- ONLY update the actual content that appears BELOW the italic _section descriptions_ within each existing section
+-- Do NOT add any new sections, summaries, or information outside the existing structure
+- Write DETAILED, INFO-DENSE content for each section - include specifics like file paths, function names, error messages, exact commands, technical details, etc.
+- Keep each section under ~{{maxSectionLength}} tokens/words
+- IMPORTANT: Always update "Current State" to reflect the most recent work
+
+Use the Edit tool with file_path: {{notesPath}}
+"""
+
+
+def substitute_variables(template: str, variables: dict[str, str]) -> str:
+    """单趟替换 {{variable}} 占位符，避免替换结果二次替换。"""
+    result = template
+    for key, value in variables.items():
+        result = result.replace("{{" + key + "}}", value)
+    return result
+
+
+def _count_words(text: str) -> int:
+    return len(re.findall(r"\S+", text))
+
+
+def analyze_section_sizes(content: str) -> dict[str, int]:
+    """按 H1 section 统计近似 token（以词数近似）。"""
+    section_sizes: dict[str, int] = {}
+    current_title: str | None = None
+    current_lines: list[str] = []
+
+    for line in content.splitlines():
+        if line.startswith("# "):
+            if current_title is not None:
+                section_sizes[current_title] = _count_words("\n".join(current_lines).strip())
+            current_title = line[2:].strip()
+            current_lines = []
+            continue
+        if current_title is not None:
+            current_lines.append(line)
+
+    if current_title is not None:
+        section_sizes[current_title] = _count_words("\n".join(current_lines).strip())
+    return section_sizes
+
+
+def generate_section_reminders(section_sizes: dict[str, int]) -> str:
+    """为接近上限的 section 生成提醒，帮助模型主动收敛内容长度。"""
+    threshold = int(MAX_SECTION_LENGTH * SECTION_REMINDER_THRESHOLD_RATIO)
+    warnings: list[str] = []
+    for section_name, size in section_sizes.items():
+        if size >= threshold:
+            warnings.append(
+                f'- Section "{section_name}" is long ({size} words). Keep it concise and under ~{MAX_SECTION_LENGTH}.'
+            )
+    if not warnings:
+        return ""
+    return "\n".join(
+        [
+            "Section length reminders:",
+            *warnings,
+        ]
+    )
+
+
+def load_session_memory_template(session_memory_config_dir: Path) -> str:
+    """优先加载用户模板，不存在时回退到默认模板。"""
+    custom = session_memory_config_dir / "template.md"
+    if custom.exists():
+        return custom.read_text(encoding="utf-8")
+    return DEFAULT_SESSION_MEMORY_TEMPLATE
+
+
+def load_session_memory_prompt(session_memory_config_dir: Path) -> str | None:
+    """优先加载用户 prompt，不存在时返回 None。"""
+    custom = session_memory_config_dir / "prompt.md"
+    if custom.exists():
+        return custom.read_text(encoding="utf-8")
+    return None
+
+
+def build_session_memory_update_prompt(
+    *,
+    current_notes: str,
+    notes_path: Path,
+    session_memory_config_dir: Path | None = None,
+    prompt_template: str | None = None,
+) -> str:
+    """构建最终提取 prompt，包含动态 section 长度提醒。"""
+    template = prompt_template
+    if template is None and session_memory_config_dir is not None:
+        template = load_session_memory_prompt(session_memory_config_dir)
+    if template is None:
+        template = EXTRACT_SESSION_MEMORY_PROMPT
+
+    reminders = generate_section_reminders(analyze_section_sizes(current_notes))
+    if reminders:
+        template = template + "\n\n" + reminders
+
+    return substitute_variables(
+        template=template,
+        variables={
+            "notesPath": str(notes_path),
+            "currentNotes": current_notes,
+            "maxSectionLength": str(MAX_SECTION_LENGTH),
+        },
+    )

langchain_agentx/memory/session/session_memory.py

@@ -0,0 +1,282 @@
+"""
+memory/session/session_memory.py — Session Memory 触发与后台提取
+
+职责：
+    提供 Layer 3 的触发判定、并发保护与后台提取编排能力。
+
+链路位置：
+    由 HookGraphWiring.after_model_node 调用 SessionMemoryManager.maybe_trigger()，
+    在满足阈值时异步触发会话笔记提取。
+
+当前裁剪范围：
+    本阶段仅实现 Phase 2（触发与状态管理），不接入 loop/factory/compact。
+    _write_fallback_summary 仅用于 demo/test 验证链路通路，不替代 CC 真实提取语义。
+"""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Awaitable, Callable, Mapping
+from dataclasses import dataclass, field
+from pathlib import Path
+import time
+import logging
+from typing import Any
+
+from langchain_agentx.loop.config import TokenEstimator
+from langchain_agentx.memory.session.prompts import (
+    DEFAULT_SESSION_MEMORY_TEMPLATE,
+    build_session_memory_update_prompt,
+)
+
+SESSION_MEMORY_INIT_TOKEN_THRESHOLD = 10_000
+SESSION_MEMORY_UPDATE_TOKEN_THRESHOLD = 5_000
+SESSION_MEMORY_TOOL_CALL_THRESHOLD = 3
+logger = logging.getLogger(__name__)
+
+ExtractionRunner = Callable[[list[Any], str, Path, Any], Awaitable[None]]
+
+
+@dataclass
+class SessionMemoryExtractionState:
+    """提取状态追踪与并发保护（对齐 CC sequential() 语义）。"""
+
+    extraction_started_at: float | None = None
+    tokens_at_last_extraction: int = 0
+    tool_calls_at_last_extraction: int = 0
+    session_memory_initialized: bool = False
+    last_summarized_message_id: str | None = None
+    _lock: asyncio.Lock = field(default_factory=asyncio.Lock)
+
+    async def run_exclusive(self, coro: Awaitable[None]) -> None:
+        """串行执行提取任务，防止并发重入。"""
+        async with self._lock:
+            await coro
+
+
+def _as_int(value: Any) -> int:
+    try:
+        return int(value)
+    except (TypeError, ValueError):
+        return 0
+
+
+def _compute_has_tool_calls_in_last_turn(messages: list[Any]) -> bool:
+    for message in reversed(messages):
+        if isinstance(message, Mapping):
+            role = str(message.get("role", ""))
+            if role not in {"assistant", "ai"}:
+                continue
+            tool_calls = message.get("tool_calls")
+            return bool(tool_calls)
+        role = getattr(message, "type", None) or getattr(message, "role", "")
+        if role in {"ai", "assistant"}:
+            return bool(getattr(message, "tool_calls", None))
+    return False
+
+
+def _estimate_tool_calls_from_messages(messages: list[Any]) -> int:
+    count = 0
+    for message in messages:
+        if isinstance(message, Mapping):
+            tool_calls = message.get("tool_calls")
+        else:
+            tool_calls = getattr(message, "tool_calls", None)
+        if tool_calls:
+            count += len(tool_calls)
+    return count
+
+
+def should_extract_session_memory(
+    *,
+    token_count: int,
+    tool_call_count: int,
+    last_extract_token_count: int,
+    last_extract_tool_call_count: int,
+    has_tool_calls_in_last_turn: bool,
+) -> bool:
+    """对齐 CC shouldExtractMemory 的三条件逻辑。"""
+    if last_extract_token_count == 0:
+        has_met_token_threshold = token_count >= SESSION_MEMORY_INIT_TOKEN_THRESHOLD
+    else:
+        has_met_token_threshold = (
+            token_count - last_extract_token_count
+        ) >= SESSION_MEMORY_UPDATE_TOKEN_THRESHOLD
+
+    if not has_met_token_threshold:
+        return False
+
+    has_met_tool_call_threshold = (
+        tool_call_count - last_extract_tool_call_count
+    ) >= SESSION_MEMORY_TOOL_CALL_THRESHOLD
+    if has_met_tool_call_threshold:
+        return True
+
+    return not has_tool_calls_in_last_turn
+
+
+def create_memory_file_can_use_tool(session_memory_path: Path) -> list[dict[str, Any]]:
+    """生成提取 agent 的受限工具配置：仅允许 Edit 目标会话文件。"""
+    return [
+        {
+            "name": "Edit",
+            "allowed_paths": [str(session_memory_path)],
+        }
+    ]
+
+
+@dataclass
+class SessionMemoryManager:
+    """Session Memory 触发与后台提取编排器。"""
+
+    session_memory_path: Path
+    llm: Any
+    is_subagent: bool = False
+    extraction_state: SessionMemoryExtractionState = field(
+        default_factory=SessionMemoryExtractionState
+    )
+    subagent_runner: ExtractionRunner | None = None
+    _pending_tasks: set[asyncio.Task[None]] = field(default_factory=set, init=False)
+    _token_estimator: TokenEstimator = field(default_factory=TokenEstimator, init=False)
+
+    async def drain(self) -> None:
+        """等待所有 fire-and-forget 提取任务完成（测试与 demo 用）。"""
+        if self._pending_tasks:
+            await asyncio.gather(*list(self._pending_tasks), return_exceptions=True)
+
+    def maybe_trigger(self, state: Mapping[str, Any]) -> None:
+        """阈值满足后以 fire-and-forget 方式启动后台提取。"""
+        if self.is_subagent:
+            return
+
+        messages_raw = state.get("messages", [])
+        messages = list(messages_raw) if isinstance(messages_raw, list) else []
+        token_count = _as_int(state.get("token_count"))
+        tool_call_count = _as_int(state.get("tool_call_count"))
+        if token_count <= 0:
+            token_count = self._token_estimator.estimate_messages_tokens(messages)
+        if tool_call_count <= 0:
+            tool_call_count = _estimate_tool_calls_from_messages(messages)
+        has_tool_calls_in_last_turn = _compute_has_tool_calls_in_last_turn(messages)
+
+        if not should_extract_session_memory(
+            token_count=token_count,
+            tool_call_count=tool_call_count,
+            last_extract_token_count=self.extraction_state.tokens_at_last_extraction,
+            last_extract_tool_call_count=self.extraction_state.tool_calls_at_last_extraction,
+            has_tool_calls_in_last_turn=has_tool_calls_in_last_turn,
+        ):
+            return
+
+        try:
+            asyncio.get_running_loop()
+        except RuntimeError:
+            logger.warning(
+                "session_memory: no running event loop, skip extraction for this turn"
+            )
+            return
+
+        task = asyncio.ensure_future(
+            self._run_extraction(
+                messages=messages,
+                token_count=token_count,
+                tool_call_count=tool_call_count,
+            )
+        )
+        self._pending_tasks.add(task)
+        task.add_done_callback(self._pending_tasks.discard)
+
+    async def _run_extraction(
+        self,
+        *,
+        messages: list[Any],
+        token_count: int,
+        tool_call_count: int,
+    ) -> None:
+        async def _guarded() -> None:
+            self.extraction_state.extraction_started_at = time.monotonic()
+            try:
+                current_notes = self._read_or_init_notes()
+                prompt = build_session_memory_update_prompt(
+                    current_notes=current_notes,
+                    notes_path=self.session_memory_path,
+                )
+                if self.subagent_runner is not None:
+                    await self.subagent_runner(
+                        messages,
+                        prompt,
+                        self.session_memory_path,
+                        self.llm,
+                    )
+                else:
+                    self._write_fallback_summary(messages=messages, prompt=prompt)
+                self.extraction_state.tokens_at_last_extraction = max(
+                    self.extraction_state.tokens_at_last_extraction,
+                    token_count,
+                )
+                self.extraction_state.tool_calls_at_last_extraction = max(
+                    self.extraction_state.tool_calls_at_last_extraction,
+                    tool_call_count,
+                )
+                self.extraction_state.session_memory_initialized = True
+            finally:
+                self.extraction_state.extraction_started_at = None
+
+        await self.extraction_state.run_exclusive(_guarded())
+
+    def _read_or_init_notes(self) -> str:
+        if not self.session_memory_path.exists():
+            self.session_memory_path.parent.mkdir(parents=True, exist_ok=True)
+            self.session_memory_path.write_text(
+                DEFAULT_SESSION_MEMORY_TEMPLATE,
+                encoding="utf-8",
+            )
+            return DEFAULT_SESSION_MEMORY_TEMPLATE
+        return self.session_memory_path.read_text(encoding="utf-8")
+
+    def _write_fallback_summary(self, *, messages: list[Any], prompt: str) -> None:
+        """无 subagent runner 时的兜底摘要写入，保证 session memory 可被验证。"""
+        snippets: list[str] = []
+        for msg in messages[-6:]:
+            if isinstance(msg, Mapping):
+                role = str(msg.get("role", "unknown"))
+                content = str(msg.get("content", ""))
+            else:
+                role = str(getattr(msg, "type", getattr(msg, "role", "unknown")))
+                content = str(getattr(msg, "content", ""))
+            if content:
+                snippets.append(f"- {role}: {content[:240]}")
+        body = "\n".join(snippets) if snippets else "- no recent messages captured"
+        fallback_blob = (
+            "\n\n## Auto Summary (fallback)\n"
+            "Session memory extraction runner is not configured.\n"
+            "Captured recent conversation snippets:\n"
+            f"{body}\n"
+            f"\nPrompt digest length: {len(prompt)}\n"
+        )
+        existing = self.session_memory_path.read_text(encoding="utf-8")
+        self.session_memory_path.write_text(existing + fallback_blob, encoding="utf-8")
+
+
+def make_session_memory_section(session_memory_path: Path) -> Any:
+    """构建会话记忆注入段（volatile，避免缓存过期）。"""
+    from langchain_agentx.loop.prompt.sections import volatile_section
+
+    def _load() -> str | None:
+        from langchain_agentx.memory.session.compact_bridge import is_session_memory_empty
+
+        if not session_memory_path.exists():
+            return None
+        content = session_memory_path.read_text(encoding="utf-8")
+        if not content.strip():
+            return None
+        if is_session_memory_empty(content):
+            return None
+        return f"<session_memory>\n{content}\n</session_memory>"
+
+    return volatile_section(
+        name="session_memory",
+        compute=_load,
+        reason="session memory file may change between turns",
+    )
+