ai-agent/MemoryAgents/backend/app/adapters/reme.py

from __future__ import annotations

import hashlib
import re
from datetime import datetime, timezone
from typing import Any

from ..embeddings import embedding_client, embedding_fingerprint
from ..config import settings
from ..schemas import ChatResponse, MemoryItem, SystemDescriptor
from .base import AdapterUnavailable, MemoryAgent


class ReMeAgent(MemoryAgent):
    """用 Markdown 文件保存记忆，并提供文件与向量搜索。"""

    id = "reme"
    _PASSIVE_JOB_OVERRIDES = {
        # 关闭后台任务，改为每次对话显式更新索引，避免隐式改文件和模型调用。
        "index_update_loop": {"backend": "base", "enable_serve": False},
        "resource_watch_loop": {"backend": "base", "enable_serve": False},
        "digest_watch_loop": {"backend": "base", "enable_serve": False},
        "dream_cron": {"backend": "base", "enable_serve": False},
    }

    # 扩展常见工程词，补强 ReMe 的 BM25 召回；语义召回单独使用 pgvector。
    _QUERY_EXPANSIONS = {
        "部署": ("deploy", "发布", "上线", "测试环境", "索引构建", "批次"),
        "上线": ("deploy", "发布", "测试环境", "索引构建"),
        "问题": ("超时", "错误", "失败", "异常", "报错", "故障", "冲突"),
        "故障": ("超时", "错误", "失败", "异常", "报错", "冲突"),
        "错误": ("失败", "异常", "报错", "故障"),
        "超时": ("timeout", "索引", "构建", "分批", "批次"),
        "连接": ("connection", "connect", "API", "数据库", "服务"),
        "索引": ("index", "向量", "构建", "分批", "批次"),
        "检索": ("retrieval", "RAG", "向量", "关键词", "重排"),
        "知识库": ("RAG", "文档", "向量", "检索", "重排"),
        "上次": ("之前", "最近", "历史"),
    }
    _QUERY_STOPWORDS = frozenset(
        {
            "什么",
            "哪些",
            "哪个",
            "怎么",
            "如何",
            "是否",
            "有没有",
            "请问",
            "告诉我",
            "根据",
            "当前",
            "相关",
            "一下",
            "遇到",
            "遇到了",
            "发生",
            "发生了",
            "的",
            "吗",
            "呢",
            "？",
            "?",
        }
    )
    _EXPLICIT_WRITE_RE = re.compile(r"记住|记下|请记录|请保存")
    _DECLARATIVE_RE = re.compile(
        r"我(?:偏好|喜欢|习惯|要求|正在|在开发|决定|选择|使用)|"
        r"项目(?:目前|现在|已经|使用|采用|数据库|前端|后端)|"
        r"(?:前端|后端|数据库|对话模型|向量模型)(?:使用|采用|选择|是|改成)|"
        r"更新一下|以后|上次|之前|曾经|长期规则|代码要求"
    )
    _QUESTION_RE = re.compile(
        r"[?？]$|(?:什么|哪些|哪个|怎么|如何|是否|有没有|多少|为什么).*(?:[?？]|$)"
    )

    def __init__(self) -> None:
        """加载 ReMe 依赖并记录当前是否可用。"""
        super().__init__()
        self._service: Any | None = None
        try:
            from reme.reme import ReMe  # type: ignore
            from reme.config import resolve_app_config  # type: ignore

            self._class = ReMe
            self._resolve_config = resolve_app_config
        except Exception:
            self._class = None
            self._resolve_config = None
        ready = bool(self._class and settings.reme_enabled)
        self.descriptor = SystemDescriptor(
            id="reme",
            name="ReMe",
            paradigm="文件即记忆",
            description="把长期记忆组织为可读 Markdown 文件，并融合 ReMe 检索与 pgvector 召回。",
            available=ready,
            mode="real-sdk" if ready else "unavailable",
            status="ready" if ready else "not-configured",
            package="reme-ai[core]>=0.4,<0.5",
            setup_hint=None if ready else "安装 reme-ai[core]，并设置 REME_ENABLED=true。",
        )

    async def memories(self, query: str | None = None) -> list[MemoryItem]:
        """把 ReMe 的 Markdown 文件转换为统一记忆模型。"""
        root = settings.data_dir / "reme"
        if not root.exists():
            return []
        items: list[MemoryItem] = []
        # 对外接口只认识 MemoryItem，因此需要把文件逐个转换。
        for path in sorted(root.glob("daily/**/*.md")):
            content = path.read_text(encoding="utf-8", errors="replace").strip()
            if not content:
                continue
            if query and query.lower() not in f"{path} {content}".lower():
                continue
            updated = datetime.fromtimestamp(path.stat().st_mtime, tz=timezone.utc)
            digest = hashlib.sha1(str(path).encode("utf-8")).hexdigest()[:16]
            items.append(MemoryItem(
                id=f"reme_{digest}",
                system=self.id,
                workspace_id=settings.workspace_id,
                content=content,
                memory_type="file-memory",
                source=str(path.relative_to(root)),
                confidence=1.0,
                created_at=updated,
                updated_at=updated,
                metadata={"path": str(path.relative_to(root)), "retrieval": "hybrid"},
            ))
        return items

    async def delete_memory(self, memory_id: str) -> bool:
        """按记忆 ID 删除文件，并同步清理索引和操作记录。"""
        root = settings.data_dir / "reme"
        deleted = False
        for path in root.glob("daily/**/*.md"):
            digest = hashlib.sha1(str(path).encode("utf-8")).hexdigest()[:16]
            if f"reme_{digest}" != memory_id:
                continue
            path.unlink(missing_ok=True)
            deleted = True
            break
        # 文件删掉后，对应向量和文件索引也要一起更新。
        if deleted:
            await self.repo.delete_embedding(self.id, memory_id)
        if deleted and self._class and settings.reme_enabled:
            service = await self._ensure()
            await service.run_job("reindex")
        if deleted:
            await self._audit("DELETE/FileMemory", target=memory_id)
        return deleted

    async def _ensure(self) -> Any:
        """按需启动 ReMe 服务，同一进程内重复使用。"""
        if not self._class or not settings.reme_enabled:
            raise AdapterUnavailable(self.descriptor.setup_hint or "ReMe 不可用")
        if self._service is None:
            working_dir = str(settings.data_dir / "reme")
            config = self._resolve_config(
                log_config=False,
                workspace_dir=working_dir,
                enable_logo=False,
                log_to_console=False,
                service={"backend": "http"},
                jobs=self._PASSIVE_JOB_OVERRIDES,
            )
            self._service = self._class(**config)
            await self._service.start()
        return self._service

    @classmethod
    def _should_memorize(cls, message: str) -> bool:
        """判断当前输入是否值得写入长期文件。"""
        normalized = " ".join(message.strip().split())
        if not normalized:
            return False
        if cls._EXPLICIT_WRITE_RE.search(normalized):
            return True
        if cls._QUESTION_RE.search(normalized):
            return False
        return bool(cls._DECLARATIVE_RE.search(normalized))

    @classmethod
    def _heuristic_queries(cls, message: str) -> list[str]:
        """为文件检索生成短查询和工程词扩展。"""
        normalized = message.strip()
        queries: list[str] = [normalized] if normalized else []

        # 保留 RAG、FastAPI、PostgreSQL 等可直接匹配的领域词。
        queries.extend(re.findall(r"[A-Za-z][A-Za-z0-9_.:/-]{1,}", normalized))

        for trigger, expansions in cls._QUERY_EXPANSIONS.items():
            # 用户用词较少时，补上文件里可能出现的近义说法。
            if trigger in normalized:
                queries.extend((trigger, *expansions))

        # 中文短词作为兜底，过滤会干扰词面检索的疑问词。
        for chunk in re.findall(r"[\u4e00-\u9fff]{2,}", normalized):
            if chunk not in cls._QUERY_STOPWORDS:
                queries.append(chunk)
                if len(chunk) > 2:
                    queries.extend(
                        chunk[index : index + 2]
                        for index in range(len(chunk) - 1)
                        if chunk[index : index + 2] not in cls._QUERY_STOPWORDS
                    )

        return cls._dedupe_queries(queries)

    @staticmethod
    def _dedupe_queries(queries: list[str]) -> list[str]:
        """清理重复或空查询，并限制单次搜索数量。"""
        unique: list[str] = []
        seen: set[str] = set()
        for query in queries:
            cleaned = " ".join(str(query).strip().split())
            if not cleaned:
                continue
            key = cleaned.casefold()
            if key in seen:
                continue
            seen.add(key)
            unique.append(cleaned)
        return unique[:12]

    async def _rewrite_query(self, message: str) -> list[str]:
        """合并规则扩展和可选的模型查询改写。"""
        queries = self._heuristic_queries(message)
        try:
            rewritten = await self.llm.json(
                (
                    "你是记忆检索查询改写器。只输出 JSON，不要回答用户问题。"
                    "把用户问题改写成 2 到 5 个适合混合文件检索的短查询，"
                    "保留专有名词，并补充可能出现在记忆里的中英文同义词。"
                    '格式必须是 {"queries": ["..."]}。'
                ),
                f"用户问题：{message}",
            )
            generated = rewritten.get("queries", [])
            if isinstance(generated, list):
                queries.extend(str(item) for item in generated if str(item).strip())
        except Exception:
            # 改写失败时继续使用规则查询，不影响基础检索。
            pass
        return self._dedupe_queries(queries)

    @staticmethod
    def _merge_search_results(results: list[str]) -> str:
        """合并多个查询结果并按正文去重。"""
        merged: list[str] = []
        seen: set[str] = set()
        for result in results:
            for block in re.split(r"(?=^========== )", result.strip(), flags=re.MULTILINE):
                block = block.strip()
                if not block:
                    continue
                fingerprint_source = block
                if block.startswith("==========") and "\n" in block:
                    fingerprint_source = block.split("\n", 1)[1]
                fingerprint = re.sub(r"\[(?:vector_)?score=[^\]]+\]", "", fingerprint_source)
                fingerprint = " ".join(fingerprint.split()).casefold()
                if fingerprint in seen:
                    continue
                seen.add(fingerprint)
                merged.append(block)
        return "\n\n".join(merged)

    async def _search_memory(self, service: Any, queries: list[str]) -> tuple[str, list[str]]:
        """逐个查询 ReMe 文件，并返回有结果的查询词。"""
        results: list[str] = []
        used_queries: list[str] = []
        for query in queries:
            search_result = await service.run_job("search", query=query, limit=5)
            answer = str(getattr(search_result, "answer", "") or "").strip()
            if answer:
                results.append(answer)
                used_queries.append(query)
        return self._merge_search_results(results), used_queries

    async def _sync_embeddings(self) -> dict[str, Any]:
        """更新已变化文件的向量，并清理已删除文件。"""
        if not embedding_client.configured:
            return {"enabled": False, "embedded": 0, "removed": 0}

        items = await self.memories()
        current_ids = {item.id for item in items}
        previous_hashes = await self.repo.embedding_hashes(self.id)
        # 文件不存在后，旧向量也要删除。
        stale_ids = set(previous_hashes) - current_ids
        for memory_id in stale_ids:
            await self.repo.delete_embedding(self.id, memory_id)

        # 只处理新增或内容有变化的文件。
        pending = [
            item
            for item in items
            if previous_hashes.get(item.id) != embedding_fingerprint(item.content)
        ]
        if not pending:
            return {"enabled": True, "embedded": 0, "removed": len(stale_ids)}

        vectors = await embedding_client.embed([item.content for item in pending])
        for item, vector in zip(pending, vectors):
            await self.repo.upsert_embedding(
                system=self.id,
                memory_id=item.id,
                content=item.content,
                source=item.source,
                content_hash=embedding_fingerprint(item.content),
                embedding=vector,
            )
        return {"enabled": True, "embedded": len(pending), "removed": len(stale_ids)}

    async def _semantic_search(self, message: str) -> tuple[str, int]:
        """查找内容相近的文件，并过滤分数过低的结果。"""
        if not embedding_client.configured:
            return "", 0
        try:
            query_vector = (await embedding_client.embed([message]))[0]
            matches = await self.repo.search_embeddings(self.id, query_vector, limit=5)
        except Exception:
            # 这一路失败时，主流程仍可使用 ReMe 文件搜索。
            return "", 0
        matches = [
            match
            for match in matches
            if float(match.get("score", 0.0)) >= settings.embedding_min_score
        ]
        blocks = [
            (
                f"========== {match['source']} [vector_score={float(match['score']):.4f}] =========="
                f"\n{match['content']}"
            )
            for match in matches
        ]
        return self._merge_search_results(blocks), len(matches)

    async def _answer_from_memory(self, message: str, evidence: str) -> str:
        """根据搜索结果生成回答，模型不可用时直接返回原文。"""
        if not evidence:
            return "没有检索到与这个问题相关的历史记忆。"
        try:
            return await self.llm.chat(
                (
                    "你是一个使用文件记忆的 AI 助手。"
                    "只能依据提供的记忆证据回答，不要编造证据中没有的事实。"
                    "如果证据不足，要明确说明不确定。直接回答用户问题，简洁自然。"
                ),
                f"用户问题：{message}\n\n记忆证据：\n{evidence}",
            )
        except Exception:
            # 保留原始内容比吞掉已经找到的结果更方便排查。
            return f"根据检索到的记忆：\n\n{evidence}"

    async def chat(self, message: str) -> ChatResponse:
        """完成文件写入、索引更新、搜索、回答和操作记录。"""
        service = await self._ensure()
        write_candidate = self._should_memorize(message)
        memory_result: Any | None = None
        # 问题和临时请求不写文件，只处理值得长期保存的陈述。
        if write_candidate:
            memory_result = await service.run_job(
                "auto_memory",
                messages=[{"name": "user", "role": "user", "content": message}],
                session_id=settings.workspace_id,
                memory_hint=(
                    "只记录稳定偏好、项目事实、进度变化、可复用的故障经验和长期流程规则。"
                    "问题、寒暄和临时请求不应写入长期记忆。"
                ),
            )
        # 写入后立即更新索引，保证本轮搜索能看到刚保存的内容。
        reindex_result = await service.run_job("reindex")
        try:
            embedding_sync = await self._sync_embeddings()
        except Exception as exc:
            # 向量服务异常时仍保留 ReMe 自带的文件检索。
            embedding_sync = {"enabled": True, "embedded": 0, "removed": 0, "error": str(exc)}
        queries = await self._rewrite_query(message)
        # 文件搜索和相近内容搜索各跑一次，最后合并并去重。
        lexical_search, used_queries = await self._search_memory(service, queries)
        semantic_search, semantic_count = await self._semantic_search(message)
        search = self._merge_search_results([lexical_search, semantic_search])
        answer = await self._answer_from_memory(message, search)
        # 记录查询词和搜索摘要，出现误匹配时可回看整个过程。
        await self._audit(
            "REME/auto_memory+search",
            details={
                "write_candidate": write_candidate,
                "memory": str(getattr(memory_result, "answer", "skipped"))[:800],
                "reindex": str(reindex_result.metadata)[:800],
                "embedding_sync": embedding_sync,
                "queries": queries,
                "used_queries": used_queries,
                "semantic_results": semantic_count,
                "search": str(search)[:1000],
            },
        )
        return ChatResponse(
            system="reme",
            answer=answer,
            mode="real-sdk",
            memory_context=await self.memories(),
            memory_events=[
                {
                    "event": "file-memory",
                    "write_candidate": write_candidate,
                    "queries": queries,
                    "retrieval": "hybrid",
                    "semantic_results": semantic_count,
                    "result": str(search)[:2000],
                }
            ],
            audit_events=await self.audit(),
        )

    async def reset(self) -> None:
        """清空公共数据，关闭服务并删除 ReMe 工作目录。"""
        await super().reset()
        # 先关闭服务，避免删除目录后仍有后台句柄占用文件。
        if self._service is not None:
            await self._service.close()
            self._service = None
        working_dir = settings.data_dir / "reme"
        if working_dir.exists():
            import shutil

            shutil.rmtree(working_dir)