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,
*,
tenant_id: str | None = None,
created_by_user_id: str | None = None,
job_id: str | None = None,
) -> 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).
tenant_id:
Tenant owning the converted files (None = public). Written on every
stored raw file and page, and stamped on each IntermediateDocument
(so a later extraction / ingestion of its JSON keeps the same tenant).
created_by_user_id, job_id:
Provenance of the upload, written and stamped like tenant_id.
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.
Note: this older batch entry point always calls extract_one_intermediate()/
extract_one_file() with fast_extraction=False — the fast_extraction
flag is only reachable via run_pipeline()'s per-document-unit engine
(_process_document_unit()). This is a deliberate scope boundary, not
an oversight.
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,
tenant_id: str | None = None,
created_by_user_id: str | None = None,
job_id: str | None = None,
) -> 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.
tenant_id:
Tenant owning the ingested documents (None = public). Part of the
:Document identity: versions are resolved, and update_mode wipes
content, within this tenant only. Overrides any tenant baked into an
ingested extract-*.json (see ingest.loader.load_file).
created_by_user_id, job_id:
Optional caller-supplied provenance metadata written onto every
:Document node created by this stage (see ingest.loader.load_file).
Returns
StageResult Stage result with one DocumentResult per ingested document.
Documents whose raw_file_id is not a stored raw file of their
effective tenant are rejected before any graph write and reported as
failed (see ingest.raw_file_check); the rest of the batch is ingested.
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,
*,
tenant_id: str | None = None,
doc_path: 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. Used to select the
document only when doc_path is not given (then every latest document
of tenant_id with that name is processed).
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.
tenant_id:
Owner of the document (None = public document). Only that tenant's
documents are read or written.
doc_path:
Path of the document — with tenant_id, the preferred selector.
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,
*,
tenant_id: str | None = None,
doc_path: str | None = None,
) -> 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. Used to select
the document only when doc_path is not given (then every latest
document of tenant_id with that name is processed).
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.
tenant_id:
Owner of the document (None = public document). Only that tenant's
documents are read or written.
doc_path:
Path of the document — with tenant_id, the preferred selector.
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,
tenant_id: str | None = None,
created_by_user_id: str | None = None,
job_id: str | None = None,
*,
files: Sequence[tuple[Path, int]] | None = None,
emitter: _ProgressEmitter = _NULL_EMITTER,
) -> StageResult
Ingest all tabular files (CSV/XLSX/XLS) in input_raw directly into Neo4j.
Bypasses Stages 0-4 entirely. For each file: 1. Scans it (headers, row count, 5-row preview). A file that cannot be read, or whose sheets have no data rows, is skipped: no raw file is stored and no :Document is created. 2. Stores the raw file and creates the :Document node. 3. Per sheet: one LLM call to decide the extraction model, one to map columns to model fields, then writes the Table + Row StructureNode subgraph to Neo4j directly.
Each file is reported in the returned :class:StageResult as one
:class:DocumentResult whose counters are sheets: nodes_processed =
tables written, nodes_failed = sheets with errors (at most one per error
the agent recorded). A file is "success" when every sheet was written
without errors, "partial" when at least one table was written but some
sheet failed, "failed" when no table was written, and "skipped"
(0 / 0, so it never fails the batch) when it was rejected by step 1.
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.
tenant_id:
Tenant owning every document this pipeline creates (None = public).
Part of the :Document identity: versions are resolved within this
tenant only, and another tenant's document at the same path is never
touched.
created_by_user_id, job_id:
Optional caller-supplied provenance metadata written verbatim onto
every :Document node this pipeline creates (leaf tabular documents and
their ancestor folder-parent nodes). Always SET (null when omitted).
files:
Tabular files already discovered by the caller, each with its
document_index for progress events. None discovers them under
input_raw (sorted, indices 0..N-1).
emitter:
Progress emitter of the run (run_pipeline() passes its own). Each
file emits document_started → stage_started →
stage_completed / stage_failed → document_finished with
stage="tabular" and its document_index; a skipped file emits
document_finished(status="skipped") without a stage outcome.
Returns
StageResult Stage result with per-document ingestion counts and errors.