ETL Module API Reference¶

def process_pdf_with_tagging(
    pdf_path: str,
    base_metadata: dict[str, Any] | None = None,
    auto_tag: bool = True,
) -> list[dict[str, Any]]:
    """Process a single PDF with optional auto-tagging.

    Args:
    ----
        pdf_path: Path to PDF file
        base_metadata: Base metadata from config
        auto_tag: Whether to auto-tag with ESG metadata

    Returns:
    -------
        List of chunked documents with metadata

    """
    # Load PDF
    loader = PyPDFLoader(pdf_path)
    pages = loader.load()

    # Add base metadata to all pages
    if base_metadata:
        for page in pages:
            page.metadata.update(base_metadata)

    # Auto-tag if enabled
    if auto_tag:
        tagger = ESGOpenAITagger()
        pages = tagger.tag_all(pages)

    # Chunk documents
    chunker = TextChunker(chunk_size=1000, chunk_overlap=100)

    chunks = []
    for page in pages:
        doc_dict = {"content": page.page_content, "metadata": page.metadata}
        page_chunks = chunker.chunk_docs([doc_dict])
        chunks.extend(page_chunks)

    return chunks

def load_config():
    """Load YAML configuration for documents."""
    with open(CONFIG_PATH, encoding="utf-8") as f:
        return yaml.safe_load(f)

def sha256sum(file_path):
    """Compute SHA256 hash of a file."""
    h = hashlib.sha256()
    with open(file_path, "rb") as f:
        for chunk in iter(lambda: f.read(DEFAULT_HASH_CHUNK_SIZE), b""):
            h.update(chunk)
    return h.hexdigest()

def download_file(
    url,
    dest_path,
    retries=DEFAULT_DOWNLOAD_RETRIES,
    backoff=DEFAULT_DOWNLOAD_BACKOFF,
) -> bool:
    """Download file with retry logic.

    Uses browser-like headers to avoid bot detection (Cloudflare, etc.).

    Args:
        url: URL to download
        dest_path: Destination file path
        retries: Number of retry attempts
        backoff: Backoff multiplier for retries

    Returns:
        True if successful, False otherwise
    """
    attempt = 0
    last_status_code = None

    while attempt < retries:
        try:
            resp = requests.get(
                url,
                timeout=DEFAULT_DOWNLOAD_TIMEOUT,
                headers=DEFAULT_HTTP_HEADERS,
                allow_redirects=True,
            )
            last_status_code = resp.status_code

            # Check for bot protection (403/503 from Cloudflare, etc.)
            if resp.status_code in (403, 503):
                # Check if it's Cloudflare protection
                is_cloudflare = (
                    "cloudflare" in resp.headers.get("Server", "").lower()
                    or "cf-ray" in resp.headers
                    or "cf-mitigated" in resp.headers
                )
                if is_cloudflare:
                    logger.warning(
                        f"Cloudflare protection detected for {url}. "
                        "Manual download required."
                    )
                    # Don't retry Cloudflare-protected URLs
                    return False

            resp.raise_for_status()

            with open(dest_path, "wb") as f:
                f.write(resp.content)
            return True

        except requests.exceptions.HTTPError as e:
            attempt += 1
            if last_status_code in (403, 403):
                logger.error(
                    f"Access denied (HTTP {last_status_code}) for {url}. "
                    "Likely bot protection. Skipping retries."
                )
                break
            logger.error(f"HTTP error downloading {url}: {e} (attempt {attempt})")
            if attempt < retries:
                time.sleep(backoff**attempt)

        except Exception as e:
            attempt += 1
            logger.error(f"Error downloading {url}: {e} (attempt {attempt})")
            if attempt < retries:
                time.sleep(backoff**attempt)

    return False

def detect_file_type(file_path: Path) -> str | None:
    """Detect file type from magic bytes.

    Args:
        file_path: Path to file

    Returns:
        File extension (.pdf, .html, etc.) or None if unknown
    """
    if not file_path.exists():
        return None

    # Read first 16 bytes for magic number detection
    try:
        with open(file_path, "rb") as f:
            header = f.read(16)
    except Exception:
        return None

    # PDF magic bytes
    if header.startswith(b"%PDF"):
        return ".pdf"

    # HTML detection (check for common HTML tags in first 1KB)
    try:
        with open(file_path, "rb") as f:
            content = f.read(1024).lower()
            if any(
                tag in content for tag in [b"<html", b"<!doctype", b"<head", b"<body"]
            ):
                return ".html"
    except Exception:
        pass

    # XML/HTML variants
    if header.startswith(b"<?xml") or header.startswith(b"<"):
        return ".html"

    return None

def safe_filename(url):
    """Generate a safe filename from a URL."""
    parsed = urlparse(url)
    filename = os.path.basename(parsed.path)
    return filename or DEFAULT_DOWNLOADED_FILENAME

def process_document(
    doc: dict[str, Any],
    storage_adapter: ETLStorageAdapter | None = None,
    use_cloud: bool = False,
) -> None:
    """Download and save a single document with metadata.

    Args:
        doc: Document configuration dictionary
        storage_adapter: Optional ETL storage adapter (for cloud storage)
        use_cloud: Whether to use cloud storage (requires storage_adapter)
    """
    title = doc.get("title", "untitled")
    jurisdiction = doc.get("jurisdiction", "unknown")
    category = doc.get("category", "misc")
    topic = doc.get("topic", "general")
    urls = doc.get("download_urls", [])

    # Extract ESG metadata if present (NGER/ISSB compliance)
    esg_metadata = doc.get("esg_metadata", {})
    spatial_metadata = doc.get("spatial_metadata", {})

    for url in urls:
        # Build metadata
        metadata = {
            "title": title,
            "jurisdiction": jurisdiction,
            "category": category,
            "topic": topic,
            "source_url": doc.get("source_url", url),  # Generic source (portal)
            "source_pdf_url": url,  # Actual PDF URL for deep linking
        }

        # Add ESG metadata if present
        if esg_metadata:
            metadata["esg_metadata"] = esg_metadata

        # Add spatial metadata if present
        if spatial_metadata:
            metadata["spatial_metadata"] = spatial_metadata

        # Use cloud storage if enabled
        if use_cloud and storage_adapter:
            try:
                doc_id = storage_adapter.download_from_url(url, metadata)
                print(f"✅ Downloaded to cloud: {url} (ID: {doc_id})")
            except Exception as e:
                logger.error(f"Failed to download {url} to cloud: {e}", exc_info=True)
                print(f"❌ Failed to download: {url}")
        else:
            # Use local filesystem (backward compatibility)
            _process_document_local(doc, metadata)

def download_documents(docs: list[dict], output_dir: str) -> list[str]:
    """Download multiple documents to output directory.

    :param docs: List of document dicts with 'download_urls' key
    :param output_dir: Directory to save downloaded files
    :return: List of downloaded file paths
    """
    import os

    os.makedirs(output_dir, exist_ok=True)

    downloaded_files = []
    for doc in docs:
        title = doc.get("title", "untitled")
        urls = doc.get("download_urls", [])

        for url in urls:
            filename = safe_filename(url)
            dest_path = os.path.join(output_dir, filename)

            if download_file(url, dest_path):
                print(f"✅ Downloaded: {title} -> {filename}")
                downloaded_files.append(dest_path)
            else:
                print(f"❌ Failed: {title} from {url}")

    return downloaded_files

def ingest_documents(
    use_cloud: bool | None = None,
    config_path: str | Path | None = None,
) -> list[str]:
    """Ingest documents from config using plugin system (single source of truth).

    This function uses the document source plugin architecture to:
    - Validate configurations
    - Generate consistent document IDs (for delta indexing)
    - Create hierarchical directory structures
    - Extract metadata using source-specific logic

    Args:
        use_cloud: Whether to use cloud storage. If None, uses CLOUD_PROVIDER setting.
        config_path: Path to config file. Defaults to CONFIG_PATH.

    Returns:
        List of document IDs (cloud mode) or file paths (local mode)
    """
    from green_gov_rag.etl.sources.factory import DocumentSourceFactory

    # Determine storage mode
    if use_cloud is None:
        use_cloud = settings.cloud_provider != "local"

    # Load config
    if config_path is None:
        config_path = CONFIG_PATH

    with open(config_path, encoding="utf-8") as f:
        config = yaml.safe_load(f)

    documents = config.get("documents", [])
    print(f"Found {len(documents)} document sources in config.")
    print(f"Storage mode: {'cloud' if use_cloud else 'local'}")

    # Initialize factory and storage adapter
    factory = DocumentSourceFactory()
    storage_adapter = ETLStorageAdapter() if use_cloud else None

    # Process documents using plugin system
    document_ids = []
    for doc_config in documents:
        try:
            # 1. Create source plugin (auto-detects type)
            source = factory.create_source(doc_config)

            # 2. Validate configuration
            validation = source.validate()
            if not validation.is_valid:
                logger.error(
                    f"Invalid config for {doc_config.get('title', 'unknown')}: {validation.errors}"
                )
                print(f"❌ Validation failed: {doc_config.get('title')}")
                for error in validation.errors:
                    print(f"   - {error}")
                continue

            # 3. Log warnings if any
            if validation.warnings:
                for warning in validation.warnings:
                    logger.warning(f"{doc_config.get('title')}: {warning}")

            # 4. Get URLs from plugin
            urls = source.get_download_urls()

            # 5. Get metadata from plugin (single source of truth)
            metadata = source.get_metadata()

            # 6. Process each URL
            for url in urls:
                try:
                    # Generate document ID using plugin (consistent with monitoring)
                    doc_id = source.get_document_id(url)

                    if use_cloud:
                        # Cloud mode - upload to S3/Azure Blob
                        if storage_adapter:
                            cloud_id = storage_adapter.download_from_url(url, metadata)
                            document_ids.append(cloud_id)
                            print(f"✅ Downloaded to cloud: {url} (ID: {cloud_id})")
                    else:
                        # Local mode - use plugin-generated path
                        dest_path = source.get_destination_path(
                            url, base_dir=str(RAW_DATA_DIR)
                        )
                        dest_path_obj = Path(dest_path)

                        # Create directory
                        dest_path_obj.parent.mkdir(parents=True, exist_ok=True)

                        # Skip if exists
                        if dest_path_obj.exists():
                            logger.info(f"Skipping {url} — already exists")
                            document_ids.append(doc_id)
                            continue

                        # Download file
                        if download_file(url, str(dest_path_obj)):
                            # Detect actual file type and fix extension if needed
                            detected_ext = detect_file_type(dest_path_obj)
                            final_path = dest_path_obj

                            if detected_ext and not dest_path_obj.name.lower().endswith(
                                detected_ext
                            ):
                                # Rename file with correct extension
                                final_filename = f"{dest_path_obj.name}{detected_ext}"
                                final_path = dest_path_obj.parent / final_filename
                                dest_path_obj.rename(final_path)
                                logger.info(
                                    f"Renamed {dest_path_obj.name} → {final_filename} "
                                    f"(detected: {detected_ext})"
                                )
                            else:
                                final_filename = dest_path_obj.name

                            # Save metadata
                            metadata_with_file = metadata.copy()
                            metadata_with_file["filename"] = final_filename
                            metadata_with_file[
                                "download_timestamp"
                            ] = datetime.utcnow().isoformat()
                            metadata_with_file["sha256"] = sha256sum(final_path)
                            metadata_with_file[
                                "document_id"
                            ] = doc_id  # NEW: for delta indexing
                            metadata_with_file[
                                "source_pdf_url"
                            ] = url  # Add PDF URL for deep linking

                            metadata_path = (
                                final_path.parent
                                / f"{final_filename}{METADATA_FILE_SUFFIX}"
                            )
                            with open(metadata_path, "w", encoding="utf-8") as mf:
                                json.dump(metadata_with_file, mf, indent=2)

                            print(f"✅ Downloaded: {url} (ID: {doc_id})")
                            document_ids.append(doc_id)
                        else:
                            # Log failed download
                            failed_urls_file = LOG_DIR / FAILED_DOWNLOADS_FILENAME
                            with open(failed_urls_file, "a", encoding="utf-8") as f:
                                f.write(
                                    f"{datetime.utcnow().isoformat()} | {url} | {metadata.get('title')} | {doc_id}\n"
                                )
                            print(f"❌ Failed: {url} (logged to {failed_urls_file})")

                except Exception as e:
                    logger.error(f"Failed to process {url}: {e}", exc_info=True)
                    print(f"❌ Error processing {url}: {e}")

        except Exception as e:
            logger.error(
                f"Failed to create source for {doc_config.get('title', 'unknown')}: {e}",
                exc_info=True,
            )
            print(f"❌ Failed to process document config: {doc_config.get('title')}")

    print(f"\n✅ Ingestion complete. Processed {len(document_ids)} document(s).")
    return document_ids

def main() -> None:
    """Main entry point for CLI usage."""
    config = load_config()
    documents = config.get("documents", [])
    print(f"Found {len(documents)} document sources in config.")

    # Detect storage mode from settings
    use_cloud = settings.cloud_provider != "local"
    print(f"Storage mode: {'cloud' if use_cloud else 'local'}")

    storage_adapter = ETLStorageAdapter() if use_cloud else None

    for doc in documents:
        process_document(doc, storage_adapter=storage_adapter, use_cloud=use_cloud)

def load_documents_config(config_path: str = "configs/documents_config.yml"):
    """Load document metadata from YAML config.

    This function maintains backward compatibility with the original API
    while supporting the new plugin-based architecture.

    Args:
        config_path: Path to documents configuration file

    Returns:
        List of document configuration dictionaries
    """
    config_file = Path(config_path)
    if not config_file.exists():
        msg = f"Config file {config_path} not found."
        raise FileNotFoundError(msg)

    with open(config_file, encoding="utf-8") as f:
        config = yaml.safe_load(f)

    return config.get("documents", [])

def load_document_sources(
    config_path: str = "configs/documents_config.yml",
) -> list[DocumentSource]:
    """Load document sources using the plugin-based architecture.

    This is the new API that returns DocumentSource objects instead of
    raw configuration dictionaries.

    Args:
        config_path: Path to documents configuration file

    Returns:
        List of DocumentSource plugin instances

    Example:
        >>> sources = load_document_sources()
        >>> for source in sources:
        ...     metadata = source.get_metadata()
        ...     urls = source.get_download_urls()
        ...     validation = source.validate()
    """
    configs = load_documents_config(config_path)
    factory = DocumentSourceFactory()
    return factory.create_sources_from_list(configs)

def get_document_sources():
    """Returns a list of all source URLs for ingestion.

    This function maintains backward compatibility with the original API.

    Returns:
        List of download URLs
    """
    documents = load_documents_config()
    sources = []
    for doc in documents:
        urls = doc.get("download_urls", [])
        sources.extend(urls)
    return sources

def get_document_sources_by_type(source_type: str) -> list[DocumentSource]:
    """Get document sources filtered by type.

    Args:
        source_type: Source type identifier (e.g., 'federal_legislation')

    Returns:
        List of DocumentSource instances matching the type

    Example:
        >>> federal_sources = get_document_sources_by_type('federal_legislation')
        >>> emissions_sources = get_document_sources_by_type('emissions_reporting')
    """
    sources = load_document_sources()
    return [s for s in sources if s.get_source_type() == source_type]

def get_document_sources_by_jurisdiction(jurisdiction: str) -> list[DocumentSource]:
    """Get document sources filtered by jurisdiction.

    Args:
        jurisdiction: Jurisdiction filter ('federal', 'state', 'local')

    Returns:
        List of DocumentSource instances matching the jurisdiction

    Example:
        >>> federal_sources = get_document_sources_by_jurisdiction('federal')
        >>> local_sources = get_document_sources_by_jurisdiction('local')
    """
    sources = load_document_sources()
    return [s for s in sources if s.get_metadata().get("jurisdiction") == jurisdiction]