pagepiper/scripts/shelve_docx.py

# scripts/shelve_docx.py
"""
cf-orch task: pagepiper/shelve_docx

Extracts text from a Word .docx file, stores section chunks in SQLite, and
(if Ollama is configured) generates embeddings in the sqlite-vec store.

Chunking strategy:
  - If the document has >=2 Heading-style paragraphs: split at each heading
    (one chunk per section, heading text included).
  - Otherwise: accumulate blocks into ~WORDS_PER_CHUNK rolling windows.

Tables are serialised as pipe-delimited rows and included in the surrounding
section chunk, preserving document order via XML tree traversal.

Entry point:
    python scripts/shelve_docx.py --doc-id X --file-path Y --db-path Z --vec-db-path W
"""
from __future__ import annotations

import logging
import os
import sqlite3
from dataclasses import dataclass, field
from pathlib import Path

logger = logging.getLogger("pagepiper.shelve_docx")

EMBED_BATCH_SIZE = 64
_WORDS_PER_CHUNK = 500


@dataclass
class _Chunk:
    page_number: int
    text: str
    source: str
    word_count: int


def _table_to_text(table) -> str:
    """Serialise a DOCX table as pipe-delimited rows."""
    lines = []
    for row in table.rows:
        cells = [c.text.strip().replace("\n", " ") for c in row.cells]
        if any(cells):
            lines.append(" | ".join(cells))
    return "\n".join(lines)


def _iter_blocks(doc):
    """
    Yield (kind, obj) pairs in document body order, where kind is
    'paragraph' or 'table'.  Walks the raw XML so that tables and
    paragraphs appear in the correct interleaved sequence.
    """
    import docx.text.paragraph as _p_mod
    import docx.table as _t_mod
    from docx.oxml.ns import qn

    for child in doc.element.body.iterchildren():
        if child.tag == qn("w:p"):
            yield "paragraph", _p_mod.Paragraph(child, doc)
        elif child.tag == qn("w:tbl"):
            yield "table", _t_mod.Table(child, doc)


def _is_heading(para) -> bool:
    return para.style.name.startswith("Heading")


def _extract_chunks(file_path: str) -> list[_Chunk]:
    import docx
    from scripts.text_clean import clean_line, is_artifact_line

    doc = docx.Document(file_path)

    # Count headings to decide strategy
    heading_count = sum(1 for p in doc.paragraphs if _is_heading(p))

    blocks: list[tuple[str, object]] = list(_iter_blocks(doc))

    if heading_count >= 2:
        return _heading_chunks(blocks)
    else:
        return _wordcount_chunks(blocks)


def _heading_chunks(blocks: list[tuple[str, object]]) -> list[_Chunk]:
    """One chunk per heading section; tables included inline."""
    from scripts.text_clean import clean_line, is_artifact_line

    chunks: list[_Chunk] = []
    current_parts: list[str] = []

    def _flush(parts: list[str]) -> None:
        text = "\n".join(parts).strip()
        if text:
            n = len(chunks) + 1
            chunks.append(_Chunk(n, text, "section", len(text.split())))

    for kind, obj in blocks:
        if kind == "paragraph":
            if _is_heading(obj):
                _flush(current_parts)
                current_parts = []
                t = obj.text.strip()
                if t:
                    current_parts.append(t)
            else:
                t = clean_line(obj.text.strip())
                if t and not is_artifact_line(t):
                    current_parts.append(t)
        elif kind == "table":
            table_text = _table_to_text(obj)
            if table_text:
                current_parts.append(table_text)

    _flush(current_parts)
    return chunks


def _wordcount_chunks(blocks: list[tuple[str, object]]) -> list[_Chunk]:
    """Accumulate blocks into ~WORDS_PER_CHUNK rolling windows."""
    from scripts.text_clean import clean_line, is_artifact_line

    chunks: list[_Chunk] = []
    current: list[str] = []
    current_count = 0

    def _flush(parts: list[str]) -> None:
        text = "\n".join(parts).strip()
        if text:
            n = len(chunks) + 1
            chunks.append(_Chunk(n, text, "text", len(text.split())))

    for kind, obj in blocks:
        if kind == "paragraph":
            t = clean_line(obj.text.strip())
            if not t or is_artifact_line(t):
                continue
        else:  # table
            t = _table_to_text(obj)
            if not t:
                continue

        words = t.split()
        if current_count + len(words) > _WORDS_PER_CHUNK and current:
            _flush(current)
            current, current_count = [], 0
        current.append(t)
        current_count += len(words)

    if current:
        _flush(current)

    return chunks


def _update_status(
    conn: sqlite3.Connection,
    doc_id: str,
    status: str,
    page_count: int | None = None,
    error_msg: str | None = None,
) -> None:
    if page_count is not None:
        conn.execute(
            "UPDATE documents SET status=?, page_count=?, updated_at=datetime('now') WHERE id=?",
            [status, page_count, doc_id],
        )
    elif error_msg is not None:
        conn.execute(
            "UPDATE documents SET status=?, error_msg=?, updated_at=datetime('now') WHERE id=?",
            [status, error_msg, doc_id],
        )
    else:
        conn.execute(
            "UPDATE documents SET status=?, updated_at=datetime('now') WHERE id=?",
            [status, doc_id],
        )
    conn.commit()


def run(doc_id: str, file_path: str, db_path: str, vec_db_path: str) -> None:
    """Run the full shelve pipeline for one DOCX. Called by cf-orch or BackgroundTasks."""
    conn: sqlite3.Connection | None = None
    try:
        conn = sqlite3.connect(db_path, timeout=30)
        conn.execute("PRAGMA journal_mode = WAL")
        conn.execute("PRAGMA foreign_keys = ON")
        _update_status(conn, doc_id, "processing")

        logger.info("Extracting sections from %s", file_path)
        chunks = _extract_chunks(file_path)
        logger.info("Extracted %d chunks", len(chunks))

        from scripts.text_clean import clean_paragraph
        conn.execute("DELETE FROM page_chunks WHERE doc_id=?", [doc_id])
        chunk_rows: list[tuple[str, int, str]] = []
        for chunk in chunks:
            cleaned = clean_paragraph(chunk.text)
            if not cleaned:
                continue
            row = conn.execute(
                """INSERT INTO page_chunks(doc_id, page_number, text, source, word_count)
                   VALUES (?,?,?,?,?) RETURNING id""",
                [doc_id, chunk.page_number, cleaned, chunk.source, len(cleaned.split())],
            ).fetchone()
            chunk_rows.append((row[0], chunk.page_number, cleaned))
        conn.commit()

        from app.config import get_llm_config
        llm_cfg = get_llm_config()
        if llm_cfg and chunks:
            try:
                logger.info("Embedding %d chunks", len(chunks))
                from circuitforge_core.llm import LLMRouter
                from circuitforge_core.vector.sqlite_vec import LocalSQLiteVecStore

                router = LLMRouter(llm_cfg)
                embed_dims = int(os.environ.get("PAGEPIPER_EMBED_DIMS", "1024"))
                vec_store = LocalSQLiteVecStore(
                    db_path=vec_db_path, table="page_vecs", dimensions=embed_dims
                )
                vec_store.delete_where({"doc_id": doc_id})

                texts = [text for _, _, text in chunk_rows]
                vectors: list[list[float]] = []
                for i in range(0, len(texts), EMBED_BATCH_SIZE):
                    vectors.extend(router.embed(texts[i : i + EMBED_BATCH_SIZE]))

                for (chunk_id, page_number, _), vector in zip(chunk_rows, vectors):
                    vec_store.upsert(
                        entry_id=chunk_id,
                        vector=vector,
                        metadata={"doc_id": doc_id, "page_number": page_number},
                    )
                logger.info("Stored %d embeddings", len(vectors))
            except Exception as embed_exc:
                logger.warning(
                    "Embedding skipped for doc %s — BM25 only (reason: %s)",
                    doc_id, embed_exc,
                )

        _update_status(conn, doc_id, "ready", page_count=len(chunks))
        logger.info("Shelve complete for doc %s (%d chunks)", doc_id, len(chunks))

    except Exception as exc:
        logger.error("Shelve failed for doc %s: %s", doc_id, exc, exc_info=True)
        if conn is not None:
            try:
                _update_status(conn, doc_id, "error", error_msg=str(exc))
            except Exception:
                logger.warning("Could not write error status for doc %s", doc_id)
        raise
    finally:
        if conn is not None:
            conn.close()


if __name__ == "__main__":
    import argparse

    logging.basicConfig(level=logging.INFO)

    parser = argparse.ArgumentParser(
        description="Shelve a Word .docx (cf-orch task entry point)"
    )
    parser.add_argument("--doc-id", required=True)
    parser.add_argument("--file-path", required=True)
    parser.add_argument("--db-path", required=True)
    parser.add_argument("--vec-db-path", required=True)
    a = parser.parse_args()
    run(
        doc_id=a.doc_id,
        file_path=a.file_path,
        db_path=a.db_path,
        vec_db_path=a.vec_db_path,
    )