memory_system.py

import os
import json
import time
from datetime import datetime

# 配置
MEMORY_DIR = "user_memory"
USER_ID = "default_user"  # 后期可做多用户隔离
EXPIRY_DAYS = {
    "profile": 365*10,    # 身份信息：永久
    "preference": 365,    # 喜好：1年
    "important": 365*2,    # 过敏/禁忌：2年
    "order": 180,         # 订单：半年
    "chat": 30            # 闲聊：1个月
}

os.makedirs(MEMORY_DIR, exist_ok=True)

# ==============================
# 1. 结构化 JSON 记忆（核心）
# ==============================
def get_memory_path():
    return os.path.join(MEMORY_DIR, f"{USER_ID}.json")

def init_memory():
    path = get_memory_path()
    if not os.path.exists(path):
        default = {
            "profile": {"name": None},           # 身份
            "preference": [],                    # 喜好
            "important": [],                     # 过敏/禁忌
            "order": [],                         # 订单
            "history": [],                       # 短期对话
            "summary": "",                       # 自动总结
            "updated_at": time.time()
        }
        save_memory(default)

def load_memory():
    try:
        with open(get_memory_path(), "r", encoding="utf-8") as f:
            return json.load(f)
    except:
        init_memory()
        return load_memory()

def save_memory(mem):
    mem["updated_at"] = time.time()
    with open(get_memory_path(), "w", encoding="utf-8") as f:
        json.dump(mem, f, ensure_ascii=False, indent=2)

# ==============================
# 2. 去重 + 冲突覆盖 + 过期清理
# ==============================
def clean_expired(items):
    now = time.time()
    valid = []
    for item in items:
        if not isinstance(item, dict):
            valid.append(item)
            continue
        ts = item.get("ts", 0)
        days = (now - ts) / 86400
        if days < EXPIRY_DAYS.get(item.get("type", "chat"), 30):
            valid.append(item)
    return valid

def unique_add(items, new_content, item_type):
    for item in items:
        if isinstance(item, dict) and item.get("content") == new_content:
            return False  # 已存在，不添加
    items.append({
        "content": new_content,
        "type": item_type,
        "ts": time.time()
    })
    return True

# ==============================
# 3. 写入接口（带去重+清理）
# ==============================
def set_name(name):
    mem = load_memory()
    mem["profile"]["name"] = name
    save_memory(mem)

def add_preference(content):
    mem = load_memory()
    mem["preference"] = clean_expired(mem["preference"])
    ok = unique_add(mem["preference"], content, "preference")
    save_memory(mem)
    return ok

def add_important(content):
    mem = load_memory()
    mem["important"] = clean_expired(mem["important"])
    ok = unique_add(mem["important"], content, "important")
    save_memory(mem)
    return ok

def add_order(content):
    mem = load_memory()
    mem["order"] = clean_expired(mem["order"])
    ok = unique_add(mem["order"], content, "order")
    save_memory(mem)
    return ok

# ==============================
# 4. 短期对话 + 自动总结（防token爆炸）
# ==============================
def add_chat_history(user, bot):
    mem = load_memory()
    mem["history"].append({
        "user": user,
        "bot": bot,
        "ts": time.time()
    })
    # 保留最近 6 轮
    if len(mem["history"]) > 6:
        mem["history"] = mem["history"][-6:]
    save_memory(mem)
    auto_summary()

def auto_summary():
    mem = load_memory()
    name = mem["profile"].get("name")
    pref = [x["content"] for x in mem["preference"]]
    imp = [x["content"] for x in mem["important"]]
    order = [x["content"] for x in mem["order"]]
    
    s = ""
    if name: s += f"姓名：{name}。"
    if pref: s += f"喜好：{'、'.join(pref)}。"
    if imp: s += f"禁忌：{'、'.join(imp)}。"
    if order: s += f"订单：{'、'.join(order)}。"
    
    mem["summary"] = s.strip()
    save_memory(mem)

# ==============================
# 5. 读取上下文（给大模型用）
# ==============================
def get_memory_prompt():
    mem = load_memory()
    hist = "\n".join([f"用户：{x['user']}\n客服：{x['bot']}" for x in mem["history"][-4:]])
    return f"""
【用户总结】
{mem['summary']}

【近期对话】
{hist}
"""

# 全局快捷调用
short_memory = type("obj", (), {"load_memory_variables": get_memory_prompt})()