""" 记忆服务 管理Agent的记忆存储和检索 """ import uuid from datetime import datetime, timedelta from typing import List, Dict, Any, Optional import numpy as np from loguru import logger from models.agent_memory import AgentMemory, MemoryType class MemoryService: """ Agent记忆服务 提供记忆的存储、检索和管理功能 """ # 嵌入模型（延迟加载） _embedding_model = None @classmethod def _get_embedding_model(cls): """ 获取嵌入模型实例（延迟加载） """ if cls._embedding_model is None: try: from sentence_transformers import SentenceTransformer cls._embedding_model = SentenceTransformer('paraphrase-multilingual-MiniLM-L12-v2') logger.info("嵌入模型加载成功") except Exception as e: logger.warning(f"嵌入模型加载失败: {e}") return None return cls._embedding_model @classmethod async def create_memory( cls, agent_id: str, content: str, memory_type: str = MemoryType.SHORT_TERM.value, importance: float = 0.5, source_room_id: Optional[str] = None, source_discussion_id: Optional[str] = None, tags: Optional[List[str]] = None, expires_in_hours: Optional[int] = None ) -> AgentMemory: """ 创建新的记忆 Args: agent_id: Agent ID content: 记忆内容 memory_type: 记忆类型 importance: 重要性评分 source_room_id: 来源聊天室 source_discussion_id: 来源讨论 tags: 标签 expires_in_hours: 过期时间（小时） Returns: 创建的AgentMemory文档 """ memory_id = f"mem-{uuid.uuid4().hex[:12]}" # 生成向量嵌入 embedding = await cls._generate_embedding(content) # 生成摘要 summary = content[:100] + "..." if len(content) > 100 else content # 计算过期时间 expires_at = None if expires_in_hours: expires_at = datetime.utcnow() + timedelta(hours=expires_in_hours) memory = AgentMemory( memory_id=memory_id, agent_id=agent_id, memory_type=memory_type, content=content, summary=summary, embedding=embedding, importance=importance, source_room_id=source_room_id, source_discussion_id=source_discussion_id, tags=tags or [], created_at=datetime.utcnow(), last_accessed=datetime.utcnow(), expires_at=expires_at ) await memory.insert() logger.debug(f"创建记忆: {memory_id} for Agent {agent_id}") return memory @classmethod async def get_memory(cls, memory_id: str) -> Optional[AgentMemory]: """ 获取指定记忆 Args: memory_id: 记忆ID Returns: AgentMemory文档或None """ return await AgentMemory.find_one(AgentMemory.memory_id == memory_id) @classmethod async def get_agent_memories( cls, agent_id: str, memory_type: Optional[str] = None, limit: int = 50 ) -> List[AgentMemory]: """ 获取Agent的记忆列表 Args: agent_id: Agent ID memory_type: 记忆类型（可选） limit: 返回数量限制 Returns: 记忆列表 """ query = {"agent_id": agent_id} if memory_type: query["memory_type"] = memory_type return await AgentMemory.find(query).sort( "-importance", "-last_accessed" ).limit(limit).to_list() @classmethod async def search_memories( cls, agent_id: str, query: str, limit: int = 10, memory_type: Optional[str] = None, min_relevance: float = 0.3 ) -> List[Dict[str, Any]]: """ 搜索相关记忆 Args: agent_id: Agent ID query: 查询文本 limit: 返回数量 memory_type: 记忆类型（可选） min_relevance: 最小相关性阈值 Returns: 带相关性分数的记忆列表 """ # 生成查询向量 query_embedding = await cls._generate_embedding(query) if not query_embedding: # 无法生成向量时，使用文本匹配 return await cls._text_search(agent_id, query, limit, memory_type) # 获取Agent的所有记忆 filter_query = {"agent_id": agent_id} if memory_type: filter_query["memory_type"] = memory_type memories = await AgentMemory.find(filter_query).to_list() # 计算相似度 results = [] for memory in memories: if memory.is_expired(): continue if memory.embedding: similarity = cls._cosine_similarity(query_embedding, memory.embedding) relevance = memory.calculate_relevance_score(similarity) if relevance >= min_relevance: results.append({ "memory": memory, "similarity": similarity, "relevance": relevance }) # 按相关性排序 results.sort(key=lambda x: x["relevance"], reverse=True) # 更新访问记录 for item in results[:limit]: memory = item["memory"] memory.access() await memory.save() return results[:limit] @classmethod async def update_memory( cls, memory_id: str, **kwargs ) -> Optional[AgentMemory]: """ 更新记忆 Args: memory_id: 记忆ID **kwargs: 要更新的字段 Returns: 更新后的AgentMemory或None """ memory = await cls.get_memory(memory_id) if not memory: return None # 如果更新了内容，重新生成嵌入 if "content" in kwargs: kwargs["embedding"] = await cls._generate_embedding(kwargs["content"]) kwargs["summary"] = kwargs["content"][:100] + "..." if len(kwargs["content"]) > 100 else kwargs["content"] for key, value in kwargs.items(): if hasattr(memory, key): setattr(memory, key, value) await memory.save() return memory @classmethod async def delete_memory(cls, memory_id: str) -> bool: """ 删除记忆 Args: memory_id: 记忆ID Returns: 是否删除成功 """ memory = await cls.get_memory(memory_id) if not memory: return False await memory.delete() return True @classmethod async def delete_agent_memories( cls, agent_id: str, memory_type: Optional[str] = None ) -> int: """ 删除Agent的记忆 Args: agent_id: Agent ID memory_type: 记忆类型（可选） Returns: 删除的数量 """ query = {"agent_id": agent_id} if memory_type: query["memory_type"] = memory_type result = await AgentMemory.find(query).delete() return result.deleted_count if result else 0 @classmethod async def cleanup_expired_memories(cls) -> int: """ 清理过期的记忆 Returns: 清理的数量 """ now = datetime.utcnow() result = await AgentMemory.find( {"expires_at": {"$lt": now}} ).delete() count = result.deleted_count if result else 0 if count > 0: logger.info(f"清理了 {count} 条过期记忆") return count @classmethod async def consolidate_memories( cls, agent_id: str, min_importance: float = 0.7, max_age_days: int = 30 ) -> None: """ 整合记忆（将重要的短期记忆转为长期记忆） Args: agent_id: Agent ID min_importance: 最小重要性阈值 max_age_days: 最大年龄（天） """ cutoff_date = datetime.utcnow() - timedelta(days=max_age_days) # 查找符合条件的短期记忆 memories = await AgentMemory.find({ "agent_id": agent_id, "memory_type": MemoryType.SHORT_TERM.value, "importance": {"$gte": min_importance}, "created_at": {"$lt": cutoff_date} }).to_list() for memory in memories: memory.memory_type = MemoryType.LONG_TERM.value memory.expires_at = None # 长期记忆不过期 await memory.save() if memories: logger.info(f"整合了 {len(memories)} 条记忆为长期记忆: Agent {agent_id}") @classmethod async def _generate_embedding(cls, text: str) -> List[float]: """ 生成文本的向量嵌入 Args: text: 文本内容 Returns: 向量嵌入列表 """ model = cls._get_embedding_model() if model is None: return [] try: embedding = model.encode(text, convert_to_numpy=True) return embedding.tolist() except Exception as e: logger.warning(f"生成嵌入失败: {e}") return [] @classmethod def _cosine_similarity(cls, vec1: List[float], vec2: List[float]) -> float: """ 计算余弦相似度 Args: vec1: 向量1 vec2: 向量2 Returns: 相似度 (0-1) """ if not vec1 or not vec2: return 0.0 try: a = np.array(vec1) b = np.array(vec2) similarity = np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)) return float(max(0, similarity)) except Exception: return 0.0 @classmethod async def _text_search( cls, agent_id: str, query: str, limit: int, memory_type: Optional[str] ) -> List[Dict[str, Any]]: """ 文本搜索（后备方案） Args: agent_id: Agent ID query: 查询文本 limit: 返回数量 memory_type: 记忆类型 Returns: 记忆列表 """ filter_query = {"agent_id": agent_id} if memory_type: filter_query["memory_type"] = memory_type # 简单的文本匹配 memories = await AgentMemory.find(filter_query).to_list() results = [] query_lower = query.lower() for memory in memories: if memory.is_expired(): continue content_lower = memory.content.lower() if query_lower in content_lower: # 计算简单的匹配分数 score = len(query_lower) / len(content_lower) results.append({ "memory": memory, "similarity": score, "relevance": score * memory.importance }) results.sort(key=lambda x: x["relevance"], reverse=True) return results[:limit]