"""记忆服务。 集成 ChromaDB 向量存储： - create_memory / 大/小周天自动创建的 Memory 写入后异步生成 embedding 存向量 - retrieve_memories 支持向量相似度检索，失败时回退 SQL 排序检索 """ from __future__ import annotations from sqlalchemy.ext.asyncio import AsyncSession from app.core.config import settings from app.core.errors import AppError from app.core.logging import logger from app.core.time import utc_now from app.db.session import AsyncSessionLocal from app.models.memory import Memory from app.modules.memory.context import MemoryRetrievalResult from app.repositories import memory as memory_repo from app.repositories.json_utils import dumps from app.schemas.common import Page from app.schemas.memory import MemoryCreateRequest, MemoryResponse, MemoryUpdateRequest from app.services.audit import record_audit_event from app.services.characters import get_active_character_model from app.services.serializers import memory_to_response def _get_embedding_service(manager): """尝试创建 EmbeddingService，失败时返回 None。""" try: models = manager.model_registry.list_by_capability("embedding") if not models: return None from app.ai_provider.embedding import EmbeddingService return EmbeddingService(manager) except Exception: # noqa: BLE001 return None async def _embed_and_store(memory: Memory) -> None: """异步后台任务：为记忆生成 embedding 并存入向量库。 embedding 服务和 ChromaDB 均不可用时静默跳过——不影响主流程。 """ text = memory.content or memory.title or "" if not text: return try: from app.ai_provider import ProviderManagerFactory manager = ProviderManagerFactory.create() service = _get_embedding_service(manager) if service is None: return import asyncio as _asyncio response = await _asyncio.wait_for( service.embed(text), timeout=10.0, ) if not response.data or not response.data[0].embedding: return embedding = response.data[0].embedding from app.integrations.vector_store import get_vector_store store = get_vector_store() if store is None: return store.add_memory( character_id=memory.character_id, memory_id=memory.id, text=text, embedding=embedding, metadata={ "type": memory.type, "source": memory.source, "title": memory.title or "", }, ) except Exception: # noqa: BLE001 logger.warning( "memory_embed_failed", memory_id=memory.id, character_id=memory.character_id, ) async def embed_memory_async(memory: Memory) -> None: """公开接口：为单条记忆生成 embedding 并存入向量库（供大/小周天后台任务调用）。""" await _embed_and_store(memory) # ── CRUD ────────────────────────────────────────────────────────────────── async def list_memory_page( *, page: int = 1, page_size: int = 20, memory_type: str | None = None, source: str | None = None, owner_scope: str = settings.default_owner_scope, ) -> Page[MemoryResponse]: async with AsyncSessionLocal() as session: character = await get_active_character_model(session, owner_scope) safe_page_size = min(max(page_size, 1), 100) safe_page = max(page, 1) items, total = await memory_repo.list_memories( session, character.id, memory_type=memory_type, source=source, page=safe_page, page_size=safe_page_size, ) return Page( items=[memory_to_response(item) for item in items], page=safe_page, page_size=safe_page_size, total=total, ) async def create_memory( session: AsyncSession, payload: MemoryCreateRequest, *, owner_scope: str = settings.default_owner_scope, ) -> MemoryResponse: character = await get_active_character_model(session, owner_scope) character_id = payload.character_id or character.id if character_id != character.id: raise AppError("not_found", "人物不存在。") memory = Memory( character_id=character_id, type=payload.type, title=payload.title, content=payload.content, source=payload.source, confidence=payload.confidence, importance=payload.importance, sensitivity=payload.sensitivity, evidence_ids_json=dumps(payload.evidence_ids), status="active", ) session.add(memory) await session.flush() # 异步 embedding if settings.vector_db_enabled: import asyncio asyncio.ensure_future(_embed_and_store(memory)) await record_audit_event( session, action="memory.created", character_id=character_id, actor="user", decision="auto", input_summary=payload.title, resource_type="memory", resource_id=memory.id, ) await session.commit() await session.refresh(memory) return memory_to_response(memory) async def update_memory( memory_id: str, payload: MemoryUpdateRequest, *, owner_scope: str = settings.default_owner_scope ) -> MemoryResponse: async with AsyncSessionLocal() as session: character = await get_active_character_model(session, owner_scope) memory = await memory_repo.get_memory(session, memory_id) if memory is None or memory.character_id != character.id or memory.deleted_at is not None: raise AppError("not_found", "记忆不存在。") updates = payload.model_dump(exclude_unset=True) evidence_ids = updates.pop("evidence_ids", None) for key, value in updates.items(): setattr(memory, key, value) if evidence_ids is not None: memory.evidence_ids_json = dumps(evidence_ids) memory.updated_at = utc_now() # 如果内容变了，异步更新向量 if ("content" in updates or "title" in updates) and settings.vector_db_enabled: import asyncio asyncio.ensure_future(_embed_and_store(memory)) await record_audit_event( session, action="memory.updated", character_id=character.id, actor="user", decision="auto", input_summary=memory.title, resource_type="memory", resource_id=memory.id, ) await session.commit() await session.refresh(memory) return memory_to_response(memory) async def delete_memory( memory_id: str, *, owner_scope: str = settings.default_owner_scope ) -> MemoryResponse: async with AsyncSessionLocal() as session: character = await get_active_character_model(session, owner_scope) memory = await memory_repo.get_memory(session, memory_id) if memory is None or memory.character_id != character.id or memory.deleted_at is not None: raise AppError("not_found", "记忆不存在。") memory.status = "deleted" memory.deleted_at = utc_now() memory.updated_at = utc_now() # 同步删除向量 if settings.vector_db_enabled: try: from app.integrations.vector_store import get_vector_store store = get_vector_store() if store is not None: store.delete_memory(character.id, memory_id) except Exception: # noqa: BLE001 pass await record_audit_event( session, action="memory.deleted", character_id=character.id, actor="user", decision="auto", input_summary=memory.title, resource_type="memory", resource_id=memory.id, ) await session.commit() await session.refresh(memory) return memory_to_response(memory) async def memory_context( session: AsyncSession, character_id: str, *, limit: int, _owner_scope: str = settings.default_owner_scope, ) -> str: """获取记忆上下文文本（供叩心枢开场、辰正等场景使用）。""" memories = await memory_repo.list_injectable_memories(session, character_id, limit=limit) if not memories: return "" lines = [f"- {item.title}: {item.content}" for item in memories] return "相关记忆：\n" + "\n".join(lines) async def retrieve_memories( session: AsyncSession, character_id: str, *, limit: int, _owner_scope: str = settings.default_owner_scope, ) -> MemoryRetrievalResult: """检索可注入上下文的相关记忆。 优先尝试向量相似度检索，失败时回退 SQL 排序检索。 """ # 尝试向量检索（5s 超时，避免 embedding endpoint 不可达时阻塞消息流） if settings.vector_db_enabled: try: from app.ai_provider import ProviderManagerFactory manager = ProviderManagerFactory.create() models = manager.model_registry.list_by_capability("embedding") if models: import asyncio as _asyncio from app.ai_provider.embedding import EmbeddingService from app.integrations.vector_store import get_vector_store query_text = f"检索 {character_id} 的最新相关记忆" service = EmbeddingService(manager) response = await _asyncio.wait_for( service.embed(query_text), timeout=5.0, ) if response.data and response.data[0].embedding: store = get_vector_store() if store is not None: vector_results = store.search( character_id, response.data[0].embedding, limit=limit, ) if vector_results: ids = [r["id"] for r in vector_results] memories = await memory_repo.get_memories_by_ids(session, ids) if memories: return MemoryRetrievalResult(memories=memories) except Exception: # noqa: BLE001 logger.warning("vector_retrieval_failed，回退 SQL 检索") # 回退 SQL memories = await memory_repo.list_injectable_memories(session, character_id, limit=limit) return MemoryRetrievalResult(memories=memories)