memory_engine.py

"""memory_engine：forgetfulMem 模块 —— 对话记忆流程（SQLite 存取在 memory.py）。

- _summarize：用 forgetfulMem 组模型总结本轮对话关键词（失败回退规则提取）；
- InjectMemoryIntoSystem：按关键词检索记忆注入 system 提示词；
- _store_memory：把本轮对话与关键词写入记忆库。
"""

import hashlib
import json
import re
import threading
from typing import List, Optional

from colorama import Fore

import config as engine_config
import defaultPattern
from api_client import _chat
from chatContent import ChatContent
from messages import _format_history


def _memory_key(system: str) -> str:
    return "sys_" + hashlib.sha1((system or "default").encode("utf-8")).hexdigest()[:16]


def _keywords_of(contents: List[ChatContent]) -> List[str]:
    """从最新非 own 消息中提取关键词（长度≥2 的中英文词）。规则提取兜底。"""
    text = ""
    for t in reversed(contents):
        if not t.ownByMyself and t.text.strip():
            text = t.text
            break
    words = re.findall(r"[\u4e00-\u9fff]{2,}|[A-Za-z0-9_]{3,}", text)
    return [w for w in words if w.strip()][:8]


def _extract_kv_keywords(text: str) -> List[str]:
    """从对象字符串中提取关键词值，如 "{'关键词': '中原雅音'}" → ['中原雅音']。"""
    try:
        import ast
        d = ast.literal_eval(text)
        if isinstance(d, dict):
            v = d.get("关键词") or d.get("keyword") or d.get("word")
            if isinstance(v, str) and v.strip():
                return [v.strip()]
    except Exception:
        pass
    return []


def _parse_keywords(raw: str) -> List[str]:
    """解析模型输出的关键词：JSON 数组 / 逗号·顿号·换行分隔 / 中文引号。解析失败返回空。"""
    if not raw:
        return []
    raw = raw.strip()
    # 先尝试整体 JSON 解析（模型常直接输出数组；元素可能是嵌套列表字符串或对象）
    try:
        obj = json.loads(raw)
        if isinstance(obj, list):
            out: List[str] = []
            for k in obj:
                if isinstance(k, dict):
                    k = k.get("关键词") or k.get("keyword") or k.get("word") or ""
                if isinstance(k, str):
                    ks = k.strip()
                    if ks.startswith("["):
                        out.extend(_parse_keywords(k))  # 嵌套列表字符串再拆一层
                    elif ks.startswith("{") and ("关键词" in ks or "keyword" in ks):
                        out.extend(_extract_kv_keywords(ks))  # 对象字符串
                    else:
                        out.append(ks)
            return [s for s in out if len(s) >= 2][:10]
    except Exception:
        pass
    # 提取 JSON 数组（容忍前后缀文字）
    m = re.search(r'\[(.*?)\]', raw, re.DOTALL)
    body = m.group(1) if m else raw
    # 去引号/标点后按分隔符切分
    body = re.sub(r'["\'“”‘’]', '', body)
    parts = re.split(r'[,，、;；\n]+', body)
    kws = [p.strip() for p in parts
           if len(p.strip()) >= 2 and not re.search(r'[{}:\[\]]', p)]
    return kws[:10]


def _summarize(contents: List[ChatContent]) -> List[str]:
    """forgetfulMem 模型总结本轮对话关键词；失败/未配置 → 规则提取兜底。

    只在 getAnswer 中调用一次，检索注入与入库共用，避免重复调用模型。
    """
    cfg = engine_config.get_config()
    group = cfg.group(engine_config.GROUP_FORGETFUL_MEM)
    fallback = _keywords_of(contents)
    if group is None or not group.models:
        return fallback
    history = _format_history(contents, max(cfg.recent_count_for_judge, 8))
    if not history:
        return fallback
    messages = [
        {"role": "system", "content": defaultPattern.patternMemoryKeywordsSystem.strip()},
        {"role": "user",
         "content": defaultPattern.patternMemoryKeywordsUser.replace("{history}", history).strip()},
    ]
    try:
        resp = _chat(engine_config.GROUP_FORGETFUL_MEM, messages, max_tokens=128)
    except Exception as e:
        print(f"{Fore.RED}[forgetfulMem] 关键词总结失败: {e}{Fore.RESET}")
        return fallback
    if not resp:
        return fallback
    kws = _parse_keywords(resp)
    print(f"{Fore.LIGHTGREEN_EX}[forgetfulMem] 关键词总结: {kws}{Fore.RESET}")
    return kws or fallback


def InjectMemoryIntoSystem(system: str, contents: List[ChatContent],
                           keywords: Optional[List[str]] = None) -> str:
    """按关键词检索记忆并注入 system 提示词（追加【记忆】块）。"""
    cfg = engine_config.get_config()
    store = _get_memory_store()
    memories = store.retrieve(
        _memory_key(system), limit=cfg.memory_inject_limit,
        keywords=keywords if keywords is not None else _keywords_of(contents))
    if not memories:
        return system
    block = "\n".join(f"- {m}" for m in memories)
    # replace 而非 format：记忆内容可能含花括号
    return system + "\n\n" + defaultPattern.patternMemoryBlock.replace("{memory}", block).strip()


def _store_memory(system: str, contents: List[ChatContent], reply: Optional[str],
                  keywords: Optional[List[str]] = None) -> None:
    """把本轮对话（最近若干条 + 回复）写入记忆，关键词一并入库。"""
    cfg = engine_config.get_config()
    store = _get_memory_store()
    entries = []
    recent = contents[-cfg.memory_store_limit:]
    for t in recent:
        s = str(t).strip()
        if s:
            entries.append(s)
    if reply and reply.strip():
        entries.append(f"[机器人回复] {reply.strip()}")
    if entries:
        store.store(_memory_key(system), entries, keywords=keywords)


_memory_store = None
_memory_store_lock = threading.Lock()


def _get_memory_store():
    global _memory_store
    if _memory_store is None:
        with _memory_store_lock:
            if _memory_store is None:
                from memory import MemoryStore
                _memory_store = MemoryStore(engine_config.get_config().memory_db)
    return _memory_store