Data Components¶

def discover_documents(
    bucket_name: str,
    prefixes: str | list[str] | None = None,
    test_data_doc_names: list[str] | None = None,
    sampling_enabled: bool = True,
    sampling_max_size_gb: float = SAMPLING_MAX_SIZE_GB,
    supported_extensions: set[str] | None = None,
    validate_test_data_keys: bool = True,
    s3_client: Any | None = None,
) -> DiscoveryResult:
    """Discover documents across one or more bucket locations and optionally sample them.

    Lists objects under every entry of *prefixes*, merges the results into a
    single corpus deduplicated by object key, filters by file extension, and
    applies size-based sampling when enabled.  The sampling budget is shared by
    the whole union, not applied per prefix.  Documents referenced by
    ``test_data_doc_names`` are prioritized during sampling so that
    benchmark-relevant files are always included when the budget permits.

    Parameters
    ----------
    bucket_name : str
        S3-compatible bucket name.
    prefixes : str | list[str] | None, default=None
        Object-key prefixes to narrow the listing.  A bare string is treated as
        a one-element list.  ``None``, an empty list, or a list containing an
        empty string lists the whole bucket.  Overlapping prefixes are safe:
        objects matched by more than one are kept once.
    test_data_doc_names : list[str] | None, default=None
        Keys of documents referenced by the benchmark test data.  Each is
        matched against the full object key, or -- when the name resolves to
        exactly one document -- against a bare file name.  Matched documents
        are sorted first so that sampling picks them before other files.
    sampling_enabled : bool, default=True
        When ``True``, only documents up to *sampling_max_size_gb* total
        are returned.
    sampling_max_size_gb : float, default=1.0
        Maximum cumulative size (in gigabytes) when sampling is enabled.
    supported_extensions : set[str] | None, default=None
        File extensions to accept.  Defaults to
        :data:`~ai4rag.utils.data.constants.SUPPORTED_EXTENSIONS`.
    validate_test_data_keys : bool, default=True
        When ``True``, every entry of *test_data_doc_names* must identify
        exactly one discovered document, otherwise :class:`BenchmarkKeyError`
        is raised.  Set to ``False`` to downgrade the failure to a warning.
    s3_client : Any | None, default=None
        Pre-configured ``boto3`` S3 client.  When ``None``, one is created
        via :func:`ai4rag.utils.clients.s3.create_s3_client`.

    Returns
    -------
    DiscoveryResult
        Discovery outcome with document metadata.

    Raises
    ------
    RuntimeError
        If no supported documents are found under any of the prefixes.
    BenchmarkKeyError
        If a benchmark key matches no discovered document, or matches more
        than one, while *validate_test_data_keys* is enabled.
    ValueError
        If sampling produces an empty selection.
    """
    if supported_extensions is None:
        supported_extensions = set(SUPPORTED_EXTENSIONS)

    resolved_prefixes = _normalize_prefixes(prefixes)
    ext_tuple = tuple(supported_extensions)
    max_size_bytes = float(sampling_max_size_gb) * 1024**3 if sampling_enabled else float(inf)

    if s3_client is None:
        s3_client = _create_s3_client_with_ssl_fallback(bucket_name, resolved_prefixes[0])
    contents = _list_objects_union(s3_client, bucket_name, resolved_prefixes)
    supported_files = [c for c in contents if c["Key"].endswith(ext_tuple)]

    if not supported_files:
        raise RuntimeError(f"No supported documents found in {_location(bucket_name, resolved_prefixes)}.")

    test_keys: set[str] = set()
    if test_data_doc_names:
        test_keys = _resolve_test_data_keys(
            test_data_doc_names,
            supported_files,
            bucket_name=bucket_name,
            prefixes=resolved_prefixes,
            strict=validate_test_data_keys,
        )
        supported_files.sort(key=lambda c: c["Key"] not in test_keys)

    total_size = 0
    selected: list[DocumentDescriptor] = []
    for file_info in supported_files:
        size = file_info["Size"]
        if total_size + size > max_size_bytes:
            continue
        selected.append(DocumentDescriptor(key=file_info["Key"], size_bytes=size))
        total_size += size

    if not selected:
        raise ValueError(
            "No documents to process. Check that the bucket/prefixes are correct and contain supported files."
        )

    dropped = sorted(test_keys - {d.key for d in selected})
    if dropped:
        _logger.warning(
            "%d benchmark-referenced document(s) could not fit within the %.2f GB sampling budget and were skipped: %s",
            len(dropped),
            sampling_max_size_gb,
            ", ".join(dropped),
        )

    result = DiscoveryResult(
        bucket=bucket_name,
        prefixes=tuple(resolved_prefixes),
        documents=selected,
        total_size_bytes=total_size,
        count=len(selected),
    )
    _logger.info(
        "Discovered %d document(s) across %d location(s), total size %d bytes",
        result.count,
        len(resolved_prefixes),
        result.total_size_bytes,
    )
    return result

def extract_text(  # pylint: disable=too-many-locals,too-many-arguments,too-many-positional-arguments,too-many-statements
    documents: list[dict],
    bucket: str,
    output_dir: str | Path,
    s3_endpoint: str | None = None,
    s3_access_key: str | None = None,
    s3_secret_key: str | None = None,
    s3_region: str | None = None,
    error_tolerance: float | None = None,
    max_extraction_workers: int | None = None,
    docling_artifacts_path: str | None = None,
    docling_config: DoclingExtractionConfig | None = None,
) -> ExtractionResult:
    """Download documents from S3 and extract text using Docling.

    Each input document is downloaded from S3, converted to a
    :class:`DoclingDocument` via the Docling library, and persisted as a
    JSON file in *output_dir*.  Conversion runs in a separate process pool
    (``multiprocess`` library, ``"spawn"`` context) while downloads happen
    concurrently in a thread pool.

    Parameters
    ----------
    documents
        List of document descriptor dicts, each with at least a ``"key"``
        and ``"size_bytes"`` entry (as produced by
        :func:`~ai4rag.utils.data.documents_discovery.discover_documents`).
        The ``"key"`` entry also names the extracted document.
    bucket
        S3-compatible bucket name.
    output_dir
        Local directory where DoclingDocument JSON files are written.
    s3_endpoint
        S3-compatible endpoint URL.  Falls back to ``AWS_S3_ENDPOINT``.
    s3_access_key
        AWS access key.  Falls back to ``AWS_ACCESS_KEY_ID``.
    s3_secret_key
        AWS secret key.  Falls back to ``AWS_SECRET_ACCESS_KEY``.
    s3_region
        AWS region.  Falls back to ``AWS_DEFAULT_REGION``.
    error_tolerance
        Fraction of documents (0.0--1.0) allowed to fail.  ``None`` means
        zero tolerance.
    max_extraction_workers
        Number of parallel worker processes.  Defaults to
        ``min(max(1, cpu_count // 2), 8)``.
    docling_artifacts_path
        Path to pre-downloaded Docling model artifacts for offline use.
        Falls back to ``DOCLING_ARTIFACTS_PATH`` environment variable.
    docling_config
        Ready :class:`DoclingExtractionConfig` controlling table-structure
        and OCR behaviour.  Callers (e.g. ``pipelines-components``) construct
        it once and pass it here; it is forwarded unchanged to every worker
        process.  ``None`` (default) uses ``DoclingExtractionConfig()`` --
        table structure and OCR both disabled.  See
        :class:`DoclingExtractionConfig` for the per-field ``what``/``why``
        and for how ``ocr_lang`` maps to bundled OCR models.

    Returns
    -------
    ExtractionResult
        Summary of the extraction run.

    Raises
    ------
    RuntimeError
        If the error count exceeds the allowed tolerance.
    """
    import tempfile

    import multiprocess as multiprocessing

    out_dir = Path(output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    if not documents:
        _logger.info("No documents to process.")
        return ExtractionResult(processed_count=0, total_documents=0, error_count=0)

    s3_creds = _resolve_s3_credentials(s3_endpoint, s3_access_key, s3_secret_key, s3_region)
    artifacts_path = _resolve_artifacts_path(docling_artifacts_path)
    pipeline_config = docling_config or DoclingExtractionConfig()

    has_custom_models = bool(
        pipeline_config.ocr_det_model_path
        or pipeline_config.ocr_cls_model_path
        or pipeline_config.ocr_rec_model_path
        or pipeline_config.ocr_rec_keys_path
    )
    _logger.info("Docling table structure parsing: %s", pipeline_config.do_table_structure)
    _logger.info(
        "Docling OCR (RapidOCR): enabled=%s lang=%s custom_models=%s",
        pipeline_config.do_ocr,
        pipeline_config.ocr_lang if pipeline_config.do_ocr else (),
        has_custom_models,
    )

    documents = sorted(documents, key=lambda d: d.get("size_bytes", 0), reverse=True)

    effective_workers = _effective_worker_count(max_extraction_workers)
    _logger.info(
        "Starting text extraction for %d documents. extraction_workers=%d, download_threads=%d.",
        len(documents),
        effective_workers,
        DOWNLOAD_MAX_THREADS,
    )

    if artifacts_path is not None:
        os.environ.setdefault("HF_HUB_OFFLINE", "1")

    mp_context = multiprocessing.get_context("spawn")  # pylint: disable=no-member
    with (
        tempfile.TemporaryDirectory() as download_dir,
        mp_context.Pool(
            processes=effective_workers,
            initializer=_text_extraction_pool_initializer,
            initargs=(pipeline_config,),
        ) as process_pool,
    ):
        download_start = time.perf_counter()
        extraction_tasks, download_errors = _download_and_submit(
            docs=documents,
            bucket=bucket,
            download_path=Path(download_dir),
            process_pool=process_pool,
            out_dir=out_dir,
            s3_creds=s3_creds,
        )
        _logger.info(
            "Downloads finished in %.1fs; %d file(s) queued for extraction, %d download error(s).",
            time.perf_counter() - download_start,
            len(extraction_tasks),
            len(download_errors),
        )
        _raise_if_threshold_exceeded(download_errors, len(documents), error_tolerance)

        extraction_errors: list[dict] = []
        processed_count = 0
        pending = list(extraction_tasks)
        completed = 0

        while pending:
            still_pending = []
            for file_path, task in pending:
                if task.ready():
                    completed += 1
                    try:
                        success, tb = task.get()
                    except Exception:
                        tb = traceback.format_exc()
                        _logger.error("Worker crashed for %s:\n%s", file_path, tb)
                        success = False
                    Path(file_path).unlink(missing_ok=True)
                    if success:
                        processed_count += 1
                    else:
                        extraction_errors.append({"file": file_path, "traceback": tb})
                    _logger.info("Extraction progress %d/%d", completed, len(extraction_tasks))
                else:
                    still_pending.append((file_path, task))
            pending = still_pending
            if pending:
                time.sleep(0.01)

    all_errors = download_errors + extraction_errors
    total_errors = len(all_errors)
    _logger.info(
        "Text extraction completed. Total processed: %d/%d, Errors: %d",
        processed_count,
        len(documents),
        total_errors,
    )
    _raise_if_threshold_exceeded(
        error_details=all_errors,
        total_docs=len(documents),
        tolerance=error_tolerance,
    )

    return ExtractionResult(
        processed_count=processed_count,
        total_documents=len(documents),
        error_count=total_errors,
    )

def load_test_data(
    bucket_name: str,
    key: str,
    benchmark_sample_size: int = BENCHMARK_SAMPLE_SIZE,
    s3_client: Any | None = None,
) -> TestDataResult:
    """Download benchmark test data from S3 and optionally sample it.

    Parameters
    ----------
    bucket_name : str
        S3-compatible bucket containing the test data file.
    key : str
        Full S3 object key to the JSON test data file.
    benchmark_sample_size : int, default=25
        Maximum number of records to keep.  When the dataset exceeds this
        limit a reproducible random sample is drawn (seed 42).  Set to
        ``0`` to disable sampling and keep all records.
    s3_client : Any | None, default=None
        Pre-configured ``boto3`` S3 client.  When ``None``, one is created
        via :func:`ai4rag.utils.clients.s3.create_s3_client`.

    Returns
    -------
    TestDataResult
        Loaded (and optionally sampled) benchmark data.

    Raises
    ------
    FileNotFoundError
        If the object does not exist in S3.
    TestDataLoaderError
        If the file is not valid JSON or the records have an unexpected
        structure.
    """
    if not bucket_name:
        raise TypeError("bucket_name must be a non-empty string")

    if s3_client is None:
        s3_client = _make_s3_client_with_ssl_fallback(bucket_name, key)

    raw_data = _download_object(s3_client, bucket_name, key)
    benchmark_data = _parse_and_validate(raw_data)

    sampled = False
    if 0 < benchmark_sample_size < len(benchmark_data):
        original_count = len(benchmark_data)
        rng = random.Random(42)
        benchmark_data = rng.sample(benchmark_data, benchmark_sample_size)
        sampled = True
        _logger.info("Sampled %d records from %d total.", benchmark_sample_size, original_count)
    else:
        _logger.info("No sampling applied; record count: %d.", len(benchmark_data))

    return TestDataResult(data=benchmark_data, record_count=len(benchmark_data), sampled=sampled)