MemoryAgents/backend/app/adapters/text2mem.py

from __future__ import annotations

import re
from typing import Any

from ..config import settings
from ..db import utcnow
from ..embeddings import embedding_client, embedding_fingerprint
from ..llm import LLMUnavailable
from ..schemas import ChatResponse, MemoryItem, SystemDescriptor
from .base import MemoryAgent


class Text2MemAgent(MemoryAgent):
    id = "text2mem"

    _EXPLICIT_WRITE_RE = re.compile(r"记住|记下|请记录|请保存")
    _DECLARATIVE_RE = re.compile(
        r"我(?:偏好|喜欢|习惯|要求|正在|在开发|决定|选择|使用)|"
        r"项目(?:目前|现在|已经|使用|采用|数据库|前端|后端)|"
        r"(?:前端|后端|数据库|对话模型|向量模型)(?:使用|采用|选择|是|改成)|"
        r"更新一下|以后|上次|之前|曾经|长期规则|代码要求"
    )
    _QUESTION_RE = re.compile(
        r"[?？]$|(?:什么|哪些|哪个|怎么|如何|是否|有没有|多少|为什么).*(?:[?？]|$)"
    )
    _STOPWORDS = frozenset(
        {
            "什么",
            "哪些",
            "哪个",
            "怎么",
            "如何",
            "是否",
            "有没有",
            "请问",
            "告诉我",
            "我的",
            "当前",
            "相关",
            "一下",
            "遇到",
            "遇到了",
            "的",
            "吗",
            "呢",
        }
    )

    def __init__(self) -> None:
        super().__init__()
        self.descriptor = SystemDescriptor(
            id="text2mem",
            name="Text2Mem",
            paradigm="IR 操作契约",
            description=(
                "当前实现 Encode / Retrieve、持久化、删除与审计；"
                "Update / Lock / Expire 等完整 IR 操作仍是后续项。"
            ),
            available=True,
            mode="native",
            status="ready",
            package="本项目内置教学实现",
        )

    @classmethod
    def _should_encode(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))

    @staticmethod
    def _memory_type(content: str) -> str:
        if re.search(r"上次|之前|曾经|遇到|发生|故障|失败", content):
            return "episodic"
        if re.search(r"流程|步骤|先.+再|以后.+按", content):
            return "procedural"
        if re.search(r"更新一下|进度|已完成|完成了|下一步|当前在做", content):
            return "task-state"
        return "semantic"

    @classmethod
    def _query_type(cls, message: str) -> str | None:
        if re.search(r"流程|步骤|怎么修复|如何修复|操作方法", message):
            return "procedural"
        if re.search(r"上次|之前|历史|曾经|遇到|发生|故障", message):
            return "episodic"
        if re.search(r"进度|完成|下一步|现在做到", message):
            return "task-state"
        if re.search(r"偏好|喜欢|技术栈|数据库|前端|后端|项目", message):
            return "semantic"
        return None

    @classmethod
    def _tokens(cls, text: str) -> set[str]:
        tokens = {item.casefold() for item in re.findall(r"[A-Za-z][A-Za-z0-9_.+-]{1,}", text)}
        for chunk in re.findall(r"[\u4e00-\u9fff]{2,}", text):
            if chunk not in cls._STOPWORDS:
                tokens.add(chunk)
            if len(chunk) > 2:
                tokens.update(
                    chunk[index : index + 2]
                    for index in range(len(chunk) - 1)
                    if chunk[index : index + 2] not in cls._STOPWORDS
                )
        return tokens

    async def _encode(self, content: str) -> tuple[MemoryItem, bool]:
        normalized = " ".join(content.strip().split()).casefold()
        for existing in await self.repo.list_memories(self.id):
            if " ".join(existing.content.strip().split()).casefold() == normalized:
                await self._audit(
                    "ENC/Encode",
                    existing.id,
                    status="deduplicated",
                    details={"source": "user_direct"},
                )
                return existing, False

        now = utcnow()
        item = MemoryItem(
            id=self.repo.memory_id("t2m"),
            system=self.id,
            workspace_id=settings.workspace_id,
            content=content.strip(),
            memory_type=self._memory_type(content),
            source="user_direct",
            confidence=1.0,
            valid_from=now,
            created_at=now,
            updated_at=now,
            metadata={
                "ir": {"stage": "ENC", "op": "Encode"},
                "policy": {"confirmation": True, "locked": False},
            },
        )
        await self.repo.add_memory(item)
        await self._audit(
            "ENC/Encode",
            item.id,
            details={"source": "user_direct", "memory_type": item.memory_type},
        )
        return item, True

    async def _sync_embeddings(self, items: list[MemoryItem]) -> dict[str, Any]:
        if not embedding_client.configured:
            return {"enabled": False, "embedded": 0, "removed": 0}

        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 pending:
            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 _retrieve(self, message: str, limit: int = 5) -> tuple[list[MemoryItem], dict[str, Any]]:
        items = await self.repo.list_memories(self.id)
        if not items:
            return [], {"strategy": "empty", "embedding": {"enabled": False}}

        query_tokens = self._tokens(message)
        expected_type = self._query_type(message)
        token_overlap: dict[str, float] = {}
        lexical_scores: dict[str, float] = {}
        for item in items:
            item_tokens = self._tokens(item.content)
            overlap = len(query_tokens & item_tokens) / max(1, len(query_tokens))
            token_overlap[item.id] = overlap
            type_bonus = 0.35 if expected_type and item.memory_type == expected_type else 0.0
            lexical_scores[item.id] = overlap + type_bonus

        semantic_scores: dict[str, float] = {}
        try:
            embedding_status = await self._sync_embeddings(items)
            if embedding_status["enabled"]:
                query_vector = (await embedding_client.embed([message]))[0]
                semantic_matches = await self.repo.search_embeddings(
                    self.id,
                    query_vector,
                    limit=min(max(limit * 2, 5), 20),
                )
                semantic_scores = {
                    str(match["memory_id"]): float(match["score"])
                    for match in semantic_matches
                }
        except Exception as exc:
            embedding_status = {"enabled": True, "error": str(exc), "embedded": 0, "removed": 0}

        ranked: list[tuple[float, MemoryItem]] = []
        for item in items:
            lexical = lexical_scores[item.id]
            semantic = semantic_scores.get(item.id, 0.0)
            if (
                expected_type
                and item.memory_type != expected_type
                and token_overlap[item.id] < 0.25
                and semantic < 0.65
            ):
                continue
            if lexical <= 0 and semantic < settings.embedding_min_score:
                continue
            ranked.append((lexical + max(semantic, 0.0), item))

        if not ranked and expected_type:
            ranked = [
                (0.1, item)
                for item in items
                if item.memory_type == expected_type
            ]

        ranked.sort(key=lambda pair: (pair[0], pair[1].updated_at), reverse=True)
        selected = [item for _, item in ranked[:limit]]
        return selected, {
            "strategy": "lexical+type+embedding" if semantic_scores else "lexical+type",
            "expected_type": expected_type,
            "embedding": embedding_status,
            "matches": [
                {
                    "memory_id": item.id,
                    "lexical": round(lexical_scores[item.id], 4),
                    "semantic": round(semantic_scores.get(item.id, 0.0), 4),
                }
                for item in selected
            ],
        }

    async def chat(self, message: str) -> ChatResponse:
        events: list[dict[str, Any]] = []
        if self._should_encode(message):
            item, created = await self._encode(message)
            events.append(
                {
                    "event": "ENC/Encode",
                    "memory_id": item.id,
                    "content": item.content,
                    "memory_type": item.memory_type,
                    "status": "created" if created else "deduplicated",
                }
            )

        context, retrieval = await self._retrieve(message)
        system_prompt = (
            "你是一个 AI 编程助手。Text2Mem 只允许通过显式 IR 记忆操作读写记忆。"
            "回答时区分当前用户陈述、历史记忆和推断，不要把推断写成事实。"
            "只能依据给定的记忆证据回答历史问题；证据不足时要明确说明。"
        )
        context_text = "\n".join(
            f"- [{item.memory_type} | {item.source} | {item.updated_at.isoformat()}] {item.content}"
            for item in context
        ) or "（暂无命中记忆）"
        try:
            answer = await self.llm.chat(system_prompt, f"记忆上下文：\n{context_text}\n\n用户：{message}")
            mode = "llm"
        except LLMUnavailable:
            mode = "local-fallback"
            answer = (
                "Text2Mem 本地模式：我已按显式 IR 规则处理本轮输入。\n\n"
                f"当前命中记忆：{context_text}\n\n"
                "配置 LLM_API_KEY 后，可让模型基于这些记忆生成自然语言回答。"
            )
        except Exception as exc:
            mode = "local-fallback"
            retrieval["llm_error"] = str(exc)[:500]
            answer = (
                "模型调用失败，Text2Mem 已退回本地证据模式；记忆写入和检索结果不受影响。\n\n"
                f"当前命中记忆：{context_text}"
            )
        await self._audit(
            "RET/Retrieve",
            details={
                "query": message,
                "count": len(context),
                "memory_ids": [item.id for item in context],
                **retrieval,
            },
        )
        return ChatResponse(
            system="text2mem",
            answer=answer,
            mode=mode,
            memory_context=context,
            memory_events=events,
            audit_events=await self.audit(),
        )