高级检索模式
高级搜索模式
相关源文件
本章引用的主要源码文件:
cognee/infrastructure/llm/prompts/extract_query_time.txtcognee/modules/retrieval/base_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/temporal_retriever.pycognee/modules/retrieval/utils/completion.pycognee/tests/integration/retrieval/test_chunks_retriever.pycognee/tests/integration/retrieval/test_graph_completion_retriever.pycognee/tests/integration/retrieval/test_graph_completion_retriever_context_extension.pycognee/tests/integration/retrieval/test_graph_completion_retriever_cot.pycognee/tests/integration/retrieval/test_structured_output.pycognee/tests/integration/retrieval/test_summaries_retriever.pycognee/tests/integration/retrieval/test_temporal_retriever.pycognee/tests/integration/retrieval/test_triplet_retriever.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 中高级搜索模式的实用指南,涵盖多阶段检索策略、专用搜索类型、会话管理以及自定义检索器实现。关于基本搜索用法,请参阅快速入门指南。关于各搜索类型的详细说明,请参阅搜索类型概览。
搜索类型选择策略
Cognee 提供 10 多种针对不同检索场景优化的搜索类型。选择取决于知识表示方式、查询特征以及期望的输出格式。
搜索类型决策矩阵
| 使用场景 | 搜索类型 | 检索器类 | 主要数据源 |
|---|---|---|---|
| 基于知识图谱的通用问答 | GRAPH_COMPLETION | GraphCompletionRetriever | 图谱边(三元组) |
| 迭代推理/优化 | GRAPH_COMPLETION_COT | GraphCompletionCotRetriever | 图谱边 |
| 深度上下文构建 | GRAPH_COMPLETION_CONTEXT_EXTENSION | GraphCompletionContextExtensionRetriever | 图谱边 |
| 基于时间的事件查询 | TEMPORAL | TemporalRetriever | 时序图谱节点 |
| 文档问答 | RAG_COMPLETION | CompletionRetriever | 向量嵌入 |
| 原始上下文检索 | CHUNKS | ChunksRetriever | 文档片段 |
| 高层概览 | SUMMARIES | SummariesRetriever | 文档摘要 |
| 代码导航 | CODE | CodeRetriever | 代码图谱 + 文件 |
| 直接图谱查询 | CYPHER | CypherSearchRetriever | 图数据库 |
来源:cognee/modules/retrieval/graph_completion_retriever.py:33-40、cognee/modules/retrieval/temporal_retriever.py:19-33、cognee/modules/retrieval/completion_retriever.py:17-20
检索器分发架构
检索系统通过分发与 UnifiedDatabaseEngine 交互的特定 BaseRetriever 实现,将自然语言查询桥接到代码实体。
来源:cognee/modules/retrieval/graph_completion_retriever.py:118-119、cognee/modules/retrieval/temporal_retriever.py:122-123、cognee/modules/retrieval/completion_retriever.py:40-40、cognee/modules/retrieval/base_retriever.py:5-13
基于图谱的搜索模式
基于图谱的检索器利用知识图谱结构来查找语义相关的实体和关系。所有高级检索器都继承自 GraphCompletionRetriever,并遵循 BaseRetriever 中定义的三步管线:get_retrieved_objects、get_context_from_objects 和 get_completion_from_context。
基础图谱补全
标准图谱检索模式对图谱节点执行向量相似度搜索以查找相关三元组,然后使用这些三元组作为上下文生成大语言模型(LLM)补全。
实现模式
GraphCompletionRetriever.get_triplets() 方法通过 _get_vector_index_collections() 从 DataPoint 子类中识别相关向量集合,从而执行核心检索。
# 实现代码位于 cognee/modules/retrieval/graph_completion_retriever.py
async def get_triplets(self, query: Optional[str] = None, query_batch: Optional[List[str]] = None):
collections = self._get_vector_index_collections()
# ...
return await brute_force_triplet_search(
query,
query_batch,
top_k = self.top_k,
collections = collections,
# ... 其他搜索参数
)
来源:cognee/modules/retrieval/graph_completion_retriever.py:154-175、cognee/modules/retrieval/graph_completion_retriever.py:88-99、cognee/modules/retrieval/graph_completion_retriever.py:101-136
思维链(CoT)检索
GraphCompletionCotRetriever 通过多轮生成推理、验证和后续问题来实现迭代优化。该过程由 _run_cot_completion 编排。
实现细节
思维链检索器使用专用提示词进行验证和后续问题生成:
| 阶段 | 提示词变量 | 用途 |
|---|---|---|
| 验证 | validation_user_prompt_path | 提取关于答案质量的推理 |
| 后续问题 | followup_user_prompt_path | 基于推理生成下一个查询 |
# 编排循环位于 cognee/modules/retrieval/graph_completion_cot_retriever.py
for reasoning_iteration in range(self.max_iter):
followup_queries = await self._run_cot_round(states)
await self._merge_followup_triplets(states, followup_queries)
if not (skip_final_completion and reasoning_iteration == self.max_iter - 1):
await self._generate_completions(states, conversation_history)
来源:cognee/modules/retrieval/graph_completion_cot_retriever.py:53-95、cognee/modules/retrieval/graph_completion_cot_retriever.py:170-175、cognee/modules/retrieval/graph_completion_cot_retriever.py:144-176
渐进式上下文扩展
GraphCompletionContextExtensionRetriever 通过将前一轮大语言模型(LLM)生成的补全结果作为新的搜索查询,逐步构建上下文,以扩展知识图谱邻域。
关键特性
- 增量增长:每一轮使用 LLM 的补全结果作为新的搜索查询来获取更多
Edge对象。 - 收敛检查:当使用
QueryState.check_convergence()未找到新三元组时停止。 - 去重:
merge_triplets确保状态中只维护唯一的边。
# 轮次逻辑位于 cognee/modules/retrieval/graph_completion_context_extension_retriever.py
async def _run_extension_round(self, states: dict):
# ...
# 使用当前补全结果作为新的搜索查询
completions = await generate_completion_batch(...)
new_triplets_batch = await self.get_triplets_batch(list(completions))
for q, new_triplets in zip(active_queries, new_triplets_batch):
states[q].merge_triplets(new_triplets)
states[q].check_convergence(prev_size)
来源:cognee/modules/retrieval/graph_completion_context_extension_retriever.py:102-129、cognee/modules/retrieval/graph_completion_context_extension_retriever.py:56-75
时序搜索
TemporalRetriever 专注于基于时间的查询,通过提取时间约束并据此过滤图谱事件。
时间提取与过滤
检索器使用 QueryInterval 模型来解析时间引用:
# 来自 cognee/modules/retrieval/temporal_retriever.py
async def extract_time_from_query(self, query: str):
# ...
system_prompt = render_prompt(prompt_path, {"time_now": time_now}, base_directory=base_directory)
interval = await LLMGateway.acreate_structured_output(query, system_prompt, QueryInterval)
return interval.starts_at, interval.ends_at
如果识别到时间戳,它会通过 collect_time_ids 和 collect_events 过滤事件。如果未找到时间戳,则通过 get_triplets(query) 回退到标准三元组搜索。
来源:cognee/modules/retrieval/temporal_retriever.py:84-104、cognee/modules/retrieval/temporal_retriever.py:119-154、cognee/modules/retrieval/temporal_retriever.py:106-117
基于向量的搜索模式
向量检索器使用嵌入文档片段上的语义相似度搜索。
检索增强生成(RAG)补全模式
CompletionRetriever 通过搜索 DocumentChunk_text 集合并将结果传递给大语言模型(LLM)来实现检索增强生成。
# 来自 cognee/modules/retrieval/completion_retriever.py
async def get_retrieved_objects(self, query: str) -> Any:
vector_engine = get_vector_engine()
try:
found_chunks = await vector_engine.search(
"DocumentChunk_text", query, limit=self.top_k, include_payload=True
)
return found_chunks
# ...
来源:cognee/modules/retrieval/completion_retriever.py:39-50、cognee/modules/retrieval/completion_retriever.py:58-82
多轮对话模式
Cognee 使用 get_session_manager() 和 CacheConfig 支持会话感知检索。
会话架构
检索器通过 session_user 上下文变量检查活动会话和用户 ID,以实现对话历史注入。
# 来自 cognee/modules/retrieval/graph_completion_retriever.py
def _use_session_cache(self) -> bool:
user = session_user.get()
user_id = getattr(user, "id", None)
return bool(user_id and CacheConfig().caching)
# 来自 cognee/modules/retrieval/completion_retriever.py
if use_session:
sm = get_session_manager()
# ...
completion = await sm.generate_completion_with_session(
session_id = self.session_id,
query = query,
context = context,
# ...
)
来源:cognee/modules/retrieval/graph_completion_retriever.py:81-85、cognee/modules/retrieval/completion_retriever.py:127-147、cognee/modules/retrieval/utils/completion.py:67-141
结构化输出与自定义模型
所有检索器都支持 response_model 参数,该参数利用 LLMGateway.acreate_structured_output 返回 Pydantic 模型而非原始字符串。
# 示例来自 cognee/modules/retrieval/utils/completion.py
async def generate_completion(
query: str,
context: str,
# ...
response_model: Type = str,
) -> Any:
# ...
result = await LLMGateway.acreate_structured_output(
text_input = user_prompt,
system_prompt = system_prompt,
response_model = response_model,
)
return result
来源:cognee/modules/retrieval/utils/completion.py:9-38、cognee/modules/retrieval/graph_completion_retriever.py:75-75
高级检索工具函数
| 函数 | 描述 | 来源 |
|---|---|---|
brute_force_triplet_search | 通过对节点进行向量搜索来查找图谱三元组的核心算法。 | cognee/modules/retrieval/utils/brute_force_triplet_search.py |
resolve_edges_to_text | 将 Edge 对象格式化为人类可读的文本,用于大语言模型(LLM)上下文。 | cognee/modules/graph/utils/__init__.py |
generate_completion_batch | 对多个查询并行执行 generate_completion。 | cognee/modules/retrieval/utils/completion.py:41-64 |
summarize_text | 使用 LLM 摘要功能压缩长上下文字符串。 | cognee/modules/retrieval/utils/completion.py:162-180 |
来源:cognee/modules/retrieval/graph_completion_retriever.py:10-23、cognee/modules/retrieval/utils/completion.py:1-181