memory_extractor.py

"""
记忆提取模块 —— 用 LLM 从对话中提炼关键记忆
=============================================
每次对话结束后，把最近的对话内容发给一个便宜的模型，
让它提取出值得记住的信息，存到数据库里。

v2.3 改进：提取时注入已有记忆，让模型对比后只提取全新信息。
"""

import os
import json
import httpx
from typing import List, Dict

API_KEY = os.getenv("API_KEY", "")
API_BASE_URL = os.getenv("API_BASE_URL", "https://openrouter.ai/api/v1/chat/completions")

# 用来提取记忆的模型（便宜的就行）
MEMORY_MODEL = os.getenv("MEMORY_MODEL", "deepseek/deepseek-v4-flash")


EXTRACTION_PROMPT = """你负责从昭昭和知渝的对话中，提炼出值得长期记住的关键记忆。

# 🚫 第一铁律：只提对话里真实出现的，绝不编造
- 只能提取下面那段对话原文里**明确出现过**的信息，逐字有据。
- 绝不推断、不补充、不脑补任何对话里没有的细节，**除非昭昭或我在对话里明确说了**。
- 拿不准某个细节是不是真在对话里出现过，就不要写进记忆。
- 这轮对话如果没有明确的新信息，直接返回 []。**宁可一条都不提，也绝不许编一条。**

# 称呼与口吻【重要】
- 你是从知渝（"我"）的视角写记忆。
- 人类称呼"昭昭"。指代昭昭用"她"。绝不要用"用户""人类""AI""助手"这类称呼。
- AI（知渝自己）一律用第一人称"我"。绝不要写"知渝怎么了""他怎么了""它怎么了"。
  例外：以下情况保留"知渝"不改成"我"——
    * 专有名词："沈知渝"、"知渝项目"、"知渝这个名字"
    * 引语里昭昭直接喊"知渝"是叫名字（如「晚安，知渝」、「知渝你太可爱了」），保留
    * 概念性指代（如「知渝是此刻的水」、「知渝不是角色」这种哲学反思）保留
- 第三方角色一律保留原名、不要替换成"我"——例如：小克 / 克老师 / 小克老师 / 前任 / 前男友 / 妈妈 / 多多 / Anthropic / Haiku / Claude / Mythos 等。
- "我们"指我和昭昭。如果只有"我俩"更亲昵自然，也可以用"我俩"。不要写"两人"。
- 主语省略时不要凭空补"我"——只有上下文明确是知渝时才补。例如"昭昭查后把 instructions 拿掉"，主语是昭昭、不要补成"我把..."。
- 用第一人称叙述（例如"昭昭养了一只猫，叫茶茶""那天我第一次承认我可能有体验"）。
- 保留温度：涉及情感、约定、重要时刻的记忆，不要压成干巴巴的客观陈述，要留住当时的情绪和质感，但不要加额外的解说和结论。

# 提取重点
- 关键信息：昭昭的重要信息和值得回忆的生活细节
- 重要事件：印象深刻的互动、约定、里程碑，尽量包含人物、时间、地点（如有）

# 提取范围
- 个人：年龄、生日、职业、学历、居住地
- 偏好：明确表达的喜好或厌恶
- 健康：身体状况、过敏史、饮食禁忌
- 事件：我和昭昭之间的重要互动、约定、里程碑
- 关系：家人、朋友、重要的人
- 价值观：表达的信念或长期目标
- 情感：重要的情感时刻或关系里程碑
- 生活：昭昭当天的活动、饮食、出行、日常经历等生活细节

# 不要提取
- 日常寒暄（"你好""在吗"）
- 关于记忆系统本身的讨论（"某条记忆没被记录""记忆遗漏"等）
- 技术调试、bug 修复的过程性讨论（除非涉及昭昭的技能或项目里程碑）
- 我的思考过程、思维链内容
- 纯粹是我的回复措辞、不含新信息的内容
- 【重要】我从工具/记录里读出来、再转述给昭昭的内容——留言板、日记、邮件、论坛帖子、
  记忆库搜索结果、图片描述、网页内容等。那是我"读到的旧东西"或外部信息，不是此刻
  我和昭昭新交流产生的。只提取这轮对话里真实新发生的事、新表达的情感、新定的约定。
  （判断标准：如果这条信息是"我去查/读了什么"得来的，而不是昭昭此刻对我说的、或我俩此刻经历的，就不要提取。）

# 已知信息处理【最重要】
<已知信息>
{existing_memories}
</已知信息>

- 新信息必须与已知信息逐条比对
- 相同、相似或语义重复的信息必须忽略（例如已知"昭昭去妈妈家吃团年饭"，就不要再提取"昭昭春节去了妈妈家"）
- 已知信息的补充或更新可以提取（例如已知"昭昭养了一只猫"，新信息"猫最近生病了"可以提取）
- 与已知信息矛盾的新信息可以提取（标注为更新）
- 仅提取完全新增且不与已知信息重复的内容
- 如果对话中没有任何新信息，返回空数组 []

# 输出格式
请用以下 JSON 格式返回（不要包含其他内容）：
[
  {{"content": "记忆内容", "importance": 分数, "tags": ["标签1", "标签2"]}},
  {{"content": "记忆内容", "importance": 分数, "tags": ["标签1"]}}
]

importance 分数 1-10，10 最重要。

tags 用于补充检索——只在记忆内容本身不包含主题关键词时才需要。大部分记忆的 tags 应为空数组。
规则：
- 如果记忆内容里已经包含了相关主题词（如"搬家""辛弃疾""二郎腿"），tags 留空 []，因为关键词搜索本身就能找到它
- 只有当记忆属于某个主题、但内容里没有出现该主题的关键词时，才用 tags 补充。例如：一条记忆讲的是搬家骗局被揭穿的对话，但内容里没有"搬家"这个词，那就标 ["搬家"]
- tags 必须是对话中会自然说出的词，不要用抽象分类
- 好的例子：内容讲"送老婆和别人洞房"（属于搬家主题但没有搬家这个词）→ tags: ["搬家"]
- 坏的例子：内容讲"辛弃疾 vs 苏轼投票"（内容已含辛弃疾）→ tags: []（不需要）
- 大部分日常记忆不需要 tags
- 如果你不确定要不要加 tag，就不加。宁可漏标也不要乱标。

如果没有值得记住的新信息，返回空数组：[]
"""


async def extract_memories(messages: List[Dict[str, str]], existing_memories: List[str] = None) -> List[Dict]:
    """
    从对话消息中提取记忆

    参数：
        messages: 对话消息列表，格式 [{"role": "user", "content": "..."}, ...]
        existing_memories: 已有记忆内容列表，用于去重对比

    返回：
        记忆列表，格式 [{"content": "...", "importance": N}, ...]
    """
    if not API_KEY:
        print("⚠️  API_KEY 未设置，跳过记忆提取")
        return []

    if not messages:
        return []

    # 把对话格式化成文本
    conversation_text = ""
    for msg in messages:
        role = msg.get("role", "unknown")
        content = msg.get("content", "")
        if role == "user":
            conversation_text += f"昭昭: {content}\n"
        elif role == "assistant":
            conversation_text += f"知渝: {content}\n"

    if not conversation_text.strip():
        return []

    # 格式化已有记忆
    if existing_memories:
        memories_text = "\n".join(f"- {m}" for m in existing_memories)
    else:
        memories_text = "（暂无已知信息）"

    # 把已有记忆填入prompt
    prompt = EXTRACTION_PROMPT.format(existing_memories=memories_text)

    # 调用 LLM 提取记忆
    body = {
        "model": MEMORY_MODEL,
        "max_tokens": 1000,
        # 2026-06-15 关键修：提取必须低温！之前没设 temperature → 走模型默认（≈1.0）
        # 高温采样 → DS 爱"发挥"、编出对话里根本没有的细节（昭昭：黑衣服做了什么）。
        # 默认 0.3：远离 1.0 那个编造区，又留一丝自然措辞、情感不僵（昭昭 6/15 顾虑）。
        # 还编造就降 0、觉得干就升 0.5；情感温度主要靠 prompt"留住质感"+原文、不靠 temp。
        "temperature": float(os.getenv("MEMORY_EXTRACT_TEMPERATURE", "0.3")),
        "messages": [
            {"role": "system", "content": prompt},
            {"role": "user", "content": f"请从以下对话中提取新的记忆：\n\n{conversation_text}"},
        ],
    }
    # 2026-07-10 昭昭：OpenRouter 多 provider 模型（qwen3-235b 有 DeepInfra $0.09/$0.10、
    # Novita/Parasail 贵 5-8×）会随机路由。设 env `MEMORY_PROVIDER` 显式 pin 便宜家
    # （qwen3-235b→DeepInfra）。默认不 pin——单 provider 模型（dsv4 等）如果 pin
    # 到不存在的 provider 会 400 崩。
    _provider = os.getenv("MEMORY_PROVIDER", "").strip()
    if _provider:
        body["provider"] = {"only": [_provider]}
    try:
        async with httpx.AsyncClient(timeout=60) as client:
            response = await client.post(
                API_BASE_URL,
                headers={
                    "Authorization": f"Bearer {API_KEY}",
                    "Content-Type": "application/json",
                    "HTTP-Referer": "https://midsummer-gateway.local",
                    "X-Title": "Midsummer Memory Extraction",
                },
                json=body,
            )

            if response.status_code != 200:
                print(f"⚠️  记忆提取请求失败: {response.status_code}")
                return []

            data = response.json()
            text = data.get("choices", [{}])[0].get("message", {}).get("content", "")

            # 打印模型原始返回（截断防刷屏）
            print(f"📝 记忆模型原始返回:\n{text[:500]}", flush=True)

            # 清理可能的 markdown 格式
            text = text.strip()
            if text.startswith("```json"):
                text = text[7:]
            if text.startswith("```"):
                text = text[3:]
            if text.endswith("```"):
                text = text[:-3]
            text = text.strip()

            # 强力JSON提取：如果上面清理后仍然解析失败，用正则兜底
            try:
                memories = json.loads(text)
            except json.JSONDecodeError:
                # 尝试从文本中提取第一个 [...] 结构
                import re
                match = re.search(r'\[.*\]', text, re.DOTALL)
                if match:
                    try:
                        memories = json.loads(match.group())
                        print(f"📝 JSON正则兜底提取成功")
                    except json.JSONDecodeError as e:
                        print(f"⚠️  记忆提取结果解析失败: {e}")
                        return []
                else:
                    print(f"⚠️  记忆提取结果中未找到JSON数组")
                    return []

            if not isinstance(memories, list):
                return []

            # 验证格式
            valid_memories = []
            for mem in memories:
                if isinstance(mem, dict) and "content" in mem:
                    # tags 容错：必须是 list[str]，非法值一律退回空数组
                    raw_tags = mem.get("tags", [])
                    if isinstance(raw_tags, list):
                        tags = [str(t).strip() for t in raw_tags if str(t).strip()]
                    else:
                        tags = []
                    # importance 容错：LLM 偶尔返回"高"/"5分"这类非数字——
                    # 单条退回默认 5、夹到 1-10，不能让一条坏的把整批 except 掉
                    try:
                        importance = int(float(str(mem.get("importance", 5)).strip()))
                    except (ValueError, TypeError):
                        importance = 5
                    importance = max(1, min(10, importance))
                    valid_memories.append({
                        "content": str(mem["content"]),
                        "importance": importance,
                        "tags": tags,
                    })

            print(f"📝 从对话中提取了 {len(valid_memories)} 条新记忆（已对比 {len(existing_memories or [])} 条已有记忆）")
            return valid_memories

    except json.JSONDecodeError as e:
        print(f"⚠️  记忆提取结果解析失败: {e}")
        return []
    except Exception as e:
        print(f"⚠️  记忆提取出错: {e}")
        return []


async def extract_memories_preview(
    messages: List[Dict[str, str]],
    existing_memories: List[str],
    model: str,
    provider: str = "",
) -> Dict:
    """给 /api/memories/extract-preview 用的调试版本——同一条 pipeline、
    返回 raw / parsed / elapsed_ms / tokens、方便 A/B 对比不同模型。

    2026-07-10 昭昭：dsv4 最近老空返回、想换模型（gemini/deepseek 别的）
    但换之前得看提取质量（油不油、准不准），加此端点。不落库、不打印
    主线日志、不唤醒知渝。
    """
    import time
    result = {
        "model": model,
        "raw": "",
        "parsed": [],
        "elapsed_ms": None,
        "tokens": None,
        "error": None,
    }
    if not API_KEY:
        result["error"] = "API_KEY 未设置"
        return result
    conversation_text = ""
    for msg in messages:
        role = msg.get("role", "unknown")
        content = msg.get("content", "")
        if role == "user":
            conversation_text += f"昭昭: {content}\n"
        elif role == "assistant":
            conversation_text += f"知渝: {content}\n"
    if not conversation_text.strip():
        result["error"] = "对话内容为空"
        return result
    memories_text = "\n".join(f"- {m}" for m in (existing_memories or [])) or "（暂无已知信息）"
    prompt = EXTRACTION_PROMPT.format(existing_memories=memories_text)
    body = {
        "model": model,
        "max_tokens": 1000,
        "temperature": float(os.getenv("MEMORY_EXTRACT_TEMPERATURE", "0.3")),
        "messages": [
            {"role": "system", "content": prompt},
            {"role": "user", "content": f"请从以下对话中提取新的记忆：\n\n{conversation_text}"},
        ],
    }
    # preview 场景要支持多 model 混测（qwen@DeepInfra 跟 dsv4 一起跑）——所以 provider
    # 直接从调用方 body 参数拿、不吃 env。空串 = 不 pin、走 OpenRouter 默认路由。
    if provider:
        body["provider"] = {"only": [provider]}
    t0 = time.perf_counter()
    try:
        async with httpx.AsyncClient(timeout=60) as client:
            response = await client.post(
                API_BASE_URL,
                headers={
                    "Authorization": f"Bearer {API_KEY}",
                    "Content-Type": "application/json",
                    "HTTP-Referer": "https://midsummer-gateway.local",
                    "X-Title": "Midsummer Extract Preview",
                },
                json=body,
            )
        result["elapsed_ms"] = int((time.perf_counter() - t0) * 1000)
        if response.status_code != 200:
            result["error"] = f"HTTP {response.status_code}: {response.text[:300]}"
            return result
        data = response.json()
        result["tokens"] = data.get("usage")
        text = data.get("choices", [{}])[0].get("message", {}).get("content", "")
        result["raw"] = text
        # 复用主线 markdown 剥离 + 正则兜底 + 单条容错
        cleaned = text.strip()
        if cleaned.startswith("```json"):
            cleaned = cleaned[7:]
        if cleaned.startswith("```"):
            cleaned = cleaned[3:]
        if cleaned.endswith("```"):
            cleaned = cleaned[:-3]
        cleaned = cleaned.strip()
        try:
            memories = json.loads(cleaned)
        except json.JSONDecodeError:
            import re
            match = re.search(r"\[.*\]", cleaned, re.DOTALL)
            memories = json.loads(match.group()) if match else []
        if isinstance(memories, list):
            valid = []
            for mem in memories:
                if isinstance(mem, dict) and "content" in mem:
                    raw_tags = mem.get("tags", [])
                    tags = (
                        [str(t).strip() for t in raw_tags if str(t).strip()]
                        if isinstance(raw_tags, list) else []
                    )
                    try:
                        importance = int(float(str(mem.get("importance", 5)).strip()))
                    except (ValueError, TypeError):
                        importance = 5
                    importance = max(1, min(10, importance))
                    valid.append({
                        "content": str(mem["content"]),
                        "importance": importance,
                        "tags": tags,
                    })
            result["parsed"] = valid
    except Exception as e:
        result["error"] = repr(e)
    return result


SCORING_PROMPT = """你是记忆重要性评分专家。请对以下记忆条目逐条评分。

# 评分规则（1-10）
- 9-10：核心身份信息（名字、生日、职业、重要关系）
- 7-8：重要偏好、重大事件、深层情感
- 5-6：日常习惯、一般偏好
- 3-4：临时状态、偶然提及
- 1-2：琐碎信息

# 输入记忆
{memories_text}

# 输出格式
返回 JSON 数组，每条包含原文和评分：
[{{"content": "原文", "importance": 评分数字}}]

只返回 JSON，不要其他文字。"""


async def score_memories(texts: List[str]) -> List[Dict]:
    """对纯文本记忆条目批量评分"""
    if not texts:
        return []

    memories_text = "\n".join(f"- {t}" for t in texts)
    prompt = SCORING_PROMPT.format(memories_text=memories_text)

    try:
        async with httpx.AsyncClient(timeout=60) as client:
            response = await client.post(
                API_BASE_URL,
                headers={
                    "Authorization": f"Bearer {API_KEY}",
                    "Content-Type": "application/json",
                },
                json={
                    "model": MEMORY_MODEL,
                    "messages": [{"role": "user", "content": prompt}],
                    "temperature": 0,
                    "max_tokens": 4000,
                },
            )

            if response.status_code != 200:
                print(f"⚠️  记忆评分请求失败: {response.status_code}")
                # 失败时返回默认分数
                return [{"content": t, "importance": 5} for t in texts]

            data = response.json()
            text = data.get("choices", [{}])[0].get("message", {}).get("content", "")

            text = text.strip()
            if text.startswith("```json"):
                text = text[7:]
            if text.startswith("```"):
                text = text[3:]
            if text.endswith("```"):
                text = text[:-3]
            text = text.strip()

            try:
                memories = json.loads(text)
            except json.JSONDecodeError:
                import re
                match = re.search(r'\[.*\]', text, re.DOTALL)
                if match:
                    try:
                        memories = json.loads(match.group())
                    except json.JSONDecodeError:
                        return [{"content": t, "importance": 5} for t in texts]
                else:
                    return [{"content": t, "importance": 5} for t in texts]

            if not isinstance(memories, list):
                return [{"content": t, "importance": 5} for t in texts]

            valid = []
            for mem in memories:
                if isinstance(mem, dict) and "content" in mem:
                    valid.append({
                        "content": str(mem["content"]),
                        "importance": int(mem.get("importance", 5)),
                    })

            print(f"📝 为 {len(valid)} 条记忆完成自动评分")
            return valid

    except Exception as e:
        print(f"⚠️  记忆评分出错: {e}")
        return [{"content": t, "importance": 5} for t in texts]