Cognee Custom Pipelines

Custom tasks and pipelines

Everything cognee does runs as a pipeline: an ordered list of tasks, each a plain Python function whose output feeds the next one. remember() is the right tool for ordinary ingestion. Build a pipeline when you need processing cognee does not ship: your own extraction, your own node types, or a post-processing step over the graph.

import cognee
from cognee.modules.pipelines import Task
from cognee.tasks.storage import add_data_points
from cognee.low_level import DataPoint


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


async def extract_people(data_items: list) -> list[Person]:
    people = []
    for item in data_items:  # always a list, see below
        text = item if isinstance(item, str) else ""
        people += [Person(name=n.strip()) for n in text.split(",") if n.strip()]
    return people


result = await cognee.run_custom_pipeline(
    tasks=[
        Task(extract_people, needs_llm=False),
        Task(add_data_points, needs_llm=False),  # store in graph + vector DBs
    ],
    data=["Ada Lovelace, Alan Turing"],
    dataset="people",
)

Use it

Pick the runner

There are three, and two share the name run_pipeline:

cognee.run_custom_pipeline(tasks, data=None, dataset="main_dataset", user=None, incremental_loading=False, data_per_batch=20, run_in_background=False, pipeline_name="custom_pipeline", data_cache=False, ...) returns {dataset_id: PipelineRunInfo} (the started run when run_in_background=True). With data=None it runs over the dataset's existing documents (Data rows).

Write a task

from cognee.modules.pipelines import Task
from cognee.modules.pipelines.models import PipelineContext
from cognee.modules.pipelines.tasks.task import task_summary
from cognee.pipelines import Drop


@task_summary("Tagged {n} chunk(s)")
async def tag_chunks(chunks: list, ctx: PipelineContext = None, label: str = "x"):
    for chunk in chunks:
        chunk.metadata["label"] = label
    return chunks  # or yield per item; return/yield Drop to discard


tag = Task(tag_chunks, label="reviewed", batch_size=10, needs_llm=False)

• A task is an async def, a generator, an async generator, or a plain def. Extra Task(fn, *args, **kwargs) arguments are passed after the pipeline data.
• needs_llm=False on tasks that never call an LLM lets an LLM-free pipeline skip the LLM connection check.
• ctx (injected by the parameter name ctx) carries user, data_item, dataset, pipeline_run_id, pipeline_name, and extras.
• task.with_config(batch_size=..., **kwargs) returns a modified copy.

How data flows

• Each document runs the whole chain on its own with run_custom_pipeline or the orchestrator, and the first task receives it as a one-element list ([data_item]), not the bare item. The lightweight run_pipeline passes data to the first task unchanged.
• data_per_batch (default 20) is how many documents run at the same time. It is a concurrency limit, not a batch size.
• batch_size belongs to the consumer. A task's batch_size decides how the previous task's generator output is grouped before it is passed in. Generator tasks always hand over lists; a coroutine or function hands over its single return value.
• Streaming: each upstream result goes down the chain immediately, so a downstream task can run many times per document.
• enriches=True: if the task returns None, its input is passed on unchanged (coroutines and functions only, not generators).
• Drop: returning or yielding it removes that item from the stream.
• Every DataPoint passing through is stamped automatically with where it came from (source_pipeline, source_task, source_user, …).

Store results

add_data_points(data_points, custom_edges=None, embed_triplets=False, graph_only=False) writes a list of DataPoints to the graph and indexes their index_fields in the vector DB. It returns the same list, so it can sit mid-chain. Give every node type identity_fields so repeated runs merge instead of duplicating (see the cognee-custom-graph-models skill).

Work on the existing graph: memify

await cognee.memify(
    extraction_tasks=["extract_subgraph_chunks"],  # names or Task objects
    enrichment_tasks=[Task(my_enrichment, needs_llm=False)],
    dataset="people",
    node_name=["AI"],  # optional subgraph filter
)

With no data, memify passes the graph (or the node_type / node_name subgraph) to the first task. Registered task names: extract_subgraph, extract_subgraph_chunks, get_triplet_datapoints, extract_user_sessions, cognify_session, extract_agent_trace_feedbacks, cognify_agent_trace_feedback, apply_feedback_weights, detect_entity_duplicates, merge_entity_duplicates, index_data_points. improve() also forwards extraction_tasks / enrichment_tasks to memify, but only inside its enrichment stage. With custom tasks that stage skips the TRIPLET_EMBEDDING gate and the has-the-graph-changed check, so they run on every improve (unless the stage is disabled, the lock is held, or the fatal persist_session_qa stage errors and stops the run first).

Check status

status = await cognee.datasets.get_status([dataset_id], pipeline_names=["custom_pipeline"])

Without pipeline_names it reports only cognify_pipeline. It returns {str(dataset_id): PipelineRunStatus} ({str(dataset_id): {pipeline_name: status}} for several pipeline_names): DATASET_PROCESSING_STARTED, _COMPLETED, or _ERRORED (_INITIATED exists only on legacy rows). The value run_custom_pipeline returns per dataset is a PipelineRunInfo instead, whose class names the outcome: PipelineRunCompleted, PipelineRunAlreadyCompleted, PipelineRunErrored, and so on.