Stages API
Individual stage runner functions. Each can be called independently or via run_pipeline().
Stage 0: Preprocess
async
run_preprocess(
input_raw: str,
output_dir: str | None = None,
context_instructions: str | None = None,
parallel_docs: int = 1,
) -> tuple[StageResult, list]
Convert raw source files to intermediate JSON using the converters module.
Also stores raw files (binary) and converted pages in MongoDB via the
storage layer when a backend is configured (STORAGE_BACKEND env var).
Parameters
input_raw:
Path to the folder containing raw source files (PDF, DOCX, XLSX, etc.).
output_dir:
Path to the output folder where intermediate JSON files will be written.
If None, converted documents are only available in memory; no files are
written to disk.
context_instructions:
Optional free-text context about the documents being processed.
parallel_docs:
Maximum number of documents converted concurrently (default: 1,
i.e. sequential — matches pre-existing behaviour).
Returns
tuple[StageResult, list[IntermediateDocument]] A StageResult with counts and errors, and a list of IntermediateDocument objects (one per successfully converted file).
Stage 1: Extraction
async
run_extraction(
input_folder: str | None = None,
output_folder: str | None = None,
intermediate_documents: list | None = None,
parallel_docs: int = 1,
) -> tuple[StageResult, list]
Process intermediate JSON files or IntermediateDocument objects through the LLM extraction pipeline and produce Document objects.
Exactly one of input_folder or intermediate_documents must be provided. If output_folder is given, the extracted Document JSON is also written to disk (mirroring subdirectory structure). If not, documents only exist in memory.
Parameters
input_folder: Path to the directory containing input JSON files from Stage 0 (searched recursively). Mutually exclusive with intermediate_documents. output_folder: Path to the directory where extraction output files will be written. If None, extraction output is only available in memory. intermediate_documents: List of IntermediateDocument objects from run_preprocess() (in-memory mode). Mutually exclusive with input_folder. parallel_docs: Maximum number of documents to process concurrently.
Returns
tuple[StageResult, list[Document]] A StageResult with counts and errors, and a list of Document objects (one per successfully extracted document).
Stage 2: Ingestion
async
run_ingestion(
output_folder: str | None = None,
files: list[Path] | None = None,
documents: list | None = None,
update_mode: bool = False,
) -> StageResult
Load extracted documents into Neo4j.
Accepts documents either from disk (via output_folder or files) or
directly as in-memory :class:~models.document_structure.Document objects
(via documents). Exactly one source must be provided.
Parameters
output_folder:
Path to the directory containing extract-*.json files. Used when
neither files nor documents is provided.
files:
Explicit list of JSON extraction file paths to ingest. Takes priority
over output_folder when both are given.
documents:
List of in-memory Document objects from run_extraction() (in-memory
mode). When provided, output_folder and files are ignored.
update_mode:
If True, wipe existing structure of the latest version and re-insert
without creating a new version.
Returns
StageResult Stage result with one DocumentResult per ingested document.
Raises
ValueError If no source is provided.
Stage 3: Annotation
async
run_annotation(
document_name: str,
manual: bool = False,
model_class: str | None = None,
parallel_docs: int = 1,
only_unannotated: bool = False,
context_instructions_override: str | None = None,
) -> StageResult
Run the annotation stage for an already-ingested document.
In normal mode delegates to :func:annotation.agent.run_annotation_agent.
In manual mode delegates to :func:annotation.agent.run_manual_annotation,
assigning model_class to every qualifying StructureNode without invoking
the LLM.
Parameters
document_name: Name of the document node already present in Neo4j. manual: If True, run in manual override mode instead of the LLM agent. model_class: CamelCase model class name required when manual is True. parallel_docs: Maximum number of leaf documents to annotate concurrently when document_name refers to a folder. only_unannotated: When True, only process StructureNodes without a :HAS_MODEL_DECISION relationship. Ignored when manual is True. context_instructions_override: When provided, use this context string instead of fetching from Neo4j.
Returns
StageResult Stage result with per-document annotation counts and errors.
Raises
ValueError If document_name is empty, or if manual is True but model_class is not provided.
Stage 4: Entity Extraction
async
run_entity_extraction(
document_name: str,
parallel_docs: int = 1,
only_unextracted: bool = False,
) -> StageResult
Run the entity extraction agent for an already-annotated document.
Delegates to :func:entity_extraction.agent.run_entity_extraction_agent.
Parameters
document_name: Name of the document node already annotated in Neo4j. parallel_docs: Maximum number of leaf documents to extract concurrently when document_name refers to a folder. only_unextracted: When True, only process StructureNodes without a :HAS_EXTRACTION relationship.
Returns
StageResult Stage result with per-document extraction counts and errors.
Raises
ValueError If document_name is empty.
Stage 5: Tabular Pipeline
async
run_tabular_pipeline(
input_raw: str,
update_mode: bool = False,
parallel_docs: int = 1,
tabular_extensions: set | None = None,
tabular_delimiter: str | None = None,
) -> StageResult
Ingest all tabular files (CSV/XLSX/XLS) in input_raw directly into Neo4j.
Bypasses Stages 0-4 entirely. For each file: 1. Reads headers + 5-row preview. 2. Makes one LLM call to decide the extraction model. 3. Makes one LLM call to map columns to model fields. 4. Writes Table + Row StructureNode subgraph to Neo4j directly.
Parameters
input_raw: Folder containing raw tabular files (searched recursively). update_mode: If True, wipe existing Table/Row subgraph and re-insert. parallel_docs: Maximum number of files to process concurrently. tabular_extensions: Set of file extensions to treat as tabular. Defaults to {'.csv', '.xlsx', '.xls'} when None. tabular_delimiter: Field delimiter for CSV files. When None, uses the default delimiter of the tabular agent.
Returns
StageResult Stage result with per-document ingestion counts and errors.