cognee 怎么组合自定义 Task 与 DataPoint 构建自定义 cognify 管道？-编程实验室

cognee 怎么组合自定义 Task 与 DataPoint 构建自定义 cognify 管道？

【免费下载链接】cogneeCognee is the open-source AI memory platform for agents. Give your AI agents persistent long-term memory across sessions with a self-hosted knowledge graph engine.项目地址: https://gitcode.com/GitHub_Trending/co/cognee

当用 Cognee 默认的add -> cognify -> search流程构建知识图谱时，提取阶段使用的是内置的实体与关系抽取逻辑。如果你的领域需要特定的实体结构（例如“人物 + knows 关系”）、或者需要显式控制管道里每个阶段的执行顺序，就可以把 cognify 中的处理阶段替换成自己编写的Task，并通过DataPoint把结果写入图谱。本文以仓库中的示例 custom_tasks_and_pipelines.py 为主路径，说明如何定义 Task 与 DataPoint、组装任务列表、用cognee.run_custom_pipeline执行，以及用图谱可视化核对结果。

准备条件

Cognee 的 API 基本都是 async 的，示例代码都运行在asyncio事件循环中（见 cognee/skill.md）。
示例中的自定义任务通过LLMGateway.acreate_structured_output调用大模型做结构化抽取，因此需要先配置可用的 LLM。cognee/skill.md 给出的配置方式是：

cognee.config.set_llm_provider("openai") cognee.config.set_llm_model("gpt-4o-mini") cognee.config.set_llm_api_key("sk-...")

其中sk-...替换为你自己的 API Key。

cognee/skill.md 的建议是：只有当默认流程不满足需求、需要显式的顺序任务控制时才使用自定义管道。如果add -> cognify已经够用，不必引入本节的机制。

两个核心构件：Task 与 DataPoint

自定义管道由两类对象组合而成：

Task是管道的一个执行阶段。从 task.py 的实现看，Task的executable支持四种形式：异步协程函数、同步函数、生成器、异步生成器，传入其他类型会抛出ValueError。Task(fn, batch_size=..., **kwargs)中除函数与batch_size之外的关键字参数会在执行时透传给函数。文件还提供了task()装饰器：

from cognee.modules.pipelines.tasks.task import task @task(batch_size=20) async def extract(chunks, graph_model=None): ... # 调用 TaskSpec 返回 BoundTask，用于绑定参数 await run_pipeline([ extract(graph_model=KnowledgeGraph), extract(graph_model=KG, batch_size=5), # 覆盖 batch_size ], data=input_data)

DataPoint是 Cognee 中表示一个知识单元的 Pydantic 模型，从 DataPoint.py 继承。按 cognee/skill.md 的说明：

metadata = {"index_fields": [...]}控制哪些字段会被嵌入向量、用于语义搜索；
关系字段可以指向其他 DataPoint 实例，从而以编程方式定义图结构；
通过cognee.tasks.storage.add_data_points插入的 DataPoint 会成为图中的节点和边，同时贡献可搜索的向量字段。

也就是说，自定义 Task 负责“从输入数据中产出 DataPoint”，add_data_points这个内置 Task 负责“把 DataPoint 落库到图与向量库”。

主路径：定义抽取任务并运行自定义管道

完整可运行的示例见 examples/guides/custom_tasks_and_pipelines.py，目标是从一段人物关系文本中抽取Person节点与knows边。下面按执行顺序拆解。

1. 定义模型：抽取用模型与 DataPoint 分开

示例为 LLM 结构化输出定义了轻量 Pydantic 模型（PersonLLM/PeopleLLM，其中knows只是名字列表），同时定义了真正入库的DataPoint：

from pydantic import BaseModel from cognee.infrastructure.engine import DataPoint class PersonLLM(BaseModel): """Lightweight Pydantic model for LLM extraction only.""" name: str knows: List[str] = [] # Just names for now, we'll resolve to Person instances later class PeopleLLM(BaseModel): """Lightweight Pydantic model for LLM extraction only.""" persons: List[PersonLLM] class Person(DataPoint): name: str # Optional relationships (we'll let the LLM populate this) knows: List["Person"] = [] # Make names searchable in the vector store metadata: Dict[str, Any] = {"index_fields": ["name"]}

Person.knows指向Person自身，字段名会成为边标签；metadata中声明index_fields: ["name"]让name字段进入向量检索。

2. 定义输入数据与抽取任务

输入数据也用 DataPoint 表示（示例中的LightweightData，携带内容哈希用的稳定 UUID），并把待处理文本包进去：

class LightweightData(DataPoint): """Lightweight DataPoint model for data ingestion only.""" id: UUID text: str def build_lightweight_data_object(text_data): return LightweightData(id=uuid5(NAMESPACE_OID, text_data), text=text_data)

自定义抽取任务是一个 async 函数，接收上一个阶段（或管道）传入的数据，内部调用LLMGateway.acreate_structured_output生成结构化结果，再把名字解析成Person实例并建立knows关系：

async def extract_people(data: LightweightData) -> List[Person]: system_prompt = ( "Extract people mentioned in the text. " "Return as `persons: Person[]` with each Person having `name` and optional `knows` relations. " "Infer ‘knows’ only when there is a clear interpersonal interaction in the text." ) person_map: Dict[str, Person] = {} for data_item in data: people_llm = await LLMGateway.acreate_structured_output( data_item.text, system_prompt, PeopleLLM ) for person_llm in people_llm.persons: person_map[person_llm.name] = Person(name=person_llm.name) for person_llm in people_llm.persons: person = person_map[person_llm.name] person.knows = [person_map[name] for name in person_llm.knows if name in person_map] return list(person_map.values())

3. 组装 Task 列表并执行 run_custom_pipeline

任务按顺序执行：前一个 Task 的输出作为后一个 Task 的输入。示例里是“自定义抽取 + 内置落库”两步：

from cognee.modules.pipelines import Task from cognee.tasks.storage import add_data_points tasks = [ Task(extract_people), # input: text -> output: list[Person] Task(add_data_points), # input: list[Person] -> output: list[Person] ] await cognee.run_custom_pipeline( tasks=tasks, data=build_lightweight_data_object(text_data), dataset="people_demo" )

run_custom_pipeline的完整签名与参数说明在 run_custom_pipeline.py，对常用参数的文档说明如下：

tasks：要执行的Task列表；data：传入第一个抽取任务的输入，可以是任意形式（配合自定义任务时）。
dataset：数据集名称或 UUID，示例中用dataset="people_demo"隔离本次运行的数据。
run_in_background：True时异步启动并立即返回，适合大数据集（文档建议 >100MB），通过返回的pipeline_run_id监控进度。
use_pipeline_cache/incremental_loading/data_cache：用于跳过已处理数据，避免重复。
data_per_batch：并行处理的数据项数量。
skip_connection_test：True时跳过首次运行的 LLM/embedding 连接检查，适用于任务本身不做 LLM 或 embedding 调用的管道（文档示例：确定性的代码图管道）。
vector_db_config/graph_db_config：为本次管道指定自定义向量库/图库配置。

4. 完整示例的运行入口

示例文件还包含清场、初始化、后续 cognify 和可视化。注意cognee.forget(everything=True)会清空已有记忆数据，示例用它保证从零开始；custom_data_models.py 中同样的调用被注释为 “Start clean (optional in your app)”，在你的应用里是可选的。

async def main(text_data): await cognee.forget(everything=True) await setup() tasks = [ Task(extract_people), # input: text -> output: list[Person] Task(add_data_points), # input: list[Person] -> output: list[Person] ] await cognee.run_custom_pipeline( tasks=tasks, data=build_lightweight_data_object(text_data), dataset="people_demo" ) await cognee.cognify() visualize_graph_path = os.path.join( os.path.dirname(__file__), ".artifacts", "custom_tasks_and_pipelines.html" ) await visualize_graph(visualize_graph_path) if __name__ == "__main__": text = "Alice knows Mark. Mark had dinner with Bob and Alice. Bob knows Mary." asyncio.run(main(text))

示例中自定义管道之后还会再执行一次cognee.cognify()，随后生成图谱 HTML。运行方式：python examples/guides/custom_tasks_and_pipelines.py（LLM 配置完成后）。

结果验证

示例的验证方式是生成图谱可视化文件（.artifacts/custom_tasks_and_pipelines.html）并在浏览器中查看Person节点与knows边是否符合输入文本。visualize_graph默认渲染有界的子图（种子节点 + k-hop 邻域），而不是整图，上限见 cognee/skill.md：neighborhood_depth=2、neighborhood_seed_top_k=10、max_nodes=500；也可以用query、seed_node_ids参数指定观察范围，或用full=True渲染整图。

如果只想知道数据集是否存在、内容规模如何，可以用 cognee/skill.md 中给出的生命周期接口做轻量检查，例如await cognee.datasets.list_datasets()。

最小形式与替代路径

如果只是想验证管道机制，cognee/skill.md 给出的最小可运行形式是一个原样返回输入的 Task：

from cognee.modules.pipelines.tasks.task import Task async def my_task(data): return data await cognee.run_custom_pipeline( tasks=[Task(my_task)], data="input", dataset="research", )

另外，如果数据源本身已经是结构化 Python 对象（不需要从文本抽取），可以跳过自定义管道，直接用add_data_points插入，见 custom_data_models.py：

from cognee.tasks.storage import add_data_points class Person(DataPoint): name: str knows: SkipValidation[Any] = None # single Person or list[Person] metadata: dict = {"index_fields": ["name"]} alice = Person(name="Alice") alice.knows = bob # 字段名成为边标签 await add_data_points([alice, bob])

限制与边界

Task只接受协程、普通函数、生成器、异步生成器四种 executable，传入其他类型直接抛ValueError（见 task.py）。
incremental_loading与data_cache的去重只对使用 Cognee 内置Data模型的数据生效（run_custom_pipeline文档说明）。
自定义管道属于“非默认任务编排”的高级用法；cognee/skill.md 明确建议优先从add -> cognify -> search的最小路径开始，确有需要时再引入自定义 Task。