向量检索与 RAG
向量搜索与检索增强生成(RAG)
相关源文件
本章引用的主要源码文件:
cognee/infrastructure/llm/prompts/extract_query_time.txtcognee/modules/graph/utils/__init__.pycognee/modules/graph/utils/get_entity_nodes_from_triplets.pycognee/modules/retrieval/base_retriever.pycognee/modules/retrieval/chunks_retriever.pycognee/modules/retrieval/completion_retriever.pycognee/modules/retrieval/graph_completion_context_extension_retriever.pycognee/modules/retrieval/graph_completion_cot_retriever.pycognee/modules/retrieval/graph_completion_retriever.pycognee/modules/retrieval/graph_summary_completion_retriever.pycognee/modules/retrieval/jaccard_retrival.pycognee/modules/retrieval/lexical_retriever.pycognee/modules/retrieval/summaries_retriever.pycognee/modules/retrieval/temporal_retriever.pycognee/modules/retrieval/triplet_retriever.pycognee/modules/retrieval/utils/access_tracking.pycognee/modules/retrieval/utils/completion.pycognee/tasks/cleanup/cleanup_unused_data.pycognee/tests/unit/modules/retrieval/chunks_retriever_test.pycognee/tests/unit/modules/retrieval/graph_completion_retriever_context_extension_test.pycognee/tests/unit/modules/retrieval/graph_completion_retriever_cot_test.pycognee/tests/unit/modules/retrieval/graph_completion_retriever_test.pycognee/tests/unit/modules/retrieval/rag_completion_retriever_test.pycognee/tests/unit/modules/retrieval/summaries_retriever_test.pycognee/tests/unit/modules/retrieval/temporal_retriever_test.pycognee/tests/unit/modules/retrieval/test_brute_force_triplet_search.pycognee/tests/unit/modules/retrieval/triplet_retriever_test.py
本文档介绍了 Cognee 中的向量搜索和检索增强生成(RAG)能力。它解释了文本和数据如何被转换为向量嵌入向量、存储在向量数据库中,并通过语义相似度搜索进行检索。检索增强生成(RAG)通过使用检索到的上下文来增强大语言模型(LLM)的响应,从而扩展了向量搜索的功能。
关于基于图的搜索和遍历,请参见 4.2 基于图的搜索。关于可用搜索类型的详细信息,请参见 4.1 搜索类型概览。关于向量数据库提供商的配置,请参见 5.1 向量数据库适配器。
目的与架构
向量搜索通过将文本转换为高维嵌入向量并在向量空间中寻找最近邻,从而实现语义相似度检索。检索增强生成(RAG)将这种检索与大语言模型(LLM)生成相结合,以提供基于上下文的响应。
该架构遵循以下模式:
- 嵌入向量生成:通过
EmbeddingEngine将文本转换为向量cognee/modules/retrieval/completion_retriever.py:5-5。 - 存储:通过
index_data_points将向量与载荷数据一起存储在向量数据库中cognee/tasks/storage/index_data_points.py:10-27。 - 搜索:对查询进行嵌入向量化,并通过
vector_engine.search进行相似度搜索以找到相关结果cognee/modules/retrieval/chunks_retriever.py:98-100。 - 检索增强生成(RAG)补全:将检索到的上下文发送给大语言模型(LLM)以生成答案,使用特定的检索器,如
CompletionRetrievercognee/modules/retrieval/completion_retriever.py:17-148或ChunksRetrievercognee/modules/retrieval/chunks_retriever.py:11-108。
来源:cognee/modules/retrieval/completion_retriever.py:5-148, cognee/tasks/storage/index_data_points.py:10-27, cognee/modules/retrieval/chunks_retriever.py:11-108
向量数据库集成
向量数据库架构
向量数据库系统使用工厂模式,其中 get_vector_engine() 创建实现 VectorDBInterface 的特定于提供商的适配器。像 ChunksRetriever 这样的检索器与 unified.vector 引擎交互,以针对特定集合(如 DocumentChunk_text)执行搜索 cognee/modules/retrieval/chunks_retriever.py:94-100。
来源:cognee/modules/retrieval/chunks_retriever.py:94-100, cognee/modules/retrieval/completion_retriever.py:40-45, cognee/infrastructure/databases/unified.py:26-26
嵌入向量引擎与数据点存储
嵌入向量生成
检索器使用 EmbeddingEngine(可通过向量引擎访问)将查询转换为向量。例如,TemporalRetriever 会显式地对查询进行嵌入向量化,以搜索相关事件 cognee/modules/retrieval/temporal_retriever.py:147-152。
数据点模式与索引
向量数据库为每个数据点存储关键组件,包括 ID、向量和载荷。任务 index_data_points 负责编排创建向量索引和并行索引数据点的过程 cognee/tasks/storage/index_data_points.py:10-72。
存储模式实现
- LanceDB:使用带有
Vector类型的LanceModel,并处理模式迁移cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py:146-158。 - PGVector:利用 SQLAlchemy 的
Mapped列和pgvector类型cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py:158-171。 - ChromaDB:将复杂类型序列化为 JSON 字符串进行存储
cognee/infrastructure/databases/vector/chromadb/ChromaDBAdapter.py:56-90。
来源:cognee/tasks/storage/index_data_points.py:10-72, cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py:146-158, cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py:158-171, cognee/infrastructure/databases/vector/chromadb/ChromaDBAdapter.py:56-90
检索策略
Cognee 通过 BaseRetriever 接口实现了多种检索策略 cognee/modules/retrieval/base_retriever.py:9-9。
片段检索(CHUNKS)
ChunksRetriever 搜索 DocumentChunk_text 集合。它返回找到的片段的原始载荷,而不是生成新的大语言模型(LLM)补全 cognee/modules/retrieval/chunks_retriever.py:33-52。
async def get_retrieved_objects(self, query: str) -> Any:
unified = await get_unified_engine()
vector_engine = unified.vector
found_chunks = await vector_engine.search(
"DocumentChunk_text", query, limit=self.top_k, include_payload=True
)
return found_chunks
来源:cognee/modules/retrieval/chunks_retriever.py:11-108
检索增强生成(RAG)补全(RAG_COMPLETION)
CompletionRetriever 处理标准的检索增强生成(RAG)。它从 DocumentChunk_text 中检索片段,将它们组合成一个上下文字符串,并生成大语言模型(LLM)补全 cognee/modules/retrieval/completion_retriever.py:39-105。如果启用了缓存,它还支持基于会话的检索,包含对话历史 cognee/modules/retrieval/completion_retriever.py:132-147。
来源:cognee/modules/retrieval/completion_retriever.py:17-148
摘要检索(SUMMARIES)
SummariesRetriever 搜索 TextSummary_text 集合,以根据向量相似度找到相关的文档摘要。
来源:cognee/modules/retrieval/summaries_retriever.py:12-117
时间检索(时序)
TemporalRetriever 将时间戳过滤与向量搜索相结合。它使用 extract_time_from_query 从查询中提取时间间隔 cognee/modules/retrieval/temporal_retriever.py:84-104,根据这些时间戳过滤图事件,然后使用向量相似度从过滤后的集合中找到最相关的事件 cognee/modules/retrieval/temporal_retriever.py:147-154。
来源:cognee/modules/retrieval/temporal_retriever.py:19-171
访问跟踪与清理
Cognee 支持跟踪数据使用情况,以促进对未使用上下文的智能清理。
访问跟踪
工具函数 update_node_access_timestamps 在检索文档片段时更新关系数据库中的 last_accessed 字段 cognee/modules/retrieval/utils/access_tracking.py:20-59。它利用 CogneeGraph.project_graph_from_db 从检索到的节点回溯到其原始 TextDocument 或 Document 节点 cognee/modules/retrieval/utils/access_tracking.py:61-87。
来源:cognee/modules/retrieval/utils/access_tracking.py:20-102, cognee/modules/graph/cognee_graph/CogneeGraph.py:13-22
未使用数据清理
cleanup_unused_data 任务会识别在指定的 minutes_threshold(根据 Data 表中的 last_accessed 列计算)内未被访问的文档,并使用 cognee.delete(mode="hard") 将其删除 cognee/tasks/cleanup/cleanup_unused_data.py:27-172。
来源:cognee/tasks/cleanup/cleanup_unused_data.py:27-172
错误处理
| 异常 | 描述 | 来源 |
|---|---|---|
CollectionNotFoundError | 当请求的向量集合(例如 DocumentChunk_text)不存在时抛出。 | cognee/infrastructure/databases/vector/exceptions/exceptions.py:5-23 |
NoDataError | 当系统中没有数据可供搜索时,由检索器抛出。 | cognee/modules/retrieval/exceptions/exceptions.py:5-6 |
来源:cognee/modules/retrieval/chunks_retriever.py:105-107, cognee/modules/retrieval/completion_retriever.py:48-50