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(), )