Stages API

scinr.newton.stages.run_preprocess async

run_preprocess(
    input_raw, output_dir=None, context_instructions=None
)

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.

Returns

tuple[StageResult, list[IntermediateDocument]] A StageResult with counts and errors, and a list of IntermediateDocument objects (one per successfully converted file).

Source code in src/scinr/newton/stages/preprocess.py

async def run_preprocess(
    input_raw: str,
    output_dir: str | None = None,
    context_instructions: 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.

    Returns
    -------
    tuple[StageResult, list[IntermediateDocument]]
        A StageResult with counts and errors, and a list of IntermediateDocument
        objects (one per successfully converted file).
    """
    import shutil
    import tempfile

    from scinr.newton.converters.main import convert_folder  # deferred import

    t0 = time.monotonic()
    raw_file_repo, page_repo = get_storage()

    # If no output_dir specified, use a temp dir and clean it up after
    _temp_dir: str | None = None
    if output_dir is None:
        _temp_dir = tempfile.mkdtemp(prefix="scinr_preprocess_")
        effective_output_dir = _temp_dir
    else:
        effective_output_dir = output_dir

    try:
        results = await convert_folder(
            Path(input_raw),
            Path(effective_output_dir),
            raw_file_repo=raw_file_repo,
            page_repo=page_repo,
            context_instructions=context_instructions,
        )
    except Exception as exc:
        duration = time.monotonic() - t0
        if _temp_dir:
            shutil.rmtree(_temp_dir, ignore_errors=True)
        return (
            StageResult(
                stage="preprocess",
                success=False,
                documents=[],
                total_processed=0,
                total_failed=0,
                duration_seconds=duration,
                errors=[str(exc)],
            ),
            [],
        )

    intermediate_docs = []
    doc_results = []
    for raw_path, _json_path, doc in results:
        intermediate_docs.append(doc)
        doc_results.append(
            DocumentResult(
                document_name=doc.document_name or raw_path.stem,
                nodes_processed=1,
                nodes_failed=0,
            )
        )

    # Clean up temp dir if we used one (docs already captured in memory)
    if _temp_dir:
        shutil.rmtree(_temp_dir, ignore_errors=True)

    duration = time.monotonic() - t0
    stage_result = StageResult(
        stage="preprocess",
        success=True,
        documents=doc_results,
        total_processed=len(doc_results),
        total_failed=0,
        duration_seconds=duration,
    )
    return stage_result, intermediate_docs

scinr.newton.stages.run_extraction async

run_extraction(
    input_folder=None,
    output_folder=None,
    intermediate_documents=None,
    parallel_docs=1,
)

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).

Source code in src/scinr/newton/stages/extraction.py

async def 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).
    """
    from scinr.newton.converters.base import IntermediateDocument  # deferred

    t0 = time.monotonic()

    # ── Validate inputs ────────────────────────────────────────────────────────
    if input_folder is not None and intermediate_documents is not None:
        raise ValueError(
            "input_folder and intermediate_documents are mutually exclusive. "
            "Provide one or the other, not both."
        )
    if input_folder is None and intermediate_documents is None:
        raise ValueError(
            "Either input_folder or intermediate_documents must be provided."
        )

    output_path = Path(output_folder) if output_folder else None

    semaphore = asyncio.Semaphore(parallel_docs)

    async def _process_intermediate(doc: IntermediateDocument) -> Document | None:
        """Process an IntermediateDocument in-memory (no disk read)."""
        async with semaphore:
            doc_name = doc.document_name or "unnamed"
            pages: list[str] = [p.markdown for p in doc.pages]
            page_ids_by_index: dict[int, str] = {
                p.index: p.page_id
                for p in doc.pages
                if p.page_id
            }
            folder_path = doc.folder_path
            context_instructions = doc.context_instructions
            total_pages = len(pages)

            logger.info("Processing in-memory document: %s (%d page(s))", doc_name, total_pages)

            if total_pages == 0:
                logger.warning("Document has no pages — skipping: %s", doc_name)
                return None

            batch_size = max(1, int(os.getenv("EXTRACTION_BATCH_SIZE", str(_DEFAULT_BATCH_SIZE))))
            chunks: list[tuple[str, list[str]]] = [("", pages[:batch_size])]
            for i in range(batch_size, total_pages, batch_size):
                chunks.append((pages[i - 1], pages[i : i + batch_size]))

            doc_path = f"{folder_path}/{doc_name}" if folder_path else doc_name
            ext = ""  # No file extension for in-memory docs

            document = Document(
                document_name=doc_name,
                document_type=ext,
                document_structure=[],
                doc_path=doc_path,
                raw_file_id=doc.raw_file_id or "",
                context_instructions=context_instructions,
            )
            llm = get_llm()

            # Set up output file if output_folder provided
            if output_path:
                if folder_path:
                    out_subdir = output_path / Path(folder_path)
                else:
                    out_subdir = output_path
                out_subdir.mkdir(parents=True, exist_ok=True)
                output_file = out_subdir / f"extract-{doc_name}.json"
            else:
                output_file = None

            for chunk_idx, (prev_page, curr_pages) in enumerate(chunks):
                active_hierarchy = get_active_hierarchy(document)
                curr_start_idx = chunk_idx * batch_size
                curr_page_ids = [
                    page_ids_by_index[curr_start_idx + i]
                    for i in range(len(curr_pages))
                    if (curr_start_idx + i) in page_ids_by_index
                ] or None

                try:
                    nodes = await extract_chunk(
                        prev_page=prev_page,
                        curr_pages=curr_pages,
                        active_hierarchy=active_hierarchy,
                        llm=llm,
                        curr_page_ids=curr_page_ids,
                        user_context=document.context_instructions or "",
                    )
                except ExtractionMaxRetriesError:
                    logger.warning(
                        "Chunk %d/%d: all extraction attempts failed — skipping chunk.",
                        chunk_idx + 1, len(chunks),
                    )
                    continue

                document = compact_extraction(document, nodes)

                # Intermediate crash-safe write if output requested
                if output_file:
                    output_file.write_text(document.model_dump_json(indent=2), encoding="utf-8")

            # Final write if output requested
            if output_file:
                output_file.write_text(document.model_dump_json(indent=2), encoding="utf-8")
                logger.info("Written: %s", output_file)

            return document

    async def _process_file(json_file: Path) -> Document | None:
        """Process a JSON file from disk."""
        async with semaphore:
            logger.info("Processing file: %s", json_file)

            raw = json.loads(json_file.read_text(encoding="utf-8"))
            pages: list[str] = [page["markdown"] for page in raw["pages"]]
            page_ids_by_index: dict[int, str] = {
                page["index"]: page["page_id"]
                for page in raw["pages"]
                if page.get("page_id")
            }
            folder_path: str | None = raw.get("folder_path")
            context_instructions: str | None = raw.get("context_instructions")
            total_pages = len(pages)
            logger.info("  Total pages: %d", total_pages)

            if total_pages == 0:
                logger.warning("  File has no pages — skipping: %s", json_file)
                return None

            batch_size = max(1, int(os.getenv("EXTRACTION_BATCH_SIZE", str(_DEFAULT_BATCH_SIZE))))
            chunks: list[tuple[str, list[str]]] = [("", pages[:batch_size])]
            for i in range(batch_size, total_pages, batch_size):
                chunks.append((pages[i - 1], pages[i : i + batch_size]))

            name, ext = os.path.splitext(json_file)
            nombre = Path(name).name
            doc_path = f"{folder_path}/{nombre}" if folder_path else nombre

            document = Document(
                document_name=nombre,
                document_type=ext,
                document_structure=[],
                doc_path=doc_path,
                raw_file_id=raw.get("raw_file_id") or "",
                context_instructions=context_instructions,
            )
            llm = get_llm()

            # Determine relative structure to preserve subdir layout
            if input_folder:
                try:
                    rel_dir = json_file.relative_to(Path(input_folder)).parent
                except ValueError:
                    rel_dir = Path(".")
            else:
                rel_dir = Path(".")

            if output_path:
                output_subdir = output_path / rel_dir
                output_subdir.mkdir(parents=True, exist_ok=True)
                output_file = output_subdir / f"extract-{json_file.stem}.json"
            else:
                output_file = None

            for chunk_idx, (prev_page, curr_pages) in enumerate(chunks):
                active_hierarchy = get_active_hierarchy(document)
                curr_start_idx = chunk_idx * batch_size
                curr_page_ids = [
                    page_ids_by_index[curr_start_idx + i]
                    for i in range(len(curr_pages))
                    if (curr_start_idx + i) in page_ids_by_index
                ] or None

                try:
                    nodes = await extract_chunk(
                        prev_page=prev_page,
                        curr_pages=curr_pages,
                        active_hierarchy=active_hierarchy,
                        llm=llm,
                        curr_page_ids=curr_page_ids,
                        user_context=document.context_instructions or "",
                    )
                except ExtractionMaxRetriesError:
                    logger.warning(
                        "  Chunk %d/%d: all extraction attempts failed — skipping chunk.",
                        chunk_idx + 1, len(chunks),
                    )
                    continue

                document = compact_extraction(document, nodes)

                if output_file:
                    output_file.write_text(document.model_dump_json(indent=2), encoding="utf-8")

            if output_file:
                output_file.write_text(document.model_dump_json(indent=2), encoding="utf-8")
                logger.info("  Written: %s", output_file)

            return document

    # ── Dispatch: in-memory or from-disk ──────────────────────────────────────
    extracted_docs: list[Document] = []
    doc_results: list[DocumentResult] = []

    if intermediate_documents is not None:
        tasks = [_process_intermediate(doc) for doc in intermediate_documents]
        results = await asyncio.gather(*tasks, return_exceptions=True)
        for i, result in enumerate(results):
            src_name = (intermediate_documents[i].document_name or f"doc_{i}")
            if isinstance(result, Exception):
                logger.error("Extraction failed for '%s': %s", src_name, result)
                doc_results.append(DocumentResult(document_name=src_name, nodes_processed=0, nodes_failed=1, errors=[str(result)]))
            elif result is None:
                doc_results.append(DocumentResult(document_name=src_name, nodes_processed=0, nodes_failed=1, errors=["No pages to process"]))
            else:
                extracted_docs.append(result)
                doc_results.append(DocumentResult(document_name=result.document_name, nodes_processed=1, nodes_failed=0))
    else:
        input_path = Path(input_folder)
        if not input_path.exists():
            raise FileNotFoundError(f"Input folder not found: '{input_folder}'.")
        json_files = sorted(input_path.rglob("*.json"))
        if not json_files:
            raise FileNotFoundError(f"No .json files found in '{input_folder}'.")
        tasks = [_process_file(f) for f in json_files]
        results = await asyncio.gather(*tasks, return_exceptions=True)
        for i, result in enumerate(results):
            src_name = json_files[i].stem
            if isinstance(result, Exception):
                logger.error("Extraction failed for '%s': %s", src_name, result)
                doc_results.append(DocumentResult(document_name=src_name, nodes_processed=0, nodes_failed=1, errors=[str(result)]))
            elif result is None:
                doc_results.append(DocumentResult(document_name=src_name, nodes_processed=0, nodes_failed=1, errors=["No pages to process"]))
            else:
                extracted_docs.append(result)
                doc_results.append(DocumentResult(document_name=result.document_name, nodes_processed=1, nodes_failed=0))

    total_failed = sum(1 for r in doc_results if r.nodes_failed > 0)
    duration = time.monotonic() - t0
    stage_result = StageResult(
        stage="extraction",
        success=total_failed == 0,
        documents=doc_results,
        total_processed=len(extracted_docs),
        total_failed=total_failed,
        duration_seconds=duration,
    )
    return stage_result, extracted_docs

scinr.newton.stages.run_ingestion async

run_ingestion(
    output_folder=None,
    files=None,
    documents=None,
    update_mode=False,
)

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.

Source code in src/scinr/newton/stages/ingestion.py

async def 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.
    """
    t0 = time.monotonic()

    if documents is None and files is None and output_folder is None:
        raise ValueError(
            "At least one of documents, files, or output_folder must be provided."
        )

    driver = get_driver()
    doc_results: list[DocumentResult] = []
    try:
        setup_schema(driver)

        if documents is not None:
            # In-memory mode
            doc_names = load_documents(documents, driver, update_mode=update_mode)
            for doc in documents:
                success = doc.document_name in doc_names
                doc_results.append(DocumentResult(
                    document_name=doc.document_name,
                    nodes_processed=1 if success else 0,
                    nodes_failed=0 if success else 1,
                    errors=[] if success else [f"Failed to ingest '{doc.document_name}'"],
                ))
        elif files is not None:
            doc_names = load_files(files, driver, update_mode=update_mode)
            for path in files:
                doc_name = path.stem.removeprefix("extract-")
                success = doc_name in doc_names
                doc_results.append(DocumentResult(
                    document_name=doc_name,
                    nodes_processed=1 if success else 0,
                    nodes_failed=0 if success else 1,
                    errors=[] if success else [f"Failed to ingest '{path.name}'"],
                ))
        else:
            output_path = Path(output_folder)
            if not output_path.exists():
                raise FileNotFoundError(
                    f"Output folder not found: '{output_folder}'. "
                    f"Run run_extraction() first."
                )
            extract_files = list(output_path.rglob("extract-*.json"))
            if not extract_files:
                raise FileNotFoundError(
                    f"No 'extract-*.json' files found in '{output_folder}'. "
                    f"Run run_extraction() first."
                )
            doc_names = load_folder(output_path, driver, update_mode=update_mode)
            for path in extract_files:
                doc_name = path.stem.removeprefix("extract-")
                success = doc_name in doc_names
                doc_results.append(DocumentResult(
                    document_name=doc_name,
                    nodes_processed=1 if success else 0,
                    nodes_failed=0 if success else 1,
                    errors=[] if success else [f"Failed to ingest '{path.name}'"],
                ))
    finally:
        driver.close()

    total_failed = sum(1 for r in doc_results if r.nodes_failed > 0)
    duration = time.monotonic() - t0
    return StageResult(
        stage="ingestion",
        success=total_failed == 0,
        documents=doc_results,
        total_processed=sum(r.nodes_processed for r in doc_results),
        total_failed=total_failed,
        duration_seconds=duration,
    )

scinr.newton.stages.run_annotation async

run_annotation(
    document_name,
    manual=False,
    model_class=None,
    parallel_docs=1,
    only_unannotated=False,
    context_instructions_override=None,
)

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.

Source code in src/scinr/newton/stages/annotation.py

async def 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.
    """
    if not document_name:
        raise ValueError("document_name must not be empty for the annotation stage.")
    if manual and not model_class:
        raise ValueError("model_class is required when manual=True.")

    t0 = time.monotonic()

    if manual:
        count = await run_manual_annotation(document_name, model_class)
        logger.info(
            "Manual annotation complete: assigned '%s' to %d nodes in '%s'",
            model_class, count, document_name,
        )
        duration = time.monotonic() - t0
        return StageResult(
            stage="annotation",
            success=True,
            documents=[DocumentResult(
                document_name=document_name,
                nodes_processed=count,
                nodes_failed=0,
            )],
            total_processed=count,
            total_failed=0,
            duration_seconds=duration,
        )

    agent_result = await run_annotation_agent(
        document_name,
        parallel_docs=parallel_docs,
        only_unannotated=only_unannotated,
        context_instructions_override=context_instructions_override,
    )

    duration = time.monotonic() - t0

    # Build DocumentResult(s) from agent output
    if "results" in agent_result:
        # Multi-leaf (folder) case
        doc_results = []
        for leaf in agent_result.get("results", []):
            nodes = leaf.get("nodes_to_annotate", [])
            errors = leaf.get("errors", [])
            doc_results.append(DocumentResult(
                document_name=leaf.get("document_name", document_name),
                nodes_processed=len(nodes),
                nodes_failed=len(errors),
                errors=errors,
            ))
    else:
        # Single document case
        nodes = agent_result.get("nodes_to_annotate", [])
        errors = agent_result.get("errors", [])
        doc_results = [DocumentResult(
            document_name=document_name,
            nodes_processed=len(nodes),
            nodes_failed=len(errors),
            errors=errors,
        )]

    total_processed = sum(r.nodes_processed for r in doc_results)
    total_failed = sum(r.nodes_failed for r in doc_results)
    return StageResult(
        stage="annotation",
        success=total_failed == 0,
        documents=doc_results,
        total_processed=total_processed,
        total_failed=total_failed,
        duration_seconds=duration,
    )

scinr.newton.stages.run_entity_extraction async

run_entity_extraction(
    document_name, parallel_docs=1, only_unextracted=False
)

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.

Source code in src/scinr/newton/stages/entity_extraction.py

async def 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.
    """
    if not document_name:
        raise ValueError("document_name must not be empty for the entity_extract stage.")

    t0 = time.monotonic()
    from scinr.newton.entity_extraction.agent import run_entity_extraction_agent

    agent_result = await run_entity_extraction_agent(
        document_name,
        parallel_docs=parallel_docs,
        only_unextracted=only_unextracted,
    )

    duration = time.monotonic() - t0

    if "results" in agent_result:
        doc_results = []
        for leaf in agent_result.get("results", []):
            targets = leaf.get("targets", [])
            errors = leaf.get("errors", [])
            doc_results.append(DocumentResult(
                document_name=leaf.get("document_name", document_name),
                nodes_processed=len(targets),
                nodes_failed=len(errors),
                errors=errors,
            ))
    else:
        targets = agent_result.get("targets", [])
        errors = agent_result.get("errors", [])
        doc_results = [DocumentResult(
            document_name=document_name,
            nodes_processed=len(targets),
            nodes_failed=len(errors),
            errors=errors,
        )]

    total_processed = sum(r.nodes_processed for r in doc_results)
    total_failed = sum(r.nodes_failed for r in doc_results)
    return StageResult(
        stage="entity_extraction",
        success=total_failed == 0,
        documents=doc_results,
        total_processed=total_processed,
        total_failed=total_failed,
        duration_seconds=duration,
    )

scinr.newton.stages.run_tabular_pipeline async

run_tabular_pipeline(
    input_raw,
    update_mode=False,
    parallel_docs=1,
    tabular_extensions=None,
    tabular_delimiter=None,
)

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.

Source code in src/scinr/newton/stages/tabular.py

async def 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.
    """
    from scinr.newton.annotation.neo4j_ops import ensure_catalog_models, ensure_theme_structure
    from scinr.newton.tabular.agent import run_tabular_agent
    from scinr.newton.utils.theme_registry import get_theme_registry
    theme_registry = get_theme_registry()

    t0 = time.monotonic()
    _TABULAR_EXTS = tabular_extensions if tabular_extensions is not None else {".csv", ".xlsx", ".xls"}
    input_path = Path(input_raw)

    if not input_path.exists():
        raise FileNotFoundError(f"--input-raw path not found: {input_raw}")

    tabular_files = sorted(
        f for f in input_path.rglob("*")
        if f.is_file() and f.suffix.lower() in _TABULAR_EXTS
    )
    if not tabular_files:
        logger.warning("run_tabular_pipeline: no tabular files found in '%s'", input_raw)
        duration = time.monotonic() - t0
        return StageResult(
            stage="tabular",
            success=True,
            documents=[],
            total_processed=0,
            total_failed=0,
            duration_seconds=duration,
        )

    logger.info(
        "run_tabular_pipeline: found %d tabular file(s) in '%s'",
        len(tabular_files), input_raw,
    )

    def _file_to_doc_path(f: Path) -> tuple[str, str]:
        doc_name = f.stem
        try:
            rel = f.relative_to(input_path)
            folder = "/".join(rel.parts[:-1]) if len(rel.parts) > 1 else ""
        except ValueError:
            folder = ""
        doc_path = f"{folder}/{doc_name}" if folder else doc_name
        return doc_name, doc_path

    file_infos = [_file_to_doc_path(f) for f in tabular_files]
    all_doc_paths = list({dp for _, dp in file_infos})
    expanded_paths: set[str] = set()
    for dp in all_doc_paths:
        expanded_paths.add(dp)
        parts = dp.split("/")
        for i in range(1, len(parts)):
            expanded_paths.add("/".join(parts[:i]))
    all_paths = list(expanded_paths)

    driver = get_driver()
    doc_results: list[DocumentResult] = []
    try:
        setup_schema(driver)

        # Fix: ensure_catalog_models and ensure_theme_structure are async and need AsyncDriver
        logger.info("Getting async neo4j driver")

        async_driver = get_async_driver()

        logger.info("Ensuring catalog models are loaded")

        await ensure_catalog_models(async_driver)

        logger.info("Ensuring themes are loaded")


        await ensure_theme_structure(async_driver, theme_registry)
        logger.info("Setup of theme and catalog complete")

        with driver.session() as session:
            logger.info("Verifying previous documents")
            if update_mode:
                result = session.run(
                    "MATCH (d:Document {latest: true}) WHERE d.path IN $paths "
                    "RETURN max(d.version) AS max_version",
                    paths=all_paths,
                )
            else:

                result = session.run(
                    "MATCH (d:Document) WHERE d.path IN $paths "
                    "RETURN max(d.version) AS max_version",
                    paths=all_paths,
                )
            record = result.single()
            max_version = record["max_version"] if record else None
            resolved_version = max_version if (update_mode and max_version is not None) else ((max_version + 1) if max_version is not None else 1)
            logger.info("Max version of ingested document selected")


        logger.info("run_tabular_pipeline: batch version=%d (update_mode=%s)", resolved_version, update_mode)

        semaphore = asyncio.Semaphore(parallel_docs)

        from scinr.newton.storage.factory import get_storage
        try:
            _raw_file_repo, _ = get_storage()
        except Exception as exc:
            logger.warning("Storage backend unavailable: %s. Continuing without persisting raw files.", exc)
            from scinr.newton.storage.null import NullRawFileRepository
            _raw_file_repo = NullRawFileRepository()

        async def _process_file(f: Path, doc_name: str, doc_path: str) -> str | None:
            async with semaphore:
                logger.info("run_tabular_pipeline: processing '%s'", f.name)
                try:
                    _driver = get_driver()
                    try:
                        raw_file_id = None
                        if _raw_file_repo is not None:
                            try:
                                raw_bytes = f.read_bytes()
                                _content_type = decide_content_type(f.suffix.lower())
                                folder_path_str = str(f.parent.relative_to(input_path)) if f.parent != input_path else ""
                                raw_file_id = await _raw_file_repo.store(
                                    filename=f.name,
                                    content=raw_bytes,
                                    content_type=_content_type,
                                    folder_path=folder_path_str,
                                )
                            except Exception as store_exc:
                                logger.warning(
                                    "run_tabular_pipeline: failed to store raw file '%s': %s",
                                    f.name, store_exc,
                                )

                        agent_kwargs = dict(
                            file_path=f,
                            document_name=doc_name,
                            doc_path=doc_path,
                            driver=_driver,
                            resolved_version=resolved_version,
                            update_mode=update_mode,
                            raw_file_id=raw_file_id,
                        )
                        if tabular_delimiter is not None:
                            agent_kwargs["delimiter"] = tabular_delimiter

                        final_state = await run_tabular_agent(**agent_kwargs)
                        if final_state.get("errors"):
                            logger.warning(
                                "run_tabular_pipeline: '%s' completed with errors: %s",
                                f.name, final_state["errors"],
                            )
                        return doc_name
                    finally:
                        _driver.close()
                except Exception as exc:
                    logger.exception("run_tabular_pipeline: failed to process '%s': %s", f.name, exc)
                    return None

        tasks = [
            _process_file(f, doc_name, doc_path)
            for f, (doc_name, doc_path) in zip(tabular_files, file_infos)
        ]
        results = await asyncio.gather(*tasks, return_exceptions=True)

        for i, result in enumerate(results):
            doc_name = file_infos[i][0]
            if isinstance(result, Exception) or result is None:
                err = str(result) if isinstance(result, Exception) else "Processing failed"
                doc_results.append(DocumentResult(document_name=doc_name, nodes_processed=0, nodes_failed=1, errors=[err]))
            else:
                doc_results.append(DocumentResult(document_name=doc_name, nodes_processed=1, nodes_failed=0))

        success_count = sum(1 for r in doc_results if r.nodes_processed > 0)
        logger.info("run_tabular_pipeline: complete. %d/%d file(s) processed.", success_count, len(tabular_files))
    finally:
        driver.close()

    total_failed = sum(r.nodes_failed for r in doc_results)
    duration = time.monotonic() - t0
    return StageResult(
        stage="tabular",
        success=total_failed == 0,
        documents=doc_results,
        total_processed=sum(r.nodes_processed for r in doc_results),
        total_failed=total_failed,
        duration_seconds=duration,
    )