Giving AI Agents Pipeline Lineage Memory: How We Built an OpenLineage Connector for Cognee

Author: Soumyajit Ghosh ( @somuai ) Target Publication: Dev.to / Substack Hackathon: Mergetober (WeMakeDevs x Cognee) Pull Request: topoteretes/cognee-community#349 (Closes topoteretes/cognee#5554 )

Modern autonomous AI agents are rapidly evolving into operational engineering partners. In data platform engineering, agents are increasingly tasked with investigating broken ETL workflows, auditing dataset provenance, and answering critical questions:

"Which upstream Apache Spark job populates the retention_metrics table?" "Why did the finance reporting pipeline fail yesterday?" "If we modify the schema of the orders table, which downstream pipelines and dashboards will break?"

However, most LLM applications possess zero awareness of operational lineage. They cannot observe how data moves through Apache Airflow DAGs, Spark transformations, dbt models, or Flink streams.

Enter OpenLineage, the Linux Foundation (LF AI & Data) open standard for operational lineage metadata, and Marquez, its reference metadata collection and lineage visualization engine.

In this walkthrough, we examine how we built an OpenLineage and Marquez data-source connector for Cognee ( topoteretes/cognee#5554 ), enabling AI agents to ingest pipeline topologies, execution runs, and dataset schemas directly into cognitive memory graphs.

Cognee is the open-source memory engine for AI agents. Rather than dumping raw unstructured text into basic vector stores:

A connector in Cognee bridges upstream systems into this shared cognitive architecture.

When designing the OpenLineage connector under packages/connector/openlineage/, we established three architectural principles:

An enterprise OpenLineage backend receives millions of low-level JSON RunEvent objects containing transient transport data. Ingesting raw JSON into an LLM produces token bloat and degraded retrieval quality.

Instead, the connector queries Marquez REST API endpoints ( /api/v1/namespaces, /api/v1/jobs, /api/v1/runs, /api/v1/datasets ) and synthesizes high-density architectural Markdown prose:

Tabular data in dlt is typically treated as rigid relational tables. Lineage, however, is richest when modeled as structural documentation.

By setting the Cognee document-mode marker:

from cognee.tasks.ingestion.dlt_utils import DOCUMENT_SOURCE_ATTR source = _openlineage () setattr ( source, DOCUMENT_SOURCE_ATTR, " openlineage " ) Enter fullscreen mode Exit fullscreen mode Cognee's ingestion pipeline routes each job and its dependencies into the full cognify entity-extraction engine. The LLM extracts explicit graph nodes ( Job, Dataset, Schema, Run ) and relationships ( READS_FROM, WRITES_TO, HAS_RUN, PRODUCES ), storing them in graph memory.

In production data stacks, pipeline DAGs are continuously renamed, refactored, or deprecated. A major failure mode in AI memory is stale ghost entities: an agent hallucinating that a pipeline job exists weeks after it was decommissioned.

To resolve this, the connector enforces a full-snapshot replacement strategy:

All connector code was validated with end-to-end tests covering namespace discovery, facet parsing, secret redaction, and snapshot replacement:

Using the connector in an autonomous agent application requires only a few lines of code:

import asyncio import os import cognee from cognee_community_connector_openlineage import openlineage_source DATASET_NAME = " pipeline_lineage_memory " async def main (): # 1. Configure the OpenLineage source pointing to Marquez source = openlineage_source ( endpoint_url = os. environ. get ( " OPENLINEAGE_URL ", " http://localhost:5000 " ), api_key = os. environ. get ( " OPENLINEAGE_API_KEY " ), namespaces = [ " analytics ", " finance " ], include_facets = True, max_runs_per_job = 5, ) # 2. Sync into Cognee Memory print ( " Syncing OpenLineage pipeline topologies into Cognee...

" ) await cognee. remember ( source, dataset_name = DATASET_NAME ) # 3. Ask cross-pipeline structural questions query = ( " Which upstream jobs populate the retention_metrics dataset, " " and did any recent runs fail? " ) answer = await cognee. search ( query_text = query, query_type = cognee.

SearchType. GRAPH_COMPLETION, datasets = [ DATASET_NAME ], ) print ( " \n Agent Answer: \n ", answer ) if __name__ == " __main__ ": asyncio. run ( main ()) Enter fullscreen mode Exit fullscreen mode When the agent queries lineage memory, it traverses the knowledge graph to synthesize a precise root-cause analysis:

Notice that the agent did not need to sift through thousands of raw JSON log events. It traversed the graph directly from dataset to job to failed run, immediately surfacing the exact error facet ( OutOfMemoryError: Java heap space ).

To ensure frictionless maintainer review and rock-solid CI pipelines:

By bridging OpenLineage and Marquez into Cognee:

The code is available in fork branch somuai/cognee-community:feat/connector-openlineage under packages/connector/openlineage/, resolving topoteretes/cognee#5554.