feat(similarity): 添加异步和批量相似度计算功能,优化性能

feat(graph_store): 增强图存储管理,添加边的注册和注销功能
feat(memory_tools): 支持批量生成嵌入向量
feat(unified_manager): 优化感知记忆和短期记忆的检索逻辑
This commit is contained in:
Windpicker-owo
2025-11-20 22:40:53 +08:00
parent 8dc754e562
commit ddc68b9257
9 changed files with 349 additions and 97 deletions

View File

@@ -21,7 +21,7 @@ import numpy as np
from src.common.logger import get_logger
from src.memory_graph.models import MemoryBlock, PerceptualMemory
from src.memory_graph.utils.embeddings import EmbeddingGenerator
from src.memory_graph.utils.similarity import cosine_similarity
from src.memory_graph.utils.similarity import cosine_similarity_async, batch_cosine_similarity_async
logger = get_logger(__name__)
@@ -430,14 +430,22 @@ class PerceptualMemoryManager:
logger.warning("查询向量生成失败,返回空列表")
return []
# 计算所有块的相似度
# 批量计算所有块的相似度(使用异步版本)
blocks_with_embeddings = [
block for block in self.perceptual_memory.blocks
if block.embedding is not None
]
if not blocks_with_embeddings:
return []
# 批量计算相似度
block_embeddings = [block.embedding for block in blocks_with_embeddings]
similarities = await batch_cosine_similarity_async(query_embedding, block_embeddings)
# 过滤和排序
scored_blocks = []
for block in self.perceptual_memory.blocks:
if block.embedding is None:
continue
similarity = cosine_similarity(query_embedding, block.embedding)
for block, similarity in zip(blocks_with_embeddings, similarities):
# 过滤低于阈值的块
if similarity >= similarity_threshold:
scored_blocks.append((block, similarity))