cortexm/pipeline.py
Rama main · réplica de lectura del repositorio original.
"""Composable ingestion pipelines (Cognee learn).
Reddit deep-dive 2026-08-29 — Cognee's "Pipelines" module
structured ingestion as composable stages:
chunk → embed → extract → graph → index
Cognee's insight: monolithic `mem.add(text)` does everything
in one pass; that's fine for chat (you want low latency) but
bad for batch (you want to control chunk size, embedder codec,
graph construction rules, etc.).
cortexm.Pipeline is the lean answer. It's a declarable list of
stages; each stage is a callable ``stage(documents, ctx) ->
documents``. The pipeline runs them in order, threading a
PipelineContext through so stages can share state.
Built-in stages:
- ``Chunk(max_tokens=512, overlap=64)`` — split long docs
- ``Extract(patterns="default")`` — μ=0 extractor
- ``Embed(codec="int8")`` — VSA hologram encode
- ``Index(backend="quadrant")`` — palace.add_batch
- ``Dedup(threshold=0.92)`` — value-match dedup
- ``Audit(stage="…")`` — log a stage marker
Usage::
from cortexm import Memory, Pipeline
from cortexm.pipeline import stages
m = Memory()
pipe = Pipeline([
stages.Chunk(max_tokens=512, overlap=64),
stages.Extract(),
stages.Embed(codec="int8"),
stages.Index(),
], memory=m)
pipe.run(["long doc 1 ...", "long doc 2 ..."])
Lean: ~250 LoC, no new deps. Stages are plain callables; users
can write their own (e.g. an LLM-enrichment stage for μ>0 paths).
"""
from __future__ import annotations
import datetime as _dt
import hashlib
import json
import re
import uuid
from dataclasses import dataclass, field
from typing import Any, Callable, Iterable
# ----------------------------- PipelineContext -----------------------------
@dataclass
class PipelineContext:
"""Threaded through every stage. Stages read/write state here.
Fields:
- memory: the live Memory object (stages call memory.add /
memory.writer.ingest_candidates / memory.palace.add_batch)
- user_id, agent_id, run_id: scope — every chunk the pipeline
produces inherits these
- stats: per-stage invocation count + last-run timing
- stage_state: scratch dict for inter-stage handoff (e.g.
Chunk() writes 'chunks', Extract() reads 'chunks' and writes
'candidates', Index() reads 'candidates')
"""
memory: Any = None
user_id: str = "default"
agent_id: str | None = None
run_id: str | None = None
stats: dict = field(default_factory=dict)
stage_state: dict = field(default_factory=dict)
def bump(self, stage_name: str, **kw) -> None:
s = self.stats.setdefault(stage_name, {"calls": 0, "last": {}})
s["calls"] += 1
s["last"] = kw
# ----------------------------- Stage contract ------------------------------
class Stage:
"""Base class for pipeline stages. Subclasses implement __call__."""
name: str = "stage"
def __call__(self, ctx: PipelineContext,
documents: Iterable[str]) -> list[str]:
raise NotImplementedError
# ----------------------------- Built-in stages -----------------------------
class Chunk(Stage):
"""Split long documents into ≤max_tokens chunks. Overlap so
facts that span a boundary don't get lost.
Tokenization: a lean regex word splitter — no transformer tokenizer
dependency (μ=0 stays intact)."""
name = "chunk"
def __init__(self, max_tokens: int = 512, overlap: int = 64):
self.max_tokens = max(max_tokens, 16)
self.overlap = max(0, min(overlap, self.max_tokens // 4))
def __call__(self, ctx, documents):
out: list[str] = []
for doc in documents:
toks = re.findall(r"\S+|\s+", doc) # word + whitespace
i = 0
while i < len(toks):
chunk_toks = toks[i:i + self.max_tokens]
chunk = "".join(chunk_toks).strip()
if chunk:
out.append(chunk)
if i + self.max_tokens >= len(toks):
break
i += self.max_tokens - self.overlap
ctx.stage_state["chunks"] = out
ctx.bump(self.name, n_chunks=len(out))
return out
class Extract(Stage):
"""Run the μ=0 deterministic extractor over each chunk. Emits
candidate (subject, relation, value) triples into
``ctx.stage_state['candidates']`` — does NOT commit them yet
(Index() does that so dedup/rbac can run first)."""
name = "extract"
def __init__(self, patterns: str = "default"):
self.patterns = patterns
def __call__(self, ctx, documents):
mem = ctx.memory
if mem is None:
ctx.stage_state["candidates"] = []
return list(documents)
extractor = mem.extractor
candidates = []
for chunk in documents:
try:
facts = extractor.extract(chunk)
except Exception:
facts = []
for f in facts:
# keep the Candidate objects as-is — Index() passes
# them straight to writer.ingest_candidates which
# expects Candidate dataclass instances, not dicts.
candidates.append(f)
ctx.stage_state["candidates"] = candidates
ctx.bump(self.name, n_candidates=len(candidates))
return list(documents)
class Embed(Stage):
"""Encode every candidate's subject/relation/value into VSA
holograms. Stores the vectors in ctx.stage_state['vectors']
keyed by candidate index; Index() consumes them."""
name = "embed"
def __init__(self, codec: str = "int8"):
self.codec = codec
def __call__(self, ctx, documents):
mem = ctx.memory
if mem is None:
return list(documents)
candidates = ctx.stage_state.get("candidates", [])
palace = mem.palace
vecs = []
for c in candidates:
try:
v = palace.vsa.encode_fact(
palace.embedder.embed(c["subject"]),
palace.embedder.embed(c["relation"]),
palace.embedder.embed(c["value"]))
vecs.append(v)
except Exception:
vecs.append(None)
ctx.stage_state["vectors"] = vecs
ctx.bump(self.name, n_vectors=len(vecs))
return list(documents)
class Index(Stage):
"""Commit candidates to the Trace + Palace via the standard
MemoryWriter.ingest_candidates path — same quarantine / lifecycle
/ palace / edges pipeline as pattern-extracted facts. Provenance
records ``pipeline`` origin so audits can distinguish pipeline-
ingested facts from chat-ingested ones."""
name = "index"
def __init__(self, backend: str | None = None):
# backend is informational only — palace already chose its
# backend via Config.index_backend. We expose it on the stage
# so the pipeline declaration documents the active backend.
self.backend = backend
def __call__(self, ctx, documents):
mem = ctx.memory
if mem is None:
return list(documents)
candidates = ctx.stage_state.get("candidates", [])
if not candidates:
ctx.bump(self.name, committed=0)
return list(documents)
try:
mem.writer.ingest_candidates(
candidates, user_id=ctx.user_id, agent_id=ctx.agent_id,
source="pipeline")
except Exception:
# fall back to per-doc mem.add
for doc in documents:
try:
mem.add(doc, user_id=ctx.user_id,
agent_id=ctx.agent_id, run_id=ctx.run_id)
except Exception:
pass
mem.reader.invalidate_caches()
ctx.bump(self.name, committed=len(candidates))
return list(documents)
class Dedup(Stage):
"""Value-match dedup: drop candidates whose (subject, relation)
pair already has a near-identical value (similarity ≥ threshold)
earlier in the candidate list. Cuts redundant writes."""
name = "dedup"
def __init__(self, threshold: float = 0.92):
self.threshold = float(threshold)
def __call__(self, ctx, documents):
candidates = ctx.stage_state.get("candidates", [])
from cortexm.util import similarity
kept = []
seen: dict[tuple, list[str]] = {}
for c in candidates:
subj = getattr(c, "subject", "") or ""
rel = getattr(c, "relation", "") or ""
val = getattr(c, "value", "") or ""
key = (subj.lower(), rel.lower())
dup = False
for prev_val in seen.get(key, []):
if similarity(prev_val, val) >= self.threshold:
dup = True
break
if not dup:
kept.append(c)
seen.setdefault(key, []).append(val)
ctx.stage_state["candidates"] = kept
ctx.bump(self.name, kept=len(kept),
dropped=len(candidates) - len(kept))
return list(documents)
class Audit(Stage):
"""Insert a marker into the audit log so the trajectory viewer
can render pipeline progress. Lean: just calls memory.audit_log.log."""
name = "audit"
def __init__(self, stage: str = "pipeline"):
self.stage = stage
def __call__(self, ctx, documents):
mem = ctx.memory
if mem is not None:
try:
mem.audit_log.log(f"pipeline.{self.stage}",
resource=ctx.user_id,
meta={"n_docs": len(list(documents))})
except Exception:
pass
ctx.bump(self.name)
return list(documents)
# ----------------------------- Pipeline -----------------------------------
class Pipeline:
"""A declarable ingestion pipeline.
pipe = Pipeline([stages.Chunk(...), stages.Extract(),
stages.Embed(), stages.Index()],
memory=m, user_id="alice")
pipe.run(["doc 1 ...", "doc 2 ..."])
Stages run in order, each receiving the output of the previous
stage (default: pass-through — most stages communicate via
ctx.stage_state, not the return value).
"""
def __init__(self, stages: list[Stage], *, memory=None,
user_id: str = "default",
agent_id: str | None = None,
run_id: str | None = None):
self.stages = stages
self.ctx = PipelineContext(memory=memory, user_id=user_id,
agent_id=agent_id, run_id=run_id)
def run(self, documents: Iterable[str]) -> dict:
docs = list(documents)
for stage in self.stages:
try:
docs = list(stage(self.ctx, docs) or docs)
except Exception as e:
self.ctx.bump(stage.name, error=str(e)[:200])
# continue — one bad stage shouldn't kill the pipeline
return {
"n_in": len(docs),
"stats": self.ctx.stats,
"stage_state_keys": list(self.ctx.stage_state.keys()),
}
# Convenience namespace
stages = type("stages", (), {
"Chunk": Chunk,
"Extract": Extract,
"Embed": Embed,
"Index": Index,
"Dedup": Dedup,
"Audit": Audit,
})