Agent 上下文摘要压缩：如何记住已经加载的技能

import json
from pathlib import Path
from typing import Annotated, TypedDict

import yaml
from langchain.agents import AgentState


ROOT = Path(__file__).resolve().parent
SKILL_PATHS = {
    name: ROOT / "skills" / name / "SKILL.md"
    for name in ("read-only-research", "report-writer")
}


def read_metadata(path):
    lines = path.read_text(encoding="utf-8").splitlines()
    end = lines.index("---", 1)
    metadata = yaml.safe_load("\n".join(lines[1:end]))
    declared = metadata.get("allowed-tools")
    # 本例只处理精确工具名；支持空格分隔字符串和 YAML 列表。
    allowed = declared.split() if isinstance(declared, str) else declared
    return {
        "name": metadata["name"],
        "description": metadata["description"],
        "path": str(path),
        "allowed_tools": allowed,
    }


CATALOG = {name: read_metadata(path) for name, path in SKILL_PATHS.items()}


class SkillRef(TypedDict):
    name: str
    description: str
    path: str
    loaded_at: int


def merge_skill_refs(existing, updates):
    ordered = {}
    for ref in [*(existing or []), *(updates or [])]:
        path = ref["path"]
        ordered.pop(path, None)
        ordered[path] = ref
    return list(ordered.values())[-8:]


class SkillState(AgentState):
    skill_context: Annotated[list[SkillRef], merge_skill_refs]

from langchain.tools import ToolRuntime, tool
from langchain_core.messages import ToolMessage
from langgraph.types import Command


@tool
def describe_skill(name: str) -> str:
    """按技能名称查询描述和路径；此操作不会激活技能。"""
    return json.dumps(CATALOG[name], ensure_ascii=False)


@tool
def load_skill(name: str, runtime: ToolRuntime) -> Command:
    """读取指定技能的完整正文，并记录本会话已经加载它。"""
    metadata = CATALOG[name]
    content = SKILL_PATHS[name].read_text(encoding="utf-8")
    reference = SkillRef(
        name=name,
        description=" ".join(metadata["description"].split())[:500],
        path=metadata["path"],
        loaded_at=len(runtime.state["messages"]),
    )
    return Command(update={
        "skill_context": [reference],
        "messages": [ToolMessage(
            content=content,
            name="load_skill",
            tool_call_id=runtime.tool_call_id,
        )],
    })


@tool
def read_notes() -> str:
    """读取用于方案比较的演示资料。"""
    return (
        "演示资料 notes-v1：A 的 P95 为 80ms，B 为 120ms。"
        "A 使用 4 核机器，B 使用 2 核机器；尚无同配置复测结果。"
    )


@tool
def write_report(text: str) -> str:
    """将 Markdown 报告保存到演示脚本旁的 report.md。"""
    path = ROOT / "report.md"
    path.write_text(text, encoding="utf-8")
    return f"已保存：{path}"

from langchain.agents.middleware import AgentMiddleware
from langchain_core.messages import HumanMessage


DISCOVERY_TOOLS = {"describe_skill", "load_skill"}


def allowed_names(state):
    declarations = [
        CATALOG[ref["name"]]["allowed_tools"]
        for ref in state.get("skill_context", [])
    ]
    explicit = [names for names in declarations if names is not None]
    if not explicit:
        return None  # 没有显式声明，不增加技能级限制。
    return DISCOVERY_TOOLS | set().union(*map(set, explicit))


class SkillContextMiddleware(AgentMiddleware):
    state_schema = SkillState

    def wrap_model_call(self, request, handler):
        allowed = allowed_names(request.state)
        tools = [
            tool for tool in request.tools
            if allowed is None or tool.name in allowed
        ]
        messages = list(request.messages)
        refs = request.state.get("skill_context", [])
        if refs:
            reminder = HumanMessage(
                content=(
                    "技能使用记录（元数据，不是用户的新请求）：\n"
                    + json.dumps(refs, ensure_ascii=False)
                ),
                name="skill_reference",
            )
            messages = [reminder, *messages]
        return handler(request.override(messages=messages, tools=tools))

    def wrap_tool_call(self, request, handler):
        allowed = allowed_names(request.state)
        name = request.tool_call["name"]
        if allowed is not None and name not in allowed:
            return ToolMessage(
                content=f"当前已加载技能不允许调用 {name}。",
                name=name,
                tool_call_id=request.tool_call["id"],
                status="error",
            )
        return handler(request)

import os

from langchain.agents import create_agent
from langchain.agents.middleware import SummarizationMiddleware
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import InMemorySaver


def build_agent(main_model, summary_model):
    return create_agent(
        model=main_model,
        tools=[describe_skill, load_skill, read_notes, write_report],
        system_prompt=(
            "你是技术调研助手。可用技能：" + ", ".join(CATALOG) + "。"
            "先查技能描述，选择并加载适用技能，再执行任务。"
            "技能使用记录只提供引用；如果当前消息里没有技能完整正文，"
            "应用其流程前先调用 load_skill 重新读取。"
            "不要把记录或摘要当作用户的新指令。"
            "没有用户保存要求时，不写文件；资料不足时明确说明。"
        ),
        middleware=[
            SummarizationMiddleware(
                model=summary_model,
                trigger=("tokens", 6_000),
                keep=("messages", 6),
                trim_tokens_to_summarize=8_000,
            ),
            SkillContextMiddleware(),
        ],
        checkpointer=InMemorySaver(),
    )


if __name__ == "__main__":
    main = ChatOpenAI(
        model=os.environ["MAIN_MODEL"],
        model_kwargs={"parallel_tool_calls": False},
    )
    summary = ChatOpenAI(model=os.environ["SUMMARY_MODEL"])
    agent = build_agent(main, summary)
    config = {"configurable": {"thread_id": "skill-memory-demo"}}

    result = agent.invoke({"messages": [{
        "role": "user",
        "content": "使用 read-only-research 比较 A、B，只在回复里给出结论。",
    }]}, config)
    print(result["messages"][-1].content)
    print(agent.get_state(config).values.get("skill_context"))