大模型 Agent 长短期记忆管理基础入门

# 确认 Python 版本（建议 3.10+）
import sys
print(f"Python 版本: {sys.version}")

!pip install -r requirements.txt

from dotenv import load_dotenv
from openai import OpenAI
import os

load_dotenv()

# 使用 DEEPSEEK_* 环境变量，与 mini-OpenClaw 源码保持一致
api_key = os.getenv("DEEPSEEK_API_KEY")
base_url = os.getenv("DEEPSEEK_BASE_URL", "https://api.deepseek.com")

if api_key:
    print(f"API Key 加载成功（前8位：{api_key[:8]}...）")
    print(f"Base URL：{base_url}")
else:
    print("API Key 未找到，请检查 .env 文件中的 DEEPSEEK_API_KEY")

# 初始化全局 client，后续所有代码直接使用此对象
# 注意：虽然用 OpenAI SDK，但 base_url 指向 DeepSeek（兼容 OpenAI 协议）
client = OpenAI(api_key=api_key, base_url=base_url)
MODEL = "deepseek-chat"   # DeepSeek 主力对话模型
print(f"client 初始化完成，模型：{MODEL}")

from openai import OpenAI

# client 和 MODEL 已在「环境准备」cell 中初始化，此处直接使用

def chat_no_memory(user_input: str) -> str:
    """无记忆 Agent：每次调用完全独立，不携带任何历史"""
    response = client.chat.completions.create(
        model=MODEL,
        messages=[
            {"role": "user", "content": user_input}
        ]
    )
    return response.choices[0].message.content

# 第一轮：用户自我介绍
print("第一轮：", chat_no_memory("我叫小明，我在学 Python，方向是数据科学"))

# 第二轮：问 Agent 记不记得
print("第二轮：", chat_no_memory("我叫什么名字？我在学什么？"))

messages = []  # 短期记忆的本质：就是这个 list

def chat_with_memory(user_input: str) -> str:
    """有记忆 Agent：把完整对话历史传给 LLM"""
    messages.append({"role": "user", "content": user_input})

    response = client.chat.completions.create(
        model=MODEL,
        messages=messages
    )

    reply = response.choices[0].message.content
    messages.append({"role": "assistant", "content": reply})
    return reply

# 同样的两轮对话
print("第一轮：", chat_with_memory("我叫小明，我在学 Python，方向是数据科学"))
print("第二轮：", chat_with_memory("我叫什么名字？我在学什么？"))

import json
import os
import time

SESSIONS_DIR = "./sessions"  # 会话文件存储目录
os.makedirs(SESSIONS_DIR, exist_ok=True)

def load_session(session_id: str) -> dict:
    """加载会话，不存在则创建新会话"""
    path = os.path.join(SESSIONS_DIR, f"{session_id}.json")
    if os.path.exists(path):
        with open(path, "r", encoding="utf-8") as f:
            return json.load(f)
    # 新会话的初始结构
    return {
        "title": "",
        "created_at": int(time.time()),
        "updated_at": int(time.time()),
        "compressed_context": "",   # 压缩摘要（暂时为空，3.3节会填充）
        "messages": []              # 当前对话历史
    }

def save_session(session_id: str, session: dict):
    """保存会话到 JSON 文件"""
    session["updated_at"] = int(time.time())
    path = os.path.join(SESSIONS_DIR, f"{session_id}.json")
    with open(path, "w", encoding="utf-8") as f:
        json.dump(session, f, ensure_ascii=False, indent=2)

# 测试：加载会话 → 添加消息 → 保存 → 重新加载验证
session = load_session("test_user_001")
session["messages"].append({"role": "user", "content": "我叫小明"})
save_session("test_user_001", session)

# 重新加载，验证持久化
loaded = load_session("test_user_001")
print(f"持久化验证：{loaded['messages'][-1]['content']}")

MAX_HISTORY = 20  # mini-OpenClaw 默认值，可根据模型上下文窗口调整

def get_messages_for_llm(session: dict) -> list:
    """
    构造实际传给 LLM 的消息列表。
    策略：如果有压缩摘要，先注入摘要；再取最近 MAX_HISTORY 条。
    """
    messages_to_send = []

    # 如果存在压缩摘要，作为 system 消息插在最前面
    if session.get("compressed_context"):
        messages_to_send.append({
            "role": "system",
            "content": session["compressed_context"]
        })

    # 只取最近 MAX_HISTORY 条对话
    recent_messages = session["messages"][-MAX_HISTORY:]
    messages_to_send.extend(recent_messages)

    return messages_to_send

# 验证：超过 20 条时只取最近 20 条
test_session = {"compressed_context": "", "messages": [{"role": "user", "content": f"消息{i}"} for i in range(30)]}
result = get_messages_for_llm(test_session)
print(f"截断验证：传入 LLM 的消息数 = {len(result)}（应为 20）")
print(f"   第一条内容：{result[0]['content']}（应为 消息10）")

COMPRESSED_CONTEXT_PREFIX = "[以下是之前对话的摘要]"

def compress_session(session: dict, client) -> dict:
    """
    压缩会话历史。
    策略：取前 50% 的消息作为待压缩部分（至少 4 条），
    调用 LLM 生成摘要后更新 compressed_context 字段，保留后 50% 的消息继续使用。
    """
    messages = session["messages"]
    if len(messages) < 4:
        return session  # 消息太少，不压缩

    # 取前 50% 作为待压缩部分（至少 4 条）
    compress_count = max(4, len(messages) // 2)
    to_compress = messages[:compress_count]
    to_keep = messages[compress_count:]

    # 构造压缩 prompt
    history_text = "\n".join(
        f"{m['role'].upper()}: {m['content']}" for m in to_compress
    )
    # 滚动摘要：把已有摘要 + 新消息一起压缩成统一摘要                                                                                                                                                 
    existing = session.get("compressed_context", "")
    if existing:                                                                                                                                                                                      
        to_compress_text = f"[之前的摘要]\n{existing}\n\n[新增对话]\n{history_text}"
    else:                                                                                                                                                                                             
        to_compress_text = history_text
                                                                                                                                                                                                    
    compress_prompt = f"""请将以下内容压缩成一段简洁的统一摘要，保留所有关键信息（用户身份、重要决策、技术细节、未完成的任务）。                                                                      
                            
    对话历史：                                                                                                                                                                                        
    {to_compress_text}                             
                            
    要求：           
    - 用第三人称描述（「用户」「助手」）
    - 保留所有关键事实，不遗漏重要细节                                                                                                                                                                
    - 100-200字以内
    - 以「[以下是之前对话的摘要]」开头"""                                                                                                                                                             
                                                                                                                                                                                                    
    response = client.chat.completions.create(
        model=MODEL,                                                                                                                                                                                  
        messages=[{"role": "user", "content": compress_prompt}],
        timeout=30          
    )                
    summary = response.choices[0].message.content
                                                                                                                                                                                                    
    # 用新的统一摘要替换旧的（不追加）
    session["compressed_context"] = summary                                                                                                                                                           
    session["messages"] = to_keep                  
    return session     

# 演示压缩效果
test_session = {
    "compressed_context": "",
    "messages": [
        {"role": "user", "content": "我叫小明，我在学 Python 数据科学"},
        {"role": "assistant", "content": "你好小明！Python 数据科学很有前途。"},
        {"role": "user", "content": "我已经学完了 pandas 的 groupby"},
        {"role": "assistant", "content": "很好！下一步可以学 sklearn。"},
        {"role": "user", "content": "sklearn 的 Pipeline 怎么用？"},
        {"role": "assistant", "content": "Pipeline 是 sklearn 的核心工具..."},
    ]
}
print(f"压缩前：{len(test_session['messages'])} 条消息")
compressed = compress_session(test_session, client)  # 取消注释可实际运行
print(f"压缩后：{len(compressed['messages'])} 条消息")
print(f"摘要内容：{compressed['compressed_context']}")
print("压缩函数定义完成（取消注释可实际调用 LLM 执行压缩）")

import json, os, time
from openai import OpenAI

SESSIONS_DIR  = "./sessions"  # session会话的存储路径
MAX_HISTORY   = 30      # 消息数达到此值时触发截取，存入messages
COMPRESS_TRIGGER = 20   # 消息数达到此值时触发压缩，存入system Prompt
os.makedirs(SESSIONS_DIR, exist_ok=True)

# 注意：此处直接使用「环境准备」cell 中已初始化的全局 client 和 MODEL 变量
# 不再重复调用 OpenAI()，避免覆盖 base_url 配置导致 DeepSeek API 连接失败

class SessionManager:
    """mini-OpenClaw 短期记忆管理器"""

    def load(self, session_id: str) -> dict:
        path = os.path.join(SESSIONS_DIR, f"{session_id}.json")
        if os.path.exists(path):
            with open(path, "r", encoding="utf-8") as f:
                data = json.load(f)
            # 向后兼容：v1 格式是纯 list，自动迁移到 v2
            if isinstance(data, list):
                return {"title": "", "created_at": 0,
                        "updated_at": 0, "compressed_context": "",
                        "messages": data}
            return data                              # v2 格式直接返回
        # 文件不存在 → 初始化空会话结构
        return {"title": "", "created_at": int(time.time()),
                "updated_at": int(time.time()),
                "compressed_context": "", "messages": []}

    def save(self, session_id: str, session: dict):
        session["updated_at"] = int(time.time())    # 刷新最后更新时间戳
        path = os.path.join(SESSIONS_DIR, f"{session_id}.json")
        with open(path, "w", encoding="utf-8") as f:
            json.dump(session, f, ensure_ascii=False, indent=2)

    def add_message(self, session: dict, role: str, content: str) -> dict:
        """添加一条消息，超过阈值时自动触发压缩"""
        session["messages"].append({"role": role, "content": content})
        # 消息数达到阈值 → 调用 LLM 生成摘要并替换旧消息，释放 context 空间
        if len(session["messages"]) >= COMPRESS_TRIGGER:
            session = compress_session(session, client)
        return session

    def get_messages_for_llm(self, session: dict) -> list:
        """构造传给 LLM 的消息列表（摘要 + 最近 N 条）"""
        result = []
        # 若存在压缩摘要，作为 system 消息放在最前面，补充历史上下文
        if session.get("compressed_context"):
            result.append({"role": "system",
                           "content": session["compressed_context"]})
                           
        # 滑动窗口：只取最近 MAX_HISTORY 条，避免超出 context window
        result.extend(session["messages"][-MAX_HISTORY:])
        return result

# ── 完整演示：存储 + 截断 + 压缩三机制联动 ──
sm = SessionManager()

# 清理旧测试数据，确保从零开始
import shutil
if os.path.exists(os.path.join(SESSIONS_DIR, "demo_full.json")):
    os.remove(os.path.join(SESSIONS_DIR, "demo_full.json"))

session = sm.load("demo_full")
print("="*60)
print("【阶段一】模拟 25 轮对话，观察三机制触发时机")
print("="*60)

for i in range(1, 26):
    question = f"第{i}轮问题：Python 的第{i}个技巧是什么？"
    answer = f"技巧{i}：这是关于 Python 的第{i}个实用技巧的详细解答。"
    session = sm.add_message(session, "user", question)
    session = sm.add_message(session, "assistant", answer)

    # 在关键节点打印状态
    if i in [5, 10, 19, 20, 25]:
        msgs_for_llm = sm.get_messages_for_llm(session)
        print(f"\n--- 第 {i} 轮结束 ---")
        print(f"  messages[] 实际条数：{len(session['messages'])}")
        print(f"  传给 LLM 的条数：  {len(msgs_for_llm)}")
        print(f"  compressed_context：{'有（已压缩）' if session.get('compressed_context') else '无'}")

sm.save("demo_full", session)

print("\n" + "="*60)
print("【阶段二】验证持久化：重新加载并检查")
print("="*60)
reloaded = sm.load("demo_full")
msgs_final = sm.get_messages_for_llm(reloaded)
print(f"  重新加载后 messages[] 条数：{len(reloaded['messages'])}")
print(f"  传给 LLM 的条数：{len(msgs_final)}")
print(f"  compressed_context 前100字：{reloaded['compressed_context'][:100]}..." if reloaded['compressed_context'] else "  compressed_context：无")
print(f"\n  传给 LLM 的第一条消息（应为压缩摘要）：")
print(f"    role: {msgs_final[0]['role']}")
print(f"    content: {msgs_final[0]['content'][:80]}...")

from langchain_deepseek import ChatDeepSeek
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
from dotenv import load_dotenv
import os

load_dotenv()

# mini-OpenClaw 的标准初始化方式
llm = ChatDeepSeek(
    model=os.getenv("DEEPSEEK_MODEL", "deepseek-chat"),
    api_key=os.getenv("DEEPSEEK_API_KEY"),
    base_url=os.getenv("DEEPSEEK_BASE_URL", "https://api.deepseek.com"),
    temperature=0.7,
)
print("ChatDeepSeek 初始化完成")

from langchain.agents import create_agent

# ── 用 create_agent 创建 Agent（与 mini-OpenClaw 源码一致）──────────
# 注意：不传 checkpointer 参数，和 mini-OpenClaw 的 agent.py 第81行用法完全一样
agent = create_agent(
    model=llm,
    tools=[],                                    # 暂不绑定工具，聚焦记忆行为
    system_prompt="你是一个测试助手，请简短回答。",
)
print(f"Agent 创建成功，类型: {type(agent).__name__}")

# ── 测试：不传 checkpointer 时，Agent 有没有自动记忆 ──────────
# 第一轮：告诉 Agent 名字
result1 = agent.invoke({"messages": [HumanMessage(content="我叫小明，请记住我的名字")]})
print("第一轮回复：", result1["messages"][-1].content)

# 第二轮：问 Agent 记不记得（注意：不手动传入第一轮的历史）
result2 = agent.invoke({"messages": [HumanMessage(content="我叫什么名字？")]})
print("第二轮回复：", result2["messages"][-1].content)
print()
# 验证结果
if "小明" in result2["messages"][-1].content:
    print("Agent 记住了（有自动记忆）")
else:
    print("Agent 没记住 → create_agent 默认无记忆，每次调用独立无状态")

from langgraph.checkpoint.memory import InMemorySaver

# ── 对照组：传入 InMemorySaver 作为 checkpointer ──────────
agent_with_ckpt = create_agent(
    model=llm,
    tools=[],
    system_prompt="你是一个测试助手，请简短回答。",
    checkpointer=InMemorySaver(),  # 关键区别：传入了 checkpointer
)

# 使用 checkpointer 时，必须通过 config 指定 thread_id（相当于会话 ID）
config = {"configurable": {"thread_id": "test-thread-001"}}

# 第一轮
r1 = agent_with_ckpt.invoke(
    {"messages": [HumanMessage(content="我叫小明，请记住我的名字")]},
    config=config,
)
print("第一轮回复：", r1["messages"][-1].content)

# 第二轮：同一个 thread_id，不手动传历史
r2 = agent_with_ckpt.invoke(
    {"messages": [HumanMessage(content="我叫什么名字？")]},
    config=config,  # 同一个 thread_id → 自动恢复上次对话状态
)
print("第二轮回复：", r2["messages"][-1].content)
print()
if "小明" in r2["messages"][-1].content:
    print("Agent 记住了 → InMemorySaver 自动保存了对话状态")

# 压缩专用 LLM（低 temperature 保证摘要质量）
llm_compress = ChatDeepSeek(
    model=os.getenv("DEEPSEEK_MODEL", "deepseek-chat"),
    api_key=os.getenv("DEEPSEEK_API_KEY"),
    base_url=os.getenv("DEEPSEEK_BASE_URL", "https://api.deepseek.com"),
    temperature=0.3,
)

def compress_session_lc(session: dict) -> dict:  
    """用 LLM 压缩对话历史（LangChain 版，prompt 与 3.3 节完全一致）"""                                                                                                                                                     
    msgs = session["messages"]                     
    if len(msgs) < 4:       
        return session

    compress_count = max(4, len(msgs) // 2)                                                                                                                                                           
    to_compress = msgs[:compress_count]
    to_keep = msgs[compress_count:]                                                                                                                                                                   
                                                    
    history_text = "\n".join(
        f"{m['role'].upper()}: {m['content']}" for m in to_compress
    )                                                                                                                                                                                                 

    # 滚动摘要：把已有摘要 + 新消息一起压缩成统一摘要                                                                                                                                                 
    existing = session.get("compressed_context", "")
    if existing:                                                                                                                                                                                      
        to_compress_text = f"[之前的摘要]\n{existing}\n\n[新增对话]\n{history_text}"
    else:                                                                                                                                                                                             
        to_compress_text = history_text
                                                                                                                                                                                                        
    prompt = f"""请将以下内容压缩成一段简洁的统一摘要，保留所有关键信息（用户身份、重要决策、技术细节、未完成的任务）。                                                                               
                              
    对话历史：                                                                                                                                                                                            
    {to_compress_text}                                 
                                
    要求：               
    - 用第三人称描述（「用户」「助手」）
    - 保留所有关键事实，不遗漏重要细节
    - 100-200字以内                                                                                                                                                                                       
    - 以「[以下是之前对话的摘要]」开头"""
                                                                                                                                                                                                        
    result = llm_compress.invoke([HumanMessage(content=prompt)])
                            
    # 用新的统一摘要替换旧的（不追加）                                                                                                                                                                
    session["compressed_context"] = result.content.strip()
    session["messages"] = to_keep                                                                                                                                                                     
    return session

print("compress_session_lc 定义完成")

# ── 测试压缩逻辑 ──────────────────────────────────────────
                                                                                                                                                                                                        
# 构造测试 session（8条消息，足够触发压缩）                                                                                                                                                           
test_session = {                                                                                                                                                                                      
    "compressed_context": "",                                                                                                                                                                         
    "messages": [                                                                                                                                                                                     
        {"role": "user",      "content": "我叫小明，正在学 Python 数据科学"},
        {"role": "assistant", "content": "你好小明！Python 数据科学很有前途。"},                                                                                                                      
        {"role": "user",      "content": "我已经学完了 pandas 的 groupby"},                                                                                                                           
        {"role": "assistant", "content": "很好！下一步可以学 sklearn。"},                                                                                                                             
        {"role": "user",      "content": "sklearn 的 Pipeline 怎么用？"},                                                                                                                             
        {"role": "assistant", "content": "Pipeline 是 sklearn 的核心工具，可以把预处理和模型串联起来。"},                                                                                             
        {"role": "user",      "content": "我现在在做一个房价预测项目"},                                                                                                                               
        {"role": "assistant", "content": "房价预测很适合用 Pipeline，建议先做特征工程。"},                                                                                                            
    ]                                                                                                                                                                                                 
}                                                                                                                                                                                                     
                                                                                                                                                                                                    
print("【第1次压缩】")                                                                                                                                                                                
print(f"压缩前：messages={len(test_session['messages'])}条，compressed_context={'有' if test_session['compressed_context'] else '无'}")
                                                                                                                                                                                                    
result1 = compress_session_lc(test_session)                                                                                                                                                           
                                                                                                                                                                                                    
print(f"压缩后：messages={len(result1['messages'])}条")                                                                                                                                               
print(f"compressed_context 内容：\n{result1['compressed_context']}")
                                                                                                                                                                                                    
# 验证：messages 减少，compressed_context 有内容                                                                                                                                                      
assert len(result1["messages"]) < 8, "压缩后消息数应减少"                                                                                                                                             
assert result1["compressed_context"], "compressed_context 不应为空"                                                                                                                                   
                                                                                                                                                                                                    
# ── 测试滚动摘要：第2次压缩是否合并了旧摘要 ──                                                                                                                                                       
# 再往 session 里加消息，模拟第2次触发                                                                                                                                                                
result1["messages"] += [                                                                                                                                                                              
    {"role": "user",      "content": "特征工程具体做哪些？"},
    {"role": "assistant", "content": "主要做缺失值填充、异常值处理、特征编码。"},                                                                                                                     
    {"role": "user",      "content": "我的数据集有 10 万条记录"},                                                                                                                                     
    {"role": "assistant", "content": "10万条足够了，注意训练集测试集 8:2 切分。"},                                                                                                                    
]                                                                                                                                                                                                     
                                                                                                                                                                                                    
print("\n【第2次压缩（滚动摘要验证）】")                                                                                                                                                              
print(f"压缩前：messages={len(result1['messages'])}条，已有摘要={'是' if result1['compressed_context'] else '否'}")
                                                                                                                                                                                                    
result2 = compress_session_lc(result1)             
                                                                                                                                                                                                    
print(f"压缩后：messages={len(result2['messages'])}条")
print(f"新摘要是否包含「小明」：{'是 ' if '小明' in result2['compressed_context'] else '否 '}")
print(f"新摘要是否提到「房价」：{'是 ' if '房价' in result2['compressed_context'] else '否 '}")                                                                                                   
print(f"compressed_context 内容：\n{result2['compressed_context']}")                                                                                                                                  
                                                                                                                                                                                                    
assert "小明" in result2["compressed_context"], "滚动摘要应保留第1次压缩中的用户身份信息"                                                                                                             
print("\n所有断言通过 ") 

from llama_index.core import Document, VectorStoreIndex
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.settings import Settings
from llama_index.embeddings.openai import OpenAIEmbedding
from dotenv import load_dotenv
import os

load_dotenv()

# 配置 Embedding 模型（与 mini-OpenClaw memory_indexer.py 一致）
# 注意：Embedding 用 OPENAI_API_KEY，与 LLM 的 DEEPSEEK_API_KEY 分开
Settings.embed_model = OpenAIEmbedding(
    model=os.getenv("EMBEDDING_MODEL", "text-embedding-3-small"),
    api_key=os.getenv("OPENAI_API_KEY"),
    api_base=os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1"),
)

# 构造记忆条目（模拟 Agent 积累的用户画像）
memory_texts = [
    "用户叫小明，是一名 Python 开发者，主要做数据分析",
    "用户正在学习 LangChain，上次讨论到 RAG 架构",
    "用户不喜欢 SQL，更偏向用 Pandas 处理数据",
    "用户的项目用 FastAPI 作为后端框架",
    "用户有 3 年 Python 经验，对机器学习基础了解",
]

# 用 Document + SentenceSplitter 构建索引（与 memory_indexer.py 相同流程）
docs = [Document(text=t) for t in memory_texts]
splitter = SentenceSplitter(chunk_size=256, chunk_overlap=32)
nodes = splitter.get_nodes_from_documents(docs)
index = VectorStoreIndex(nodes)
print(f"向量索引创建完成，共 {len(nodes)} 个节点")

# 语义检索：用 retriever 找最相关的 2 条记忆
retriever = index.as_retriever(similarity_top_k=2)
query = "用户用什么数据处理工具"
results = retriever.retrieve(query)

print(f"\n查询：'{query}'")
print("\n检索结果（分数越高越相似）：")
for node in results:
    print(f"  [{node.score:.4f}] {node.text}")

import json
import os

SESSIONS_DIR = "sessions"
os.makedirs(SESSIONS_DIR, exist_ok=True)

def kv_save(key: str, value: dict) -> None:
    """KV 写入：key → JSON 文件"""
    path = os.path.join(SESSIONS_DIR, f"{key}.json")
    with open(path, "w", encoding="utf-8") as f:
        json.dump(value, f, ensure_ascii=False, indent=2)

def kv_load(key: str) -> dict:
    """KV 读取：通过 key 精确取回 value"""
    path = os.path.join(SESSIONS_DIR, f"{key}.json")
    if not os.path.exists(path):
        return {}  # key 不存在，返回空
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)

# 演示：用 user_id 作为 key，存储用户画像
kv_save("user_xiaoming", {
    "name": "小明",
    "role": "Python 开发者",
    "preferences": ["Pandas", "FastAPI"],
    "skill_level": "中级"
})

profile = kv_load("user_xiaoming")
print(f"用户画像：{profile}")

from langchain_deepseek import ChatDeepSeek
from langchain_core.messages import HumanMessage
from dotenv import load_dotenv
import json
import os

load_dotenv()

# LLM 工厂函数：统一管理连接配置，仅 temperature 按场景不同
def create_llm(temperature: float = 0.7) -> ChatDeepSeek:
    return ChatDeepSeek(
        model=os.getenv("DEEPSEEK_MODEL", "deepseek-chat"),
        api_key=os.getenv("DEEPSEEK_API_KEY"),
        base_url=os.getenv("DEEPSEEK_BASE_URL", "https://api.deepseek.com"),
        temperature=temperature,
    )

# 记忆筛选用低 temperature（0.1），减少随机性
llm_judge = create_llm(0.1)

def is_worth_memorizing(conversation_snippet: str) -> tuple[bool, str]:
    """
    判断对话片段是否值得写入长期记忆。
    返回：(是否写入, 提炼后的记忆文本)
    """
    prompt = f"""你是一个记忆筛选助手。请分析以下对话片段，判断其中是否包含值得长期记忆的信息。

    值得长期记忆的信息特征：
    - 用户的身份、职业、技术背景
    - 用户的明确偏好或厌恶
    - 正在进行的项目的关键背景
    - 用户提出的明确要求或约束

    不值得长期记忆的信息：
    - 临时性的任务（「帮我写一段代码」执行完就结束了）
    - 常识性问题的问答
    - 纯粹的闲聊

    对话片段：
    {conversation_snippet}

    请严格用 JSON 格式回复，不要输出其他内容：
    {{"worth_memorizing": true/false, "memory_text": "如果值得记忆，提炼成一句话；否则留空"}}"""

    result = llm_judge.invoke([HumanMessage(content=prompt)])
    parsed = json.loads(result.content)
    return parsed["worth_memorizing"], parsed.get("memory_text", "")

# 测试：两个对话片段，一个值得记忆，一个不值得
snippet1 = "用户：我叫小明，做了 3 年 Python，现在在做一个 FastAPI 项目。\nAgent：了解！"
snippet2 = "用户：Python 的 list comprehension 怎么写？\nAgent：[x for x in range(10)] 这样写。"

for snippet in [snippet1, snippet2]:
    worth, text = is_worth_memorizing(snippet)
    print(f"值得记忆：{worth}")
    if worth:
        print(f"记忆内容：{text}")
    print()

import hashlib
import os

# === 配置常量 ===
MEMORY_FILE = "MEMORY.md"
# 当 MEMORY.md 的估算 Token 数超过此阈值时，从 Direct 注入切换为 RAG 检索
# 阈值设置依据：主流模型 System Prompt 上限约 4K tokens，预留一半给对话历史
MEMORY_TOKEN_THRESHOLD = 2000

# === MD5 缓存 ===
# 避免每次对话轮次都重新读磁盘：只有文件内容变化（MD5 不同）时才重新加载
# 这是 mini-OpenClaw 的性能优化手段——高频调用场景下减少 IO 开销
_memory_cache = {"content": "", "md5": ""}


def _file_md5(path: str) -> str:
    """计算文件 MD5，用于检测 MEMORY.md 是否被外部修改（如用户手动编辑）"""
    if not os.path.exists(path):
        return ""
    with open(path, "rb") as f:
        return hashlib.md5(f.read()).hexdigest()


def load_memory_direct() -> str:
    """
    Direct 注入模式：读取 MEMORY.md 全文，拼接到 System Prompt 中。
    
    工作流程：
    1. 计算当前文件 MD5
    2. 与缓存 MD5 比较——相同则直接返回缓存（跳过磁盘读取）
    3. 不同则重新读取文件，更新缓存
    
    返回值：记忆全文字符串，由调用方拼接到 System Prompt 末尾
    """
    current_md5 = _file_md5(MEMORY_FILE)
    print(f"MD5转换后的值：{current_md5}")

    # 缓存命中：文件未变化，直接返回上次读取的内容
    if current_md5 == _memory_cache["md5"]:
        return _memory_cache["content"]

    if not os.path.exists(MEMORY_FILE):
        return ""

    with open(MEMORY_FILE, "r", encoding="utf-8") as f:
        content = f.read().strip()

    # 更新缓存：下次调用时如果 MD5 相同就不再读磁盘
    _memory_cache["content"] = content
    _memory_cache["md5"] = current_md5
    return content


def append_to_memory(memory_text: str) -> None:
    """
    向 MEMORY.md 追加一条记忆。
    写入后立即使缓存失效（md5 置空），确保下次 load 能读到最新内容。
    """
    with open(MEMORY_FILE, "a", encoding="utf-8") as f:
        f.write(f"- {memory_text}\n")
    # 关键：写入后必须使缓存失效，否则 load_memory_direct 会返回旧内容
    _memory_cache["md5"] = ""


def estimate_tokens(text: str) -> int:
    """
    粗估 Token 数（不依赖 tiktoken，零依赖实现）。
    中文约 1.5 字符/token，英文约 0.75 词/token。
    生产环境建议用 tiktoken 精确计算。
    """
    return int(len(text) / 1.5)


def should_use_rag() -> bool:
    """
    判断是否应从 Direct 切换为 RAG 模式。
    当 MEMORY.md 内容超过 MEMORY_TOKEN_THRESHOLD 时返回 True。
    调用方根据此结果决定：
    - False → load_memory_direct() 全文注入 System Prompt
    - True  → 走 RAG 检索路径，只注入相关片段
    """
    content = load_memory_direct()
    return estimate_tokens(content) > MEMORY_TOKEN_THRESHOLD


# === 演示：模拟 Agent 写入和读取长期记忆 ===
# 确保 MEMORY.md 存在（首次运行时创建）
if not os.path.exists(MEMORY_FILE):
    with open(MEMORY_FILE, "w", encoding="utf-8") as f:
        f.write("# 长期记忆\n\n")

# 模拟 Agent 在对话中积累的用户画像
append_to_memory("用户叫小明，Python 开发者，3 年经验")
append_to_memory("用户偏好 FastAPI + Pandas，不喜欢 SQL")

# 读取并展示当前记忆状态
memory = load_memory_direct()
print(f"当前记忆内容（{estimate_tokens(memory)} tokens 估算）：")
print(memory)
print(f"\n是否需要切换 RAG 模式：{should_use_rag()}")

# ============================================================                                                                                                                                  
# Cell 1：依赖导入 & 写入模拟记忆文件
# ============================================================
from llama_index.core import Document, VectorStoreIndex
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.settings import Settings
from llama_index.embeddings.openai import OpenAIEmbedding
from dotenv import load_dotenv
import os, hashlib

load_dotenv()

# 配置全局 Embedding 模型
Settings.embed_model = OpenAIEmbedding(
    model=os.getenv("EMBEDDING_MODEL", "text-embedding-3-small"),
    api_key=os.getenv("OPENAI_API_KEY"),
    api_base=os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1"),
)

MEMORY_FILE = "MEMORY.md"

# 写入模拟长期记忆（实际项目中由 Agent 自动维护）
memory_content = """# 长期记忆

- 用户叫小明，Python 开发者，3 年经验
- 用户偏好 FastAPI + Pandas，不喜欢 SQL
- 用户正在开发一个数据分析 API 项目
- 用户上次讨论了 LangChain RAG 架构
- 用户习惯用 conda 管理 Python 环境
"""
with open(MEMORY_FILE, "w", encoding="utf-8") as f:
    f.write(memory_content)

print("✅ MEMORY.md 写入完成")
print("文件内容预览：")
print(memory_content)

# ============================================================
# Cell 2：文档切分（Chunking）
# 目标：把 MEMORY.md 切成若干小块（Node），每块是后续 Embedding 的最小单元
# ============================================================

content = open(MEMORY_FILE, "r", encoding="utf-8").read()

# 将整个文件包装成 LlamaIndex Document
doc = Document(text=content, metadata={"source": "MEMORY.md"})

# SentenceSplitter：按句子边界切割
# chunk_size=256   → 每块最多 256 个 token
# chunk_overlap=32 → 相邻块重叠 32 token，防止语义在边界处断裂
splitter = SentenceSplitter(chunk_size=256, chunk_overlap=32)
nodes = splitter.get_nodes_from_documents([doc])

print(f"切分结果：共 {len(nodes)} 个文本块\n")
print("=" * 50)

# 逐一打印每个 Node 的内容（这就是将来被 Embedding 的原始文本）
for i, node in enumerate(nodes):
    print(f"【Node {i}】")
    print(f"  node_id : {node.node_id}")          # LlamaIndex 自动生成的唯一 ID
    print(f"  文本内容: {repr(node.get_text())}")  # repr() 保留换行符可见
    print(f"  token数 : 约 {len(node.get_text().split())} 词")
    print(f"  metadata: {node.metadata}")
    print()

# ============================================================
# Cell 3：对每个 Node 做 Embedding，打印向量形状与前几个维度
# 目标：直观感受"文本 → 浮点数向量"的转换结果
# ============================================================

embed_model = Settings.embed_model

print("正在调用 Embedding API，请稍候...\n")

for i, node in enumerate(nodes):
    text = node.get_text().strip()

    # 调用 Embedding 模型，返回一个浮点数列表（即向量）
    vector = embed_model.get_text_embedding(text)

    print(f"【Node {i}】文本片段（前30字）: {text[:30]}...")
    print(f"  向量维度  : {len(vector)}")                        # text-embedding-3-small → 1536 维
    print(f"  前5个分量 : {[round(v, 6) for v in vector[:5]]}") # 展示向量开头几个值
    print(f"  值域范围  : [{min(vector):.4f}, {max(vector):.4f}]")
    print()

print("💡 说明：每个文本块被压缩成一个固定长度的向量（1536 维）")
print("   语义相近的文本，向量之间的余弦相似度越接近 1.0")

# ============================================================
# Cell 4：构建向量索引，执行语义检索，展示召回结果的溯源信息
# 目标：看清"检索到的内容"究竟对应哪个 Node，相似度得分是多少
# ============================================================

# 构建索引（内部自动对所有 Node 做 Embedding 并建立向量数据库）
index = VectorStoreIndex(nodes)
print("✅ 向量索引构建完成\n")

# 执行语义检索
query = "用户用什么工具做数据处理"
retriever = index.as_retriever(similarity_top_k=2)  # 召回最相似的 2 个 Node
result_nodes = retriever.retrieve(query)

print(f"Query: {query}\n")
print("=" * 50)
print(f"召回 {len(result_nodes)} 个文本块：\n")

for rank, node_with_score in enumerate(result_nodes):
    node   = node_with_score.node        # 原始 Node 对象
    score  = node_with_score.score       # 余弦相似度得分（越高越相关）

    # 找出该 Node 在原始 nodes 列表中的索引位置
    original_idx = next(
        (i for i, n in enumerate(nodes) if n.node_id == node.node_id), -1
    )

    print(f"【召回第 {rank+1} 名】")
    print(f"  来源 Node 索引: Node {original_idx}")   # 对应 Cell 2 中打印的编号
    print(f"  node_id       : {node.node_id}")
    print(f"  相似度得分    : {score:.4f}")            # 0~1，越接近1越相关
    print(f"  文本内容      :\n    {node.get_text().strip()}")
    print()

# ============================================================
# Cell 5：组装最终注入 LLM 的记忆字符串
# 目标：展示从"检索结果"到"可注入 prompt"的最后一步格式化
# ============================================================

def format_rag_memory(result_nodes) -> str:
    """
    将检索结果格式化为 Markdown 字符串。
    该字符串会被拼接到 LLM 的 system prompt 头部，
    让 LLM 在回答时"记得"用户的历史信息。
    """
    if not result_nodes:
        return ""
    lines = [f"- {n.node.get_text().strip()}" for n in result_nodes]
    return "## 相关记忆（语义检索）\n" + "\n".join(lines)


injected_memory = format_rag_memory(result_nodes)

print("最终注入 LLM 的记忆字符串（放入 system prompt）：")
print("-" * 50)
print(injected_memory)
print("-" * 50)

# 模拟拼接 system prompt（实际 Agent 代码中的用法）
system_prompt = f"""你是一个智能助手。以下是用户的历史记忆，请结合这些信息回答问题：

{injected_memory}

请根据以上记忆，回答用户的问题。
"""
print("\n组装后的 system_prompt 示例（前200字）：")
print(system_prompt[:200], "...")

# ============================================================
# Cell 1：构造一份包含"重复、矛盾、冗余"的脏记忆文件
# 目标：为 sleep-time 重组提供有明显整理空间的测试数据
# ============================================================
MEMORY_FILE = "MEMORY.md"

# 故意设计以下几类问题（用注释标出，实际文件不含注释）：
#   [重复] 用户偏好 Pandas 出现了两次
#   [矛盾] 用户喜欢 SQL vs 不喜欢 SQL
#   [冗余] 用户信息分散在多条而非一条
#   [过时] 用户"正在学习 Python"vs"已有 3 年经验"
dirty_memory = """# 长期记忆

- 用户叫小明
- 用户是 Python 开发者
- 用户有大约 3 年的 Python 开发经验
- 用户正在学习 Python 基础（过时）
- 用户偏好使用 Pandas 做数据处理
- 用户不喜欢写 SQL，更喜欢用 Pandas 操作数据
- 用户喜欢用 SQL 做数据查询（与上条矛盾）
- 用户在开发一个数据分析 API 项目
- 用户的项目使用 FastAPI 作为后端框架
- 用户正在开发基于 FastAPI 的数据分析接口（与上条重复）
- 用户上次讨论了 LangChain RAG 架构
- 用户对 LangChain 的 RAG 实现方式很感兴趣（与上条重复）
- 用户习惯用 conda 管理 Python 环境
- 用户使用 conda 创建虚拟环境（与上条重复）
- 用户询问过如何优化向量检索速度
- 用户希望部署到云端，询问过 AWS 和阿里云的费用对比
- 用户表示预算有限，倾向于选择性价比高的方案
"""

with open(MEMORY_FILE, "w", encoding="utf-8") as f:
    f.write(dirty_memory)

# 打印整理前的状态统计
lines = [l for l in dirty_memory.strip().split("\n") if l.startswith("- ")]
print("=" * 50)
print("📋 整理前 MEMORY.md 状态")
print("=" * 50)
print(f"总条目数  : {len(lines)} 条")
print(f"总字符数  : {len(dirty_memory)} 字符")
print(f"\n已知问题（人工标注）：")
print("  [重复] 'Pandas'条目 × 2，'FastAPI 项目'条目 × 2，'conda'条目 × 2，'LangChain RAG'条目 × 2")
print("  [矛盾] 喜欢 SQL vs 不喜欢 SQL")
print("  [过时] '正在学习 Python 基础' vs '3 年经验'")
print(f"\n原始内容：\n{dirty_memory}")

# ============================================================
# Cell 2：执行 sleep-time 记忆重组
# 流程：读取 MEMORY.md → 构造整理 prompt → 调用 LLM → 备份原文件 → 写回整理结果
# ============================================================
from langchain_deepseek import ChatDeepSeek
from langchain_core.messages import HumanMessage
from dotenv import load_dotenv
import os

load_dotenv()

def create_llm(temperature: float = 0.7) -> ChatDeepSeek:
    """
    LLM 工厂函数。
    temperature=0.2：低随机性，保证整理结果稳定、不乱发挥
    """
    return ChatDeepSeek(
        model=os.getenv("DEEPSEEK_MODEL", "deepseek-chat"),
        base_url=os.getenv("DEEPSEEK_BASE_URL", "https://api.deepseek.com"),
        temperature=temperature,
    )

# 整理任务用低 temperature，避免 LLM 随机删掉重要条目
llm_organizer = create_llm(0.2)

def sleep_time_reorganize() -> str:
    """
    sleep-time 记忆重组主函数。

    整理策略（通过 prompt 指令 LLM 执行）：
    1. 去重：相同含义的条目只保留一条
    2. 去矛盾：时间靠后或更具体的信息优先
    3. 去冗余：把分散的同一主题信息合并为一条
    4. 精炼：每条压缩为一句话
    5. 限制总量：最多 20 条，强制聚焦核心信息
    """
    if not os.path.exists(MEMORY_FILE):
        return ""

    with open(MEMORY_FILE, "r", encoding="utf-8") as f:
        current_memory = f.read().strip()

    # 内容太少（<100字符）说明刚写入，不值得整理
    if not current_memory or len(current_memory) < 100:
        return current_memory

    # ---- 构造整理 prompt ----
    # 关键：用明确的规则约束 LLM 的整理行为，避免随意删改
    prompt = f"""你是一个记忆整理助手。请对以下 Agent 长期记忆进行整理：

    整理规则：
    1. 去除重复信息，合并相似条目
    2. 删除过时或矛盾的信息（保留更新的、更具体的）
    3. 将每条记忆精炼为一句话，以 '- ' 开头
    4. 保留文件头 '# 长期记忆'
    5. 最多保留 20 条最重要的记忆

    当前记忆：
    {current_memory}

    请直接输出整理后的完整 Markdown 内容，不要添加任何解释。"""

    print("正在调用 LLM 进行记忆整理，请稍候...\n")
    result = llm_organizer.invoke([HumanMessage(content=prompt)])
    reorganized = result.content.strip()

    # ---- 写回前备份原始内容 ----
    # 防止 LLM 整理出错时无法恢复
    backup_path = MEMORY_FILE + ".bak"
    with open(backup_path, "w", encoding="utf-8") as f:
        f.write(current_memory)

    # 写回整理后的内容
    with open(MEMORY_FILE, "w", encoding="utf-8") as f:
        f.write(reorganized)

    before_lines = current_memory.count("\n")
    after_lines  = reorganized.count("\n")
    print(f" 整理完成：{before_lines} 行 → {after_lines} 行（压缩率 {(1 - after_lines/before_lines)*100:.0f}%）")
    print(f" 原始备份：{backup_path}")
    return reorganized

reorganized = sleep_time_reorganize()

# ============================================================
# Cell 3：整理前后对比，量化整理效果
# ============================================================

# 读取备份（整理前）和当前文件（整理后）
with open(MEMORY_FILE + ".bak", "r", encoding="utf-8") as f:
    before = f.read().strip()

with open(MEMORY_FILE, "r", encoding="utf-8") as f:
    after = f.read().strip()

# 提取条目列表（以"- "开头的行）
before_items = [l.strip() for l in before.split("\n") if l.strip().startswith("- ")]
after_items  = [l.strip() for l in after.split("\n")  if l.strip().startswith("- ")]

# ---- 打印对比报告 ----
print("=" * 60)
print(" sleep-time 记忆重组效果对比")
print("=" * 60)

print(f"\n【数量变化】")
print(f"  整理前：{len(before_items)} 条")
print(f"  整理后：{len(after_items)} 条")
print(f"  减少了：{len(before_items) - len(after_items)} 条（{(1 - len(after_items)/len(before_items))*100:.0f}% 压缩）")

print(f"\n【字符量变化】")
print(f"  整理前：{len(before)} 字符")
print(f"  整理后：{len(after)} 字符")

print("\n" + "-" * 60)
print("整理前（原始脏记忆）：")
print("-" * 60)
for i, item in enumerate(before_items, 1):
    print(f"  {i:2d}. {item}")

print("\n" + "-" * 60)
print("整理后（LLM 提炼结果）：")
print("-" * 60)
for i, item in enumerate(after_items, 1):
    print(f"  {i:2d}. {item}")

print("\n" + "-" * 60)
print("  预期效果验证：")
print("   重复条目（Pandas/FastAPI/conda/LangChain）是否合并为 1 条？")
print("   矛盾条目（喜欢SQL vs 不喜欢SQL）是否保留了正确的一条？")
print("   过时条目（正在学习Python基础）是否被删除？")
print("   分散的用户信息是否合并为一条？")