🗄️ Long-term Memory（长期记忆）

"""
InMemoryStore 基础使用
功能：创建内存存储并执行基本操作
"""
from langgraph.store.memory import InMemoryStore

# 定义嵌入函数（用于向量搜索）
def embed(texts: list[str]) -> list[list[float]]:
    """简单的嵌入函数（生产环境使用真实模型）"""
    return [[1.0, 2.0, 3.0] for _ in texts]

# 创建 Store
store = InMemoryStore(
    index={
        "embed": embed,  # 嵌入函数
        "dims": 3        # 向量维度
    }
)

# 存储用户信息
store.put(
    ("users",),           # Namespace（元组）
    "user_123",          # Key
    {                    # Value（JSON 可序列化）
        "name": "张三",
        "language": "zh-CN",
        "preferences": {
            "theme": "dark",
            "notifications": True
        }
    }
)

# 读取用户信息
user_info = store.get(("users",), "user_123")

if user_info:
    print(f"用户名: {user_info.value['name']}")
    print(f"语言: {user_info.value['language']}")
    print(f"主题: {user_info.value['preferences']['theme']}")
else:
    print("用户不存在")

"""
put() 方法：写入或更新数据
"""
from langgraph.store.memory import InMemoryStore

store = InMemoryStore()

# 写入用户偏好
store.put(
    ("users", "preferences"),  # Namespace（可以是多层）
    "user_123",               # Key
    {
        "language": "zh-CN",
        "timezone": "Asia/Shanghai",
        "email_notifications": True
    }
)

# 写入订单历史
store.put(
    ("users", "orders"),
    "user_123",
    {
        "orders": [
            {"id": "ORDER-001", "product": "iPhone 15"},
            {"id": "ORDER-002", "product": "AirPods"}
        ],
        "total_spent": 15000.0
    }
)

print("✅ 数据写入成功")

"""
get() 方法：读取特定 Key 的数据
"""
# 读取用户偏好
preferences = store.get(("users", "preferences"), "user_123")

if preferences:
    print(f"语言: {preferences.value['language']}")
    print(f"时区: {preferences.value['timezone']}")
    print(f"邮件通知: {preferences.value['email_notifications']}")

    # 访问元数据
    print(f"创建时间: {preferences.created_at}")
    print(f"更新时间: {preferences.updated_at}")
else:
    print("未找到用户偏好")

"""
search() 方法：跨 Namespace 搜索数据
功能：支持过滤和向量相似度搜索
"""
# 示例 1：使用过滤条件搜索
results = store.search(
    ("users", "preferences"),
    filter={"language": "zh-CN"}  # 过滤条件
)

print(f"找到 {len(results)} 个中文用户")
for item in results:
    print(f"- 用户 {item.key}: {item.value}")

# 示例 2：向量相似度搜索
results = store.search(
    ("knowledge_base",),
    query="如何重置密码？",  # 查询字符串
    limit=5                  # 返回最相关的 5 个结果
)

for i, item in enumerate(results):
    print(f"\n结果 {i+1}:")
    print(f"  Key: {item.key}")
    print(f"  内容: {item.value}")
    # print(f"  相似度: {item.score}")

"""
delete() 方法：删除特定数据
"""
# 删除用户的订单历史
store.delete(("users", "orders"), "user_123")

# 验证删除
orders = store.get(("users", "orders"), "user_123")
if orders is None:
    print("✅ 订单历史已删除")

"""
工具中读取长期记忆示例
功能：根据用户 ID 查询用户偏好
"""
from langchain.tools import tool, ToolRuntime
from langchain.agents import create_agent, AgentState
from langgraph.store.memory import InMemoryStore
from typing_extensions import TypedDict

# 定义上下文 Schema
class Context(TypedDict):
    user_id: str

@tool
def get_user_language(runtime: ToolRuntime[Context]) -> str:
    """获取用户的语言偏好"""
    # 从 Runtime 获取 Store
    store = runtime.store

    # 从 Context 获取 user_id
    user_id = runtime.context["user_id"]

    # 读取用户偏好
    preferences = store.get(("users", "preferences"), user_id)

    if preferences:
        language = preferences.value.get("language", "en-US")
        return f"用户 {user_id} 的语言偏好是: {language}"
    else:
        return f"用户 {user_id} 没有设置语言偏好"

# 创建 Store
store = InMemoryStore()

# 预先存储一些数据
store.put(
    ("users", "preferences"),
    "user_123",
    {"language": "zh-CN", "theme": "dark"}
)

# 创建 Agent
agent = create_agent(
    model="gpt-4o",
    tools=[get_user_language],
    store=store,               # 传递 Store
    context_schema=Context     # 传递 Context Schema
)

# 调用 Agent（提供 context）
result = agent.invoke(
    {
        "messages": [{"role": "user", "content": "查询我的语言设置"}]
    },
    {
        "context": {"user_id": "user_123"}
    }
)

print(result["messages"][-1].content)

"""
工具中写入长期记忆示例
功能：保存用户偏好设置
"""
from langchain.tools import tool, ToolRuntime
from typing_extensions import TypedDict

class UserPreferences(TypedDict):
    language: str
    theme: str
    notifications: bool

@tool
def save_user_preferences(
    preferences: UserPreferences,
    runtime: ToolRuntime[Context]
) -> str:
    """保存用户偏好设置"""
    store = runtime.store
    user_id = runtime.context["user_id"]

    # 写入 Store
    store.put(
        ("users", "preferences"),
        user_id,
        preferences
    )

    return f"成功保存用户 {user_id} 的偏好设置"

@tool
def get_user_preferences(runtime: ToolRuntime[Context]) -> str:
    """获取用户偏好设置"""
    store = runtime.store
    user_id = runtime.context["user_id"]

    preferences = store.get(("users", "preferences"), user_id)

    if preferences:
        prefs = preferences.value
        return f"""
        语言: {prefs.get('language', '未设置')}
        主题: {prefs.get('theme', '未设置')}
        通知: {'开启' if prefs.get('notifications') else '关闭'}
        """
    else:
        return "未找到用户偏好设置"

# 创建支持读写的 Agent
agent = create_agent(
    model="gpt-4o",
    tools=[save_user_preferences, get_user_preferences],
    store=store,
    context_schema=Context
)

# 保存偏好
result1 = agent.invoke(
    {
        "messages": [{
            "role": "user",
            "content": "帮我设置：语言中文、深色主题、开启通知"
        }]
    },
    {"context": {"user_id": "user_456"}}
)

# 查询偏好
result2 = agent.invoke(
    {
        "messages": [{"role": "user", "content": "查询我的偏好设置"}]
    },
    {"context": {"user_id": "user_456"}}
)

print(result2["messages"][-1].content)

"""
Namespace 设计模式示例
功能：演示不同的命名空间组织方式
"""
from langgraph.store.memory import InMemoryStore

store = InMemoryStore()

# 模式 1：按用户组织
store.put(
    ("users", "user_123", "profile"),
    "basic_info",
    {"name": "张三", "email": "[email protected]"}
)

store.put(
    ("users", "user_123", "preferences"),
    "settings",
    {"language": "zh-CN", "theme": "dark"}
)

# 模式 2：按功能模块组织
store.put(
    ("orders", "user_123"),
    "ORDER-001",
    {"product": "iPhone 15", "status": "shipped"}
)

store.put(
    ("orders", "user_123"),
    "ORDER-002",
    {"product": "MacBook Pro", "status": "processing"}
)

# 模式 3：按时间组织
from datetime import datetime
today = datetime.now().strftime("%Y-%m-%d")

store.put(
    ("analytics", "daily", today),
    "user_123",
    {"visits": 5, "actions": 12}
)

# 模式 4：多租户组织
store.put(
    ("tenants", "company_A", "users"),
    "user_001",
    {"name": "员工A", "role": "admin"}
)

store.put(
    ("tenants", "company_B", "users"),
    "user_001",
    {"name": "员工B", "role": "viewer"}
)

# 查询示例：获取用户的所有订单
orders = store.search(("orders", "user_123"))
print(f"用户有 {len(orders)} 个订单")
for order in orders:
    print(f"- {order.key}: {order.value['product']}")

"""
个性化推荐系统完整示例
功能：基于用户历史行为和偏好，提供智能推荐
"""
from langchain.agents import create_agent
from langchain.tools import tool, ToolRuntime
from langgraph.store.memory import InMemoryStore
from typing_extensions import TypedDict
from datetime import datetime
import json

# ==================== 上下文定义 ====================

class Context(TypedDict):
    user_id: str
    session_id: str

# ==================== 初始化 Store ====================

store = InMemoryStore()

# 预置一些用户数据
store.put(
    ("users", "user_001", "profile"),
    "basic",
    {
        "name": "张三",
        "email": "[email protected]",
        "joined_date": "2024-01-15"
    }
)

store.put(
    ("users", "user_001", "preferences"),
    "settings",
    {
        "language": "zh-CN",
        "categories": ["电子产品", "图书", "运动"],
        "price_range": {"min": 100, "max": 5000}
    }
)

store.put(
    ("users", "user_001", "history"),
    "browsed_products",
    {
        "products": [
            {"id": "P001", "name": "iPhone 15", "category": "电子产品", "viewed_at": "2024-12-20"},
            {"id": "P002", "name": "Python 编程", "category": "图书", "viewed_at": "2024-12-21"},
            {"id": "P003", "name": "跑步鞋", "category": "运动", "viewed_at": "2024-12-22"}
        ]
    }
)

store.put(
    ("users", "user_001", "history"),
    "purchased_products",
    {
        "products": [
            {"id": "P002", "name": "Python 编程", "price": 89, "purchased_at": "2024-12-21"}
        ],
        "total_spent": 89
    }
)

# ==================== 工具定义 ====================

@tool
def get_user_profile(runtime: ToolRuntime[Context]) -> str:
    """获取用户基本信息"""
    store = runtime.store
    user_id = runtime.context["user_id"]

    profile = store.get(("users", user_id, "profile"), "basic")

    if profile:
        info = profile.value
        return f"""
        用户姓名：{info['name']}
        邮箱：{info['email']}
        注册时间：{info['joined_date']}
        """
    else:
        return "未找到用户信息"

@tool
def get_user_preferences(runtime: ToolRuntime[Context]) -> str:
    """获取用户偏好设置"""
    store = runtime.store
    user_id = runtime.context["user_id"]

    prefs = store.get(("users", user_id, "preferences"), "settings")

    if prefs:
        data = prefs.value
        return f"""
        偏好语言：{data['language']}
        感兴趣的类别：{', '.join(data['categories'])}
        价格范围：{data['price_range']['min']} - {data['price_range']['max']} 元
        """
    else:
        return "未找到用户偏好"

@tool
def get_browsing_history(runtime: ToolRuntime[Context]) -> str:
    """查询用户浏览历史"""
    store = runtime.store
    user_id = runtime.context["user_id"]

    history = store.get(("users", user_id, "history"), "browsed_products")

    if history:
        products = history.value["products"]
        result = "最近浏览的商品：\n"
        for p in products[-5:]:  # 最近5个
            result += f"- {p['name']} ({p['category']}) - {p['viewed_at']}\n"
        return result
    else:
        return "暂无浏览历史"

@tool
def get_purchase_history(runtime: ToolRuntime[Context]) -> str:
    """查询用户购买历史"""
    store = runtime.store
    user_id = runtime.context["user_id"]

    purchases = store.get(("users", user_id, "history"), "purchased_products")

    if purchases:
        data = purchases.value
        products = data["products"]
        result = f"购买记录（总消费：¥{data['total_spent']}）：\n"
        for p in products:
            result += f"- {p['name']} - ¥{p['price']} - {p['purchased_at']}\n"
        return result
    else:
        return "暂无购买记录"

@tool
def record_product_view(
    product_id: str,
    product_name: str,
    category: str,
    runtime: ToolRuntime[Context]
) -> str:
    """记录用户浏览商品"""
    store = runtime.store
    user_id = runtime.context["user_id"]

    # 获取现有历史
    history = store.get(("users", user_id, "history"), "browsed_products")

    if history:
        products = history.value["products"]
    else:
        products = []

    # 添加新记录
    products.append({
        "id": product_id,
        "name": product_name,
        "category": category,
        "viewed_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    })

    # 保留最近 50 条
    products = products[-50:]

    # 更新存储
    store.put(
        ("users", user_id, "history"),
        "browsed_products",
        {"products": products}
    )

    return f"已记录浏览：{product_name}"

@tool
def update_user_preferences(
    categories: list[str],
    runtime: ToolRuntime[Context]
) -> str:
    """更新用户偏好类别"""
    store = runtime.store
    user_id = runtime.context["user_id"]

    # 获取现有偏好
    prefs = store.get(("users", user_id, "preferences"), "settings")

    if prefs:
        current = prefs.value
        current["categories"] = categories
    else:
        current = {
            "language": "zh-CN",
            "categories": categories,
            "price_range": {"min": 0, "max": 10000}
        }

    # 更新存储
    store.put(
        ("users", user_id, "preferences"),
        "settings",
        current
    )

    return f"已更新偏好类别为：{', '.join(categories)}"

@tool
def get_personalized_recommendations(runtime: ToolRuntime[Context]) -> str:
    """基于用户历史生成个性化推荐"""
    store = runtime.store
    user_id = runtime.context["user_id"]

    # 获取用户偏好
    prefs = store.get(("users", user_id, "preferences"), "settings")

    # 获取浏览历史
    history = store.get(("users", user_id, "history"), "browsed_products")

    if not prefs:
        return "无法生成推荐：未找到用户偏好"

    categories = prefs.value.get("categories", [])

    # 简化的推荐逻辑（实际应用会更复杂）
    recommendations = {
        "电子产品": ["MacBook Pro", "AirPods Pro", "iPad Air"],
        "图书": ["深度学习", "算法导论", "设计模式"],
        "运动": ["网球拍", "健身手环", "瑜伽垫"]
    }

    result = "🎯 为您推荐：\n\n"
    for category in categories:
        if category in recommendations:
            result += f"【{category}】\n"
            for item in recommendations[category][:2]:  # 每个类别推荐2个
                result += f"  • {item}\n"
            result += "\n"

    return result

# ==================== 创建 Agent ====================

recommendation_agent = create_agent(
    model="gpt-4o",
    tools=[
        get_user_profile,
        get_user_preferences,
        get_browsing_history,
        get_purchase_history,
        record_product_view,
        update_user_preferences,
        get_personalized_recommendations
    ],
    store=store,
    context_schema=Context,
    system_prompt="""你是个性化推荐助手。

职责：
1. 了解用户的偏好和历史行为
2. 记录用户的浏览行为
3. 基于用户数据生成个性化推荐
4. 帮助用户更新偏好设置

工作流程：
- 查询用户信息时，使用相应的查询工具
- 用户浏览商品时，记录到历史中
- 生成推荐时，综合考虑偏好和历史
- 始终保持友好和有帮助"""
)

# ==================== 测试场景 ====================

def simulate_user_interaction(user_id: str, query: str):
    """模拟用户交互"""
    print(f"\n{'='*60}")
    print(f"用户: {query}")
    print('='*60)

    result = recommendation_agent.invoke(
        {"messages": [{"role": "user", "content": query}]},
        {"context": {"user_id": user_id, "session_id": "session_001"}}
    )

    answer = result["messages"][-1].content
    print(f"\n🤖 助手:\n{answer}\n")
    return answer

# ==================== 执行测试 ====================

if __name__ == "__main__":
    user_id = "user_001"

    # 场景 1：查询用户信息
    simulate_user_interaction(
        user_id,
        "我的个人信息是什么？"
    )

    # 场景 2：查看偏好和历史
    simulate_user_interaction(
        user_id,
        "我之前浏览过哪些商品？我的偏好是什么？"
    )

    # 场景 3：记录新的浏览
    simulate_user_interaction(
        user_id,
        "我正在查看 MacBook Pro（电子产品类别），帮我记录一下"
    )

    # 场景 4：获取个性化推荐
    simulate_user_interaction(
        user_id,
        "根据我的偏好和历史，给我一些推荐吧"
    )

    # 场景 5：更新偏好
    simulate_user_interaction(
        user_id,
        "我想把偏好类别改成：电子产品、健康、旅游"
    )

    # 场景 6：再次获取推荐（基于新偏好）
    simulate_user_interaction(
        user_id,
        "现在给我推荐一些商品"
    )

    print("\n" + "="*60)
    print("个性化推荐系统测试完成！")
    print("="*60)

# 同时使用 Checkpointer 和 Store
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.store.memory import InMemoryStore

agent = create_agent(
    model="gpt-4o",
    tools=[...],
    checkpointer=InMemorySaver(),  # 短期记忆
    store=InMemoryStore(),         # 长期记忆
    context_schema=Context
)

# 安装 PostgreSQL Store（示例，实际包名可能不同）
pip install langgraph-store-postgres

# 使用 PostgreSQL Store
from langgraph.store.postgres import PostgresStore

DB_URI = "postgresql://user:pass@localhost:5432/langchain"

with PostgresStore.from_conn_string(DB_URI) as store:
    store.setup()  # 创建表结构

    agent = create_agent(
        model="gpt-4o",
        tools=[...],
        store=store
    )

# 个人知识
store.put(("users", "user_123", "notes"), "note_001", {...})

# 团队共享知识
store.put(("teams", "team_A", "knowledge"), "doc_001", {...})

# 全局知识库
store.put(("global", "faq"), "how_to_reset_password", {...})

# 健壮的数据读取
@tool
def get_user_setting(runtime: ToolRuntime[Context]) -> str:
    """获取用户设置（带错误处理）"""
    try:
        store = runtime.store
        user_id = runtime.context["user_id"]

        settings = store.get(("users", user_id, "settings"), "config")

        if settings:
            return f"设置：{settings.value}"
        else:
            # 返回默认值
            return "使用默认设置"
    except Exception as e:
        # 记录错误并返回友好消息
        print(f"Error reading settings: {e}")
        return "无法读取设置，请稍后重试"