Writing a Custom Stage

This guide shows how to write your own pipeline stage and slot it into a DirectoryPipeline. It covers the Stage protocol, the rules every stage must follow, and how to manage the stage list.

A stage is a single, replaceable unit of work in the pipeline. The six built-ins are walk → parse → chunk → relate → enrich → embed-pack. When they don’t cover what you need, write your own. Common cases:

• A redaction pass before enrichment.
• A custom dedup step.
• A domain-specific relation resolver.

The Stage protocol

A stage is any object with a name and a run method. There’s no base class to inherit from — indx uses structural typing, so anything that fits the shape drops straight in.

from typing import Protocol, runtime_checkable
from indx import SpaceContext

@runtime_checkable
class Stage(Protocol):
    name: str  # stable identifier, e.g. "pii-redact"

    def run(self, ctx: SpaceContext) -> SpaceContext: ...

Two members, both mandatory:

Member	Type	Purpose
name	str	Stable identifier, used to address the stage in replace() / drop() and in progress output. Use a lowercase, hyphenated value (e.g. pii-redact).
run	(ctx: SpaceContext) -> SpaceContext	Does the work and returns the context.

The pipeline executes stages in registration order, threading the context out of each stage straight into the next. So whatever earlier stages produced is already on ctx when your stage runs:

• The directory graph from Walk.
• ctx.parsed from Parse.
• ctx.space.chunks from Chunk.
• ctx.space.relations from Relate.
• Enriched metadata from Enrich.
• Per-chunk embeddings on each chunk.embedding, plus ctx.space.manifest.embedding_model / ctx.space.manifest.embedding_dim, from Embed+Pack.

Return the SAME context object

A stage must return the exact SpaceContext instance it received, mutated in place. Never return a new object, a copy, or None.

The whole pipeline relies on one shared, append-only context threading through every stage. Returning a fresh object — or forgetting the return entirely — breaks that contract and silently drops everything downstream stages depend on.

def run(self, ctx: SpaceContext) -> SpaceContext:
    ...
    return ctx          # ✅ same object, mutated
    # return SpaceContext(...)   # ❌ new object — breaks the pipeline
    # (no return)                # ❌ returns None

A worked example: PII redaction

A redaction pass is the canonical custom stage. Run it after Chunk (so chunk text exists) and before any egress-capable component can see the raw content, such as an external enrichment LLM.

from indx import SpaceContext

class PiiRedactStage:
    name = "pii-redact"

    def run(self, ctx: SpaceContext) -> SpaceContext:
        for chunk in ctx.space.chunks:
            chunk.text = redact(chunk.text)
        return ctx            # MUST return the same context

Insert it at index 3. It lands after chunk (index 2) and before relate, which shifts from index 3 to 4:

pipeline.insert(3, PiiRedactStage())   # run after Chunk, before Relate

The default stage order is 0:walk, 1:parse, 2:chunk, 3:relate, 4:enrich, 5:embed-pack. Inserting at index 3 places your stage between Chunk and Relate, so by the time the external LLM in Enrich runs, every chunk’s text is already redacted.

Local profile

Redaction matters specifically when you use an egress-capable LLM, VLM, or hosted store. The default stack is cloud-backed: it uses openai models and requires OPENAI_API_KEY.

If your goal is purely to keep data on-device, switch to the local profile or the zero-dependency core path. Neither sends anything over the network, so you may not need redaction at all.

Stage rules

Custom stages obey the same contract as the built-ins. These are enforceable rules, not suggestions.

1. Communicate only through SpaceContext

Read your inputs from the context and write your outputs back to it. Use no module-level globals, no side channels, and no shared mutable singletons. The context is the single source of truth for everything the pipeline has accumulated.

# ❌ side channel — invisible to other stages, breaks resume & determinism
GLOBAL_CHUNKS.extend(new_chunks)

# ✅ everything flows through ctx
for chunk in self._dedup(ctx.space.chunks):
    ...

2. Be idempotent and resume-aware

Re-running a stage on its own output must not corrupt state or duplicate work. Skip what’s already recorded on the context or in the resume cache, rather than redoing it. This is what makes --resume cheap and keeps index.json byte-stable across reruns.

3. Never silently swallow errors

A per-item failure must be visible. You have two options. Pick one, but never use an empty except:

• Raise a typed error, an IndxError subclass such as StageError, with an actionable message. Say what failed, on which item or stage, and what to do about it. This is the fail-loud path for fatal problems.
• Record a non-fatal failure on ctx.errors and continue. Append a typed StageErrorRecord so the item is skipped visibly rather than disappearing.

from indx.core.context import StageErrorRecord

def run(self, ctx: SpaceContext) -> SpaceContext:
    for chunk in ctx.space.chunks:
        try:
            chunk.text = redact(chunk.text)
        except RedactionError as exc:
            ctx.errors.append(
                StageErrorRecord(
                    stage=self.name,
                    kind="skip",
                    item=chunk.id,
                    message=f"could not redact {chunk.id}: {exc}",
                )
            )
    return ctx

StageErrorRecord carries stage, kind, item, and message. The item is the offending file path or chunk id, and kind is either "skip" or "fatal".

Per-item failures recorded on ctx.errors stay on the run context. Under --strict (or strict=True), the first skip is promoted to a fatal StageError that aborts the run. See errors and exit codes for the full hierarchy.

4. Emit progress

Report progress through the shared reporting hook so the user sees per-stage output. That means Rich progress on the CLI and structured logs otherwise. Don’t run silently and don’t print(). Library code logs and reports; it never prints.

5. Be replaceable and optional

A stage must function without assuming a specific neighbor implementation. Don’t depend on a particular parser, store, or upstream stage being present. Depend on the shape of the context instead. A well-behaved stage can then be inserted, reordered, or dropped without breaking the rest of the pipeline.

Managing the stage list

A fresh DirectoryPipeline registers the six built-ins in canonical order. These methods let you reshape that list; each returns the pipeline for chaining.

Method	Effect
insert(index, stage)	Insert a stage at a 0-based position.
append(stage)	Add a stage to the end of the pipeline.
replace(name, stage)	Swap out the stage with the given name.
drop(name)	Remove the named stage entirely.
stages()	Return the current ordered list of stages (handy for checking indices).

from indx import DirectoryPipeline

pipeline = DirectoryPipeline(embedder="bge-m3", store="qdrant")

# Inspect current order
for i, s in enumerate(pipeline.stages()):
    print(i, s.name)        # 0 walk, 1 parse, 2 chunk, 3 relate, 4 enrich, 5 embed-pack

# Add a redaction pass before Relate, drop LLM enrichment, then run
space = (
    pipeline
    .insert(3, PiiRedactStage())   # after Chunk, before Relate
    .drop("enrich")                # skip LLM work entirely
    .run("./docs", "./ai-ready")
)

A few common patterns:

• insert(3, …) — a redaction or normalization pass between Chunk and Relate.
• append(MyExportStage()) — an extra export or notification step after the archive is sealed.
• replace("relate", CustomRelationStage()) — your own relation resolver in place of the built-in.
• drop("enrich") — skip LLM enrichment when no model is available or wanted. Using drop("embed-pack") yields a graph-only space with no vectors.

Note

Stage management changes the order and membership of stages. Swapping the heavy-lifting component inside a built-in stage — a different parser, embedder, or store — is a separate concern handled with use() or constructor arguments. See bring your own stack and configuration.

Packaging a stage for reuse

Inserting a stage in code is perfect for project-specific logic. To make a stage installable and resolvable by name anywhere a built-in is, register it as an entry-point plugin. Put it under the indx.stages group:

# pyproject.toml of your package
[project.entry-points."indx.stages"]
pii-redact = "my_pkg.stages:PiiRedactStage"

Once installed, indx discovers it at import and merges it into the registry. See authoring a plugin for the full packaging recipe, including extras, lazy imports, and contract tests.

See also

• Pipeline and stages — how the shared context flows through the six stages.
• Protocols reference — the full Stage protocol and component protocols.
• Authoring a plugin — package and distribute stages and components by name.