3.5 · 自定义 Graph 模型（Custom Graph Models）

自定义 Graph 模型

自定义图模型

Cognee 中的自定义图模型是 Pydantic 类（继承自 DataPoint），通过 graph_model 参数传递给 cognify() 或 remember() 函数。这些模型定义了模式，用于指导大语言模型（LLM）从文档中提取实体和关系。通过指定自定义模型，你可以精确控制系统提取的实体和关系类型，从而生成特定领域的知识图谱。

范围：本页面涵盖 graph_model 参数、默认的 KnowledgeGraph 模型、如何定义自定义 DataPoint 提取模型，以及将这些模型转换为图结构的内部逻辑。

概述

cognify() 中的 graph_model 参数接受任何定义了 LLM 提取结构的 Pydantic BaseModel（具体来说是 DataPoint）。在 extract_graph_from_data 任务中，Cognee 会将这些模型与结构化输出框架结合使用，确保 LLM 返回的数据符合你指定的模式。

来源：cognee/tasks/graph/extract_graph_from_data.py:129-136、cognee/modules/graph/utils/get_graph_from_model.py:178-184、examples/guides/custom_graph_model.py:14-33

Cognify 管线中的图模型流程

下图连接了自然语言空间（用户的提取意图）和代码实体空间（处理模型的任务和函数）。

图：提取模型数据流

Cognee · 图：提取模型数据流 · 图 1

来源：cognee/tasks/graph/extract_graph_from_data.py:56-125、cognee/tasks/storage/add_data_points.py:30-72、cognee/infrastructure/llm/extraction.py:20-20

默认的 KnowledgeGraph 模型

当未指定 graph_model 时，Cognee 会默认使用 KnowledgeGraph 模型。如果提供了自定义模型，系统会跳过标准的本体扩展逻辑，将提供的模型视为严格的提取模式。

# integrate_chunk_graphs 中的逻辑检查
if graph_model is not KnowledgeGraph:
    for chunk_index, chunk_graph in enumerate(chunk_graphs):
        data_chunks[chunk_index].contains = chunk_graph

    return data_chunks

来源：cognee/tasks/graph/extract_graph_from_data.py:99-103、cognee/shared/data_models.py:19-19

内部图提取逻辑

当 Cognee 存储数据点（无论是来自自定义模型还是默认模型）时，它会使用 get_graph_from_model 将 Pydantic 对象转换为由节点和边组成的图结构。这个过程会递归地将嵌套的 DataPoint 对象视为相关节点。

图：get_graph_from_model 逻辑

Cognee · 图：get_graph_from_model 逻辑 · 图 2

提取过程中的关键函数：

1. _extract_field_data：判断字段值是否包含 DataPoint、DataPoint 列表或 Tuple[Edge, DataPoint]。这样可以附加特定的边元数据。cognee/modules/graph/utils/get_graph_from_model.py:49-92
2. _create_edge_properties：构建图边的属性字典，包括 source_node_id、target_node_id 和 relationship_name。cognee/modules/graph/utils/get_graph_from_model.py:95-117
3. get_graph_from_model：递归遍历模型的主入口。它使用 added_nodes 和 visited_properties 来防止循环图中的无限循环。cognee/modules/graph/utils/get_graph_from_model.py:178-184
4. 记忆化：Cognee 使用 _simple_model_for 来缓存 Pydantic 类，以防止在大型提取周期中因重复调用 copy_model 而导致内存泄漏。cognee/modules/graph/utils/get_graph_from_model.py:28-47

来源：cognee/modules/graph/utils/get_graph_from_model.py:28-184

本体集成与扩展

对于默认的 KnowledgeGraph 模型，Cognee 会使用 ontology_resolver 执行本体验证和扩展。

• 本体解析器：通常是加载 OWL/RDF 文件的 RDFLibOntologyResolver。cognee/modules/ontology/rdf_xml/RDFLibOntologyResolver.py:14-20
• 扩展：expand_with_nodes_and_edges 函数会检查提取的实体在本体中是否存在父类或相关个体。cognee/modules/graph/utils/expand_with_nodes_and_edges.py:110-112
• 标准化：它使用 generate_node_id 和 generate_node_name 对节点名称和 ID 进行标准化，以确保一致性。cognee/modules/graph/utils/expand_with_nodes_and_edges.py:39-40

来源：cognee/tasks/graph/extract_graph_from_data.py:110-112、cognee/modules/graph/utils/expand_with_nodes_and_edges.py:31-65

存储与索引

从自定义模型中提取图后，会由 add_data_points 任务进行处理。

1. 去重：deduplicate_nodes_and_edges 确保如果多个片段提取到相同的实体，数据库中只会创建一个节点。cognee/tasks/storage/add_data_points.py:78-78
2. 混合写入：如果数据库支持，add_nodes_with_vectors 会一步完成节点及其嵌入向量的保存。cognee/tasks/storage/add_data_points.py:87-88
3. 向量索引：index_data_points 根据模型 metadata 字典中定义的 index_fields 为节点生成嵌入向量。cognee/tasks/storage/add_data_points.py:91-95、examples/guides/custom_graph_model.py:21-21
4. 三元组嵌入：如果启用了 embed_triplets，系统会根据图结构创建 Triplet 对象，用于关系感知搜索。cognee/tasks/storage/add_data_points.py:143-147

来源：cognee/tasks/storage/add_data_points.py:78-147、cognee/modules/engine/models.py:16-16

示例：自定义编程语言模型

自定义模型通过类型提示定义关系。在下面的示例中，used_in 成为一条边，因为它指向 Field 对象列表（这些对象是 DataPoint）。

from typing import List
from cognee.low_level import DataPoint

class Field(DataPoint):
    name: str
    metadata: dict = {"index_fields": ["name"]}

class ProgrammingLanguage(DataPoint):
    name: str
    used_in: List[Field] = None
    metadata: dict = {"index_fields": ["name"]}

# 使用方式
await cognee.remember(text, graph_model = ProgrammingLanguage)

1. extract_graph_from_data 使用 LLM 填充 ProgrammingLanguage 实例。cognee/tasks/graph/extract_graph_from_data.py:161-173
2. get_graph_from_model 将 used_in 识别为关系，因为 Field 继承自 DataPoint。cognee/modules/graph/utils/get_graph_from_model.py:49-53
3. add_data_points 将结果存储在图数据库和向量数据库中。cognee/tasks/storage/add_data_points.py:30-45

来源：examples/guides/custom_graph_model.py:14-33、cognee/tasks/graph/extract_graph_from_data.py:129-173