Long-term Memory
Module Summary

# ── 依赖导入 ──────────────────────────────────────────────
import uuid
from typing import Optional, Literal
from pydantic import BaseModel, Field
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
from langgraph.graph import StateGraph, START, END, MessagesState
from langgraph.store.base import BaseStore
from langgraph.store.memory import InMemoryStore   # 【L1】跨线程 Store
from langgraph.checkpoint.memory import MemorySaver   # 【L1】对话内 Checkpointer
from trustcall import create_extractor             # 【L2】结构化提取库

# ════════════════════════════════════════════════
# 【L2】Profile Schema：单一结构化用户画像
#   · 固定 key "user_profile" → 每次覆写同一份
#   · 所有字段 Optional，新对话只更新提及的字段
# ════════════════════════════════════════════════
class LearnerProfile(BaseModel):
    name:           Optional[str]       = Field(None, description="用户姓名")
    expertise_level: Optional[str]     = Field(None, description="技术水平: beginner/intermediate/advanced")
    current_topics: Optional[list[str]] = Field(None, description="正在学习的技术主题列表")
    preferred_style: Optional[str]     = Field(None, description="偏好的解释风格: concise/detailed/example-driven")
    background:     Optional[str]      = Field(None, description="职业背景或编程经验简述")

# ════════════════════════════════════════════════
# 【L3】Collection Schema：独立学习笔记条目
#   · 每条笔记独立，UUID key → 追加不覆盖
#   · LLM 决定对话中哪些知识点值得记录
# ════════════════════════════════════════════════
class LearningNote(BaseModel):
    topic:   str = Field(description="知识点所属主题（如 LangGraph、Python、ML）")
    content: str = Field(description="具体知识点内容或用户的理解/疑问")
    type:    Literal["concept", "question", "insight"] = Field(
        description="记忆类型：concept=新概念, question=待解答疑问, insight=用户洞察"
    )

# ── 【L1】初始化双存储 ──────────────────────────────────
in_memory_store = InMemoryStore()   # 跨会话长期记忆
checkpointer    = MemorySaver()     # 对话内短期记忆（断点/时间旅行）
llm             = ChatOpenAI(model="gpt-4o-mini")

# ── 【L2】trustcall 提取器：从对话历史提取 LearnerProfile ──
profile_extractor = create_extractor(llm, tools=[LearnerProfile], tool_choice="LearnerProfile")

# ── 【L3】trustcall 提取器：从对话提取多条 LearningNote ──
notes_extractor = create_extractor(llm, tools=[LearningNote], tool_choice="LearningNote")

class LearnerState(MessagesState):
    user_id:        str     # 用户唯一 ID，用于构建 namespace
    profile_ctx:    str     # 从 Store 读出的 Profile 摘要，注入 System
    notes_ctx:      str     # 从 Store 检索出的相关笔记，注入 System

# ════════════════════════════════════════════════
# 【L1 + L2 + L3】retrieve_memories 节点
#   · 接收 store 参数（LangGraph 自动注入）
#   · 从 Store 读取 Profile（L2）和相关笔记（L3）
#   · 将记忆格式化后存入 State，后续节点直接使用
# ════════════════════════════════════════════════
def retrieve_memories(state: LearnerState, store: BaseStore):
    user_id   = state["user_id"]
    last_msg  = state["messages"][-1].content

    # ── 【L2】读取用户画像（Profile）────────────────────────
    # 用固定 key "user_profile" get 单条记忆，无则返回 None
    profile_ns = (user_id, "profile")
    profile_item = store.get(profile_ns, "user_profile")
    profile_ctx = ""
    if profile_item:
        p = profile_item.value
        profile_ctx = f"用户画像：姓名={p.get('name','未知')}, 水平={p.get('expertise_level','未知')}, "
                      f"当前主题={p.get('current_topics','')}, 偏好风格={p.get('preferred_style','')}"

    # ── 【L3】语义检索最相关学习笔记（Collection）──────────
    # 用当前问题作为 query，找到最相关的 3 条历史笔记
    notes_ns    = (user_id, "notes")
    results     = store.search(notes_ns, query=last_msg, limit=3)
    notes_ctx   = ""
    if results:
        notes_lines = [f"[{r.value['type']}] {r.value['topic']}: {r.value['content']}"
                       for r in results]
        notes_ctx = "相关学习记录：\n" + "\n".join(notes_lines)

    return {"profile_ctx": profile_ctx, "notes_ctx": notes_ctx}

# ════════════════════════════════════════════════
# 【L4】respond 节点：Memory Agent 的核心
#   · LLM 读取 Profile + 相关笔记，个性化作答
#   · 可选地调用 save_memory 工具主动触发记忆保存
#   · 也可以主动调用 search_memory 工具获取更多记忆
# ════════════════════════════════════════════════

# ── 【L4】供 LLM 调用的记忆工具 ──────────────────────────
from langchain_core.tools import tool

@tool
def save_memory(topic: str, content: str, memory_type: str) -> str:
    """将重要知识点或用户信息保存到长期记忆。
    用于：用户分享了重要背景信息、学到了新概念、提出了关键疑问时。
    memory_type: 'concept' / 'question' / 'insight'
    """
    return f"已请求保存记忆：[{memory_type}] {topic}: {content}"

@tool
def search_memory(query: str) -> str:
    """主动检索用户的历史学习记录。
    用于：需要回顾用户之前的问题或理解时。
    """
    return f"已请求检索记忆：{query}"

agent_llm = llm.bind_tools([save_memory, search_memory])

def respond(state: LearnerState):
    """【L4】核心回复节点：根据用户画像和历史笔记个性化作答"""
    profile_ctx = state.get("profile_ctx", "")
    notes_ctx   = state.get("notes_ctx", "")

    system_parts = ["你是一个个性化学习伴侣，根据用户的知识背景和历史笔记定制回答。"]
    if profile_ctx:
        system_parts.append(f"\n{profile_ctx}")
    if notes_ctx:
        system_parts.append(f"\n{notes_ctx}")
    system_parts.append(
        "\n当对话中出现值得记录的新知识点、用户疑问或重要背景信息时，"
        "请主动调用 save_memory 工具。这些记录将在未来对话中帮助你更好地辅导用户。"
    )

    response = agent_llm.invoke([
        SystemMessage(content="\n".join(system_parts)),
        *state["messages"],
    ])
    return {"messages": [response]}

# ════════════════════════════════════════════════
# 【L2】update_profile 节点
#   · 用 trustcall 从完整对话历史提取 LearnerProfile
#   · 固定 key "user_profile"：原地覆写，始终只有一份画像
#   · patch 模式：只更新对话中提及的字段
# ════════════════════════════════════════════════
async def update_profile(state: LearnerState, store: BaseStore):
    user_id    = state["user_id"]
    profile_ns = (user_id, "profile")

    # 先读取已有 Profile 作为上下文，让 trustcall 做增量更新
    existing = store.get(profile_ns, "user_profile")
    existing_profile = existing.value if existing else {}

    result = await profile_extractor.ainvoke({
        "messages": state["messages"],
        "existing": {"LearnerProfile": existing_profile},  # patch 模式的关键
    })

    if result["responses"]:
        updated = result["responses"][0].model_dump(exclude_none=True)
        # 【L1】写入 Store：固定 key，下次 put 覆写同一条记录
        store.put(profile_ns, "user_profile", updated)

# ════════════════════════════════════════════════
# 【L3 + L4】save_notes 节点
#   · 解析 respond 节点的工具调用结果
#   · 也用 notes_extractor 从对话中提取知识点
#   · UUID key：每条笔记独立追加，不覆盖历史
# ════════════════════════════════════════════════
async def save_notes(state: LearnerState, store: BaseStore):
    user_id  = state["user_id"]
    notes_ns = (user_id, "notes")

    # ── 方式 1：解析 LLM 工具调用（L4：Agent 主动触发）──────
    last_msg = state["messages"][-1]
    if hasattr(last_msg, "tool_calls") and last_msg.tool_calls:
        for tc in last_msg.tool_calls:
            if tc["name"] == "save_memory":
                args = tc["args"]
                note = {"topic": args["topic"], "content": args["content"],
                        "type": args.get("memory_type", "concept")}
                # 【L1 + L3】UUID key：每条笔记独立存储，追加不覆盖
                store.put(notes_ns, str(uuid.uuid4()), note)

    # ── 方式 2：无论 LLM 是否主动触发，也用 extractor 自动提取 ──
    result = await notes_extractor.ainvoke({"messages": state["messages"][-4:]})
    for note_obj in result["responses"]:
        store.put(notes_ns, str(uuid.uuid4()), note_obj.model_dump())

# ════════════════════════════════════════════════
# 【L1 + L4】图构建：双 Store 挂载
#   · checkpointer：对话内状态（断点/时间旅行）
#   · store：跨会话长期记忆（Profile + Notes）
# ════════════════════════════════════════════════
builder = StateGraph(LearnerState)
builder.add_node("retrieve_memories", retrieve_memories)
builder.add_node("respond",           respond)
builder.add_node("update_profile",    update_profile)   # 【L2】异步记忆节点
builder.add_node("save_notes",         save_notes)        # 【L3】异步记忆节点

builder.add_edge(START,                "retrieve_memories")
builder.add_edge("retrieve_memories",  "respond")
builder.add_edge("respond",            "update_profile")  # respond 后更新 Profile
builder.add_edge("respond",            "save_notes")      # respond 后保存笔记
builder.add_edge("update_profile",    END)
builder.add_edge("save_notes",         END)

# 【L1 + L4】同时挂载 checkpointer 和 store
graph = builder.compile(checkpointer=checkpointer, store=in_memory_store)

# ── 多轮对话示例 ──────────────────────────────────────────
user_id   = "learner_tony"
config_1  = {"configurable": {"thread_id": "session_001"}}

# ── 第一次会话：介绍自己 ──
result = graph.invoke({
    "user_id":   user_id,
    "messages": [HumanMessage(content="我叫 Tony，是个 Python 工程师，最近在学 LangGraph，对 Memory 机制特别感兴趣")]
}, config=config_1)
# → Profile 已更新: name=Tony, background=Python工程师, current_topics=[LangGraph]
# → Notes 新增: [concept] LangGraph/Memory 机制

# ── 第二次会话（新 thread，但 Store 记忆持久化）──
config_2 = {"configurable": {"thread_id": "session_002"}}
result2 = graph.invoke({
    "user_id":   user_id,
    "messages": [HumanMessage(content="InMemoryStore 和 Checkpointer 有什么区别？")]
}, config=config_2)
# retrieve_memories 自动搜索到：Tony 是 Python 工程师，正在学 LangGraph + Memory
# respond 节点用这些背景个性化回答，不再是"对陌生人的通用解释"
# → LLM 可能调用 save_memory 记录这个"question"类型的笔记

# ── 验证 Profile 持久化 ──
stored_profile = in_memory_store.get((user_id, "profile"), "user_profile")
print(stored_profile.value)
# {'name': 'Tony', 'background': 'Python工程师', 'current_topics': ['LangGraph'], ...}

# ── 验证 Notes 积累 ──
all_notes = in_memory_store.search((user_id, "notes"), query="LangGraph memory")
for note in all_notes:
    print(f"[{note.value['type']}] {note.value['topic']}: {note.value['content']}")
# [concept] LangGraph: Memory 机制 - Checkpointer vs Store 的区别
# [question] LangGraph: InMemoryStore 与 Checkpointer 的具体区分