extraction/generic_json_extractor.py

from __future__ import annotations

from pathlib import Path
from typing import Any, Dict, Optional, Union

from logger import get_logger
from .label_mapper import validate_mappings
from .spacy_extractor import extract_fields

logger = get_logger(__name__)

TABULAR_EXTENSIONS = {".csv", ".xls", ".xlsx"}


def is_tabular(filename: Union[str, Path]) -> bool:
    return Path(filename).suffix.lower() in TABULAR_EXTENSIONS


def extract(
    filename: Union[str, Path],
    markdown_text: str,
    mappings: Optional[Dict[str, Dict[str, Any]]],
    file_data: Optional[bytes] = None,
) -> Dict[str, Any]:
    ext = Path(filename).suffix.lower()

    if ext in TABULAR_EXTENSIONS:
        from .json_extractor import extract_json_from_file
        result = extract_json_from_file(filename, file_data)
        if "error" not in result:
            result["extractor"] = "pandas"
        return result

    if not mappings:
        return {
            "error": (
                f"Cannot extract JSON from '{ext}' files without field mappings. "
                "Provide a 'mappings' object with field extraction rules."
            ),
            "file_type": ext,
        }

    valid, mapper_error = validate_mappings(mappings)
    if not valid:
        return {
            "error": "invalid_spacy_labels",
            "label_mapper": mapper_error,
            "file_type": ext,
        }

    try:
        data = extract_fields(markdown_text, mappings)
        logger.info("spaCy extraction completed for %s: %d fields", ext, len(data))
        return {
            "success": True,
            "extractor": "spacy",
            "file_type": ext,
            "data": data,
        }
    except Exception as exc:
        logger.exception("spaCy extraction failed for %s", ext)
        return {
            "error": f"spaCy extraction failed: {exc}",
            "file_type": ext,
            "exception_type": type(exc).__name__,
        }