ingestion.py

from pathlib import Path
from typing import Any

from dotenv import load_dotenv

load_dotenv()

from chroma import get_collection  # noqa: E402 — after load_dotenv

DOCUMENTS_DIR = Path(__file__).parent / "documents"


def parse_document(file_path: str | Path) -> list[dict[str, Any]]:
    """Return list of {text, page_number} dicts from a .pdf, .docx, or .md file."""
    file_path = Path(file_path)
    suffix = file_path.suffix.lower()

    if suffix == ".pdf":
        from pypdf import PdfReader

        reader = PdfReader(str(file_path))
        pages = []
        for i, page in enumerate(reader.pages):
            text = page.extract_text() or ""
            if text.strip():
                pages.append({"text": text, "page_number": i + 1})
        return pages

    if suffix == ".docx":
        from docx import Document

        doc = Document(str(file_path))
        full_text = "\n".join(p.text for p in doc.paragraphs if p.text.strip())
        return [{"text": full_text, "page_number": None}]

    if suffix == ".md":
        text = file_path.read_text(encoding="utf-8")
        return [{"text": text, "page_number": None}]

    raise ValueError(f"Unsupported file type: {suffix!r}. Supported: .pdf, .docx, .md")


def chunk_text(
    text: str,
    chunk_size: int = 500,
    overlap: int = 50,
) -> list[str]:
    """Split text into overlapping word-based chunks."""
    words = text.split()
    if not words:
        return []

    chunks: list[str] = []
    start = 0
    while start < len(words):
        end = min(start + chunk_size, len(words))
        chunks.append(" ".join(words[start:end]))
        if end == len(words):
            break
        start += chunk_size - overlap

    return chunks


def embed_and_store(
    chunks: list[str],
    metadata: dict[str, Any],
) -> None:
    """Embed chunks and upsert them into ChromaDB with per-chunk metadata."""
    if not chunks:
        return

    collection = get_collection()

    source_file = metadata.get("source_file", "unknown")
    page_number = metadata.get("page_number")  # may be None for docx

    ids = [
        f"{source_file}__p{page_number}__c{i}" if page_number is not None else f"{source_file}__c{i}"
        for i in range(len(chunks))
    ]

    chunk_metadata = [
        {
            "source_file": source_file,
            "chunk_index": i,
            **({"page_number": page_number} if page_number is not None else {}),
        }
        for i in range(len(chunks))
    ]

    collection.upsert(documents=chunks, metadatas=chunk_metadata, ids=ids)


def ingest_all_documents(extra_dirs: list[Path] | None = None) -> dict[str, Any]:
    """Scan documents/ folder (and any extra_dirs) for supported files and index them."""
    supported = {".pdf", ".docx", ".md"}
    search_dirs = [DOCUMENTS_DIR] + (extra_dirs or [])
    files: list[Path] = []
    for d in search_dirs:
        if d.exists() and d.is_dir():
            files.extend(f for f in d.iterdir() if f.suffix.lower() in supported)

    if not files:
        print("No documents found in", DOCUMENTS_DIR)
        return {"ingested": 0, "files": []}

    results: list[str] = []
    total_chunks = 0

    for file_path in files:
        print(f"Ingesting: {file_path.name}")
        try:
            pages = parse_document(file_path)
            file_chunks = 0

            for page in pages:
                chunks = chunk_text(page["text"])
                embed_and_store(
                    chunks,
                    metadata={
                        "source_file": file_path.name,
                        "page_number": page["page_number"],
                    },
                )
                file_chunks += len(chunks)

            total_chunks += file_chunks
            results.append(file_path.name)
            print(f"  -> {file_chunks} chunks stored")

        except Exception as e:
            print(f"  -> ERROR: {e}")

    print(f"\nDone. {len(results)} file(s), {total_chunks} total chunks.")
    return {"ingested": len(results), "total_chunks": total_chunks, "files": results}


def ingest_scraped_content(data: dict) -> None:
    """Chunk and upsert scraped website content into ChromaDB."""
    parts = []
    if data.get("title"):
        parts.append(data["title"])
    if data.get("description"):
        parts.append(data["description"])
    if data.get("headings"):
        parts.append("\n".join(data["headings"]))
    if data.get("body_text"):
        parts.append(data["body_text"])

    full_text = "\n\n".join(parts)
    chunks = chunk_text(full_text)
    if not chunks:
        return

    url = data.get("url", "scraped")
    scraped_at = data.get("scraped_at", "")
    embed_and_store(
        chunks,
        metadata={"source_file": url, "scraped_at": scraped_at},
    )


if __name__ == "__main__":
    ingest_all_documents()