pagepiper/scripts/shelve_odt.py

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

Extracts text from an OpenDocument Text (.odt) 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 paragraphs (<text:h>): 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. odfpy already yields body children in document order, so no
raw XML tree-walk is needed (unlike the DOCX shelver).

Entry point:
    python scripts/shelve_odt.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
from pathlib import Path

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

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 an ODT table as pipe-delimited rows."""
    from odf.table import TableCell, TableRow
    from odf import teletype

    lines = []
    for row in table.getElementsByType(TableRow):
        cells = [teletype.extractText(c).strip().replace("\n", " ") for c in row.getElementsByType(TableCell)]
        if any(cells):
            lines.append(" | ".join(cells))
    return "\n".join(lines)


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

    doc = load(file_path)
    blocks = list(doc.text.childNodes)

    heading_count = sum(1 for b in blocks if b.qname[1] == "h")

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


def _heading_chunks(blocks: list) -> list[_Chunk]:
    """One chunk per heading section; tables included inline."""
    from odf import teletype
    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 block in blocks:
        kind = block.qname[1]
        if kind == "h":
            _flush(current_parts)
            current_parts = []
            t = teletype.extractText(block).strip()
            if t:
                current_parts.append(t)
        elif kind == "p":
            t = clean_line(teletype.extractText(block).strip())
            if t and not is_artifact_line(t):
                current_parts.append(t)
        elif kind == "table":
            table_text = _table_to_text(block)
            if table_text:
                current_parts.append(table_text)

    _flush(current_parts)
    return chunks


def _wordcount_chunks(blocks: list) -> list[_Chunk]:
    """Accumulate blocks into ~WORDS_PER_CHUNK rolling windows."""
    from odf import teletype
    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 block in blocks:
        kind = block.qname[1]
        if kind in ("p", "h"):
            t = clean_line(teletype.extractText(block).strip())
            if not t or is_artifact_line(t):
                continue
        elif kind == "table":
            t = _table_to_text(block)
            if not t:
                continue
        else:
            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 ODT. 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 an OpenDocument Text .odt (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,
    )