memory.py

"""forgetfulMem 模块：SQLite 记忆存取。

记忆按 key 隔离（默认为 system 提示词的哈希，不同机器人设定各自独立）。
- 回复前：retrieve() 提取相关记忆注入 system 提示词
- 回复后（或不回复时）：store() 把本轮对话写入记忆
"""

import sqlite3
import threading
from datetime import datetime
from typing import List, Optional

_SCHEMA = """
CREATE TABLE IF NOT EXISTS memories (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    key TEXT NOT NULL,
    content TEXT NOT NULL,
    keywords TEXT DEFAULT '',
    created_at TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_memories_key ON memories(key);
"""


class MemoryStore:
    def __init__(self, db_path: str = "memories.db"):
        self.db_path = db_path
        self._lock = threading.Lock()
        # 持久连接：跨操作保留（:memory: 也适用）；所有访问都在 _lock 内，线程安全
        self._conn = sqlite3.connect(db_path, timeout=30, check_same_thread=False)
        self._init_db()

    def _init_db(self) -> None:
        with self._lock:
            self._conn.executescript(_SCHEMA)
            # 旧库迁移：补 keywords 列
            cols = [r[1] for r in self._conn.execute("PRAGMA table_info(memories)")]
            if "keywords" not in cols:
                self._conn.execute(
                    "ALTER TABLE memories ADD COLUMN keywords TEXT DEFAULT ''")
            self._conn.commit()

    # ---------- 写入 ----------
    def store(self, key: str, contents: List[str],
              keywords: Optional[List[str]] = None) -> int:
        """写入记忆条目（同 key + 同内容去重）。返回实际插入条数。

        keywords: 模型总结的关键词列表，一并入库供后续检索（逗号分隔存储）。
        """
        if not key or not contents:
            return 0
        now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        kw_text = ",".join(k.strip() for k in (keywords or []) if k.strip())
        inserted = 0
        with self._lock:
            cur = self._conn.cursor()
            for c in contents:
                c = (c or "").strip()
                if not c:
                    continue
                exists = cur.execute(
                    "SELECT COUNT(*) FROM memories WHERE key=? AND content=?",
                    (key, c)).fetchone()[0]
                if exists:
                    continue
                cur.execute(
                    "INSERT INTO memories (key, content, keywords, created_at) "
                    "VALUES (?, ?, ?, ?)",
                    (key, c, kw_text, now))
                inserted += 1
            self._conn.commit()
        return inserted

    # ---------- 读取 ----------
    def retrieve(self, key: str, limit: int = 5,
                 keywords: Optional[List[str]] = None) -> List[str]:
        """提取记忆。

        - 提供 keywords：按命中打分（记忆自带的关键词列命中权重更高），
          不足再补最近记录；
        - 未提供：直接取最近的 limit 条。
        """
        if not key:
            return []
        with self._lock:
            cur = self._conn.cursor()
            if keywords:
                kw = [k for k in keywords if len(k.strip()) >= 2]
                if kw:
                    rows = cur.execute(
                        "SELECT content, keywords FROM memories WHERE key=? "
                        "ORDER BY id DESC LIMIT ?", (key, limit * 5)).fetchall()
                    scored = []
                    for content, stored_kw in rows:
                        score = 0
                        for k in kw:
                            if k in (stored_kw or ""):
                                score += 2       # 模型总结的关键词命中：权重高
                            elif k in content:
                                score += 1       # 仅正文命中
                        if score:
                            scored.append((score, content))
                    scored.sort(key=lambda x: -x[0])
                    hits = [c for _, c in scored[:limit]]
                    if len(hits) < limit:
                        fill = cur.execute(
                            "SELECT content FROM memories WHERE key=? "
                            "ORDER BY id DESC LIMIT ?",
                            (key, limit - len(hits))).fetchall()
                        hits += [r[0] for r in fill if r[0] not in hits]
                    return hits[:limit]
            rows = cur.execute(
                "SELECT content FROM memories WHERE key=? "
                "ORDER BY id DESC LIMIT ?", (key, limit)).fetchall()
            return [r[0] for r in rows]

    def clear(self, key: str) -> int:
        """清空某个 key 的记忆，返回删除条数。"""
        with self._lock:
            cur = self._conn.cursor()
            cur.execute("DELETE FROM memories WHERE key=?", (key,))
            self._conn.commit()
            return cur.rowcount

    def count(self, key: str) -> int:
        with self._lock:
            return self._conn.execute(
                "SELECT COUNT(*) FROM memories WHERE key=?",
                (key,)).fetchone()[0]


if __name__ == "__main__":
    store = MemoryStore(":memory:")
    store.store("demo", ["第一条记忆", "第二条记忆", "第一条记忆"])
    print("count:", store.count("demo"))
    print("retrieve:", store.retrieve("demo", limit=5))
    print("keyword '二':", store.retrieve("demo", keywords=["二"]))