plenum/scripts/make_embeddings.py

"""
Chunks and embeds all speeches that are not yet represented in the 'speech_chunks' table.


Pipeline:
  1. Find all speeches with no chunk rows in PostgreSQL via LEFT JOIN.
  2. Split each talk's text into speech_chunks (max 500 chars) using TextChunker.
  3. Generate embeddings via Ollama in parallel (3 workers, batches of 20).
  4. Insert chunk rows into the 'speech_chunks' table with id = "{speech_id}:{chunk_index}".

The search_vector column on speeches is kept in sync by a trigger – no manual update needed.

Usage:
  python scripts/make_embeddings.py
"""
from pathlib import Path

import logging
import os
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Dict, List

logging.getLogger("httpx").setLevel(logging.WARNING)

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import bootstrap  # noqa: E402,F401  — sets cwd and sys.path to the project root

from parliament import PARLIAMENT
from postgres_client import pg
from utils import TextChunker

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)
logger = logging.getLogger(__name__)

EMBED_DIM = PARLIAMENT.embeddings.dimension
EMBED_BATCH = 20
MAX_WORKERS = 3
INSERT_BATCH = 100


def _make_embeddings(texts: List[str]) -> List[List[float]]:
    # OpenAI-compatible vLLM endpoint (VLLM_EMBEDDING_HOST/LLM_MODEL_EMBEDDING),
    # same as query-time embeddings — keeps documents and queries in one space.
    return pg.make_embeddings(texts)


def _embed_batch(chunk_batch: List[Dict]) -> List[Dict]:
    """Embed a batch of chunk dicts (adds 'embedding' field)."""
    texts = [c["text"] for c in chunk_batch]
    embeddings = _make_embeddings(texts)
    for i, chunk in enumerate(chunk_batch):
        chunk["embedding"] = embeddings[i]
    return chunk_batch


def _chunk_docs(missing: List[Dict], text_key: str, parent_key: str) -> List[List[Dict]]:
    """Split docs into chunk dicts and group them into embed batches."""
    all_batches: List[List[Dict]] = []
    for doc in missing:
        doc_id = doc["id"]
        text = (doc.get(text_key) or "").strip()
        if not text:
            continue
        speech_chunks = TextChunker(chunk_limit=500).chunk(text)
        _chunks = [
            {
                "id": f"{doc_id}:{idx}",
                parent_key: doc_id,
                "chunk_index": idx,
                "text": content,
            }
            for idx, content in enumerate(speech_chunks)
            if content and content.strip()
        ]
        for i in range(0, len(_chunks), EMBED_BATCH):
            batch = _chunks[i : i + EMBED_BATCH]
            if batch:
                all_batches.append(batch)
    return all_batches


def _embed_and_insert(all_batches: List[List[Dict]], insert_sql: str, parent_key: str) -> int:
    """Embed all batches in parallel and bulk-insert the chunk rows."""
    if not all_batches:
        logger.info("No speech_chunks to embed.")
        return 0

    logger.info(f"Embedding {len(all_batches)} batches via Ollama …")

    def _rows(speech_chunks: List[Dict]) -> List[tuple]:
        return [
            (c["id"], c[parent_key], c["chunk_index"], c["text"], c["embedding"])
            for c in speech_chunks
        ]

    total_inserted = 0
    pending: List[Dict] = []
    completed = 0

    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
        futures = [executor.submit(_embed_batch, batch) for batch in all_batches]
        for future in as_completed(futures):
            result = future.result()
            completed += 1
            pending.extend(result)
            print(
                f"  batches embedded: {completed}/{len(all_batches)} | "
                f"speech_chunks pending insert: {len(pending)}",
                end="\r",
            )
            if len(pending) >= INSERT_BATCH:
                pg.execute_values(insert_sql, _rows(pending))
                total_inserted += len(pending)
                pending = []

    if pending:
        pg.execute_values(insert_sql, _rows(pending))
        total_inserted += len(pending)

    print()
    logger.info(f"Done. Inserted {total_inserted} speech_chunks into PostgreSQL.")
    return total_inserted


def make_embeddings() -> int:
    """
    Find speeches with no speech_chunks and generate + insert embeddings.
    Returns the total number of chunk rows inserted.
    """
    # Find speeches that have no rows in the speech_chunks table.
    # SET LOCAL disables parallel hash join for this query — it otherwise spills
    # into the Postgres container's small /dev/shm and fails with DiskFull.
    missing = pg.execute(
        """
        SET LOCAL max_parallel_workers_per_gather = 0;
        SELECT t.id, t.text
        FROM speeches t
        LEFT JOIN speech_chunks c ON c.speech_id = t.id
        WHERE c.id IS NULL
          AND t.text IS NOT NULL
          AND t.text != ''
        """
    )
    logger.info(f"Found {len(missing)} speeches without speech_chunks")

    all_batches = _chunk_docs(missing, text_key="text", parent_key="speech_id")

    INSERT_SQL = """
    INSERT INTO speech_chunks (id, speech_id, chunk_index, text, embedding)
    VALUES %s
    ON CONFLICT (id) DO NOTHING
    """
    return _embed_and_insert(all_batches, INSERT_SQL, parent_key="speech_id")


def make_motion_embeddings(year: int | None = None) -> int:
    """
    Find documents with no speech_chunks and generate + insert embeddings.
    Optional year filter for running the backfill in slices.
    Returns the total number of chunk rows inserted.
    """
    sql = """
        SET LOCAL max_parallel_workers_per_gather = 0;
        SELECT m.doc_id AS id, m.text
        FROM documents m
        LEFT JOIN document_chunks c ON c.doc_id = m.doc_id
        WHERE c.id IS NULL
          AND m.has_text
        """
    params = None
    if year is not None:
        sql += " AND m.year = %s"
        params = (year,)
    missing = pg.execute(sql, params)
    logger.info(f"Found {len(missing)} documents without speech_chunks")

    all_batches = _chunk_docs(missing, text_key="text", parent_key="doc_id")

    INSERT_SQL = """
    INSERT INTO document_chunks (id, doc_id, chunk_index, text, embedding)
    VALUES %s
    ON CONFLICT (id) DO NOTHING
    """
    return _embed_and_insert(all_batches, INSERT_SQL, parent_key="doc_id")


def make_yrkande_embeddings(limit: int | None = None) -> int:
    """
    Embed document_proposals rows that have no embedding yet. Each yrkande (text)
    is a short, self-contained proposal, so it is embedded whole (no chunking).
    Rows already exist (populated from the proposals_raw JSONB), so this UPDATEs them.
    Returns the number of yrkanden embedded.
    """
    sql = (
        "SELECT id, text FROM document_proposals "
        "WHERE embedding IS NULL AND text IS NOT NULL AND text != ''"
    )
    if limit is not None:
        sql += f" LIMIT {int(limit)}"
    missing = pg.execute(sql)
    logger.info(f"Found {len(missing)} yrkanden without embeddings")

    items = [{"id": r["id"], "text": r["text"]} for r in missing]
    batches = [items[i : i + EMBED_BATCH] for i in range(0, len(items), EMBED_BATCH)]
    if not batches:
        logger.info("No yrkanden to embed.")
        return 0

    logger.info(f"Embedding {len(batches)} yrkande batches via vLLM …")
    UPDATE_SQL = "UPDATE document_proposals SET embedding = %s WHERE id = %s"

    total = 0
    pending: List[Dict] = []
    completed = 0
    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
        futures = [executor.submit(_embed_batch, batch) for batch in batches]
        for future in as_completed(futures):
            pending.extend(future.result())
            completed += 1
            print(f"  batches embedded: {completed}/{len(batches)} | pending update: {len(pending)}", end="\r")
            if len(pending) >= INSERT_BATCH:
                pg.execute_many(UPDATE_SQL, [(c["embedding"], c["id"]) for c in pending])
                total += len(pending)
                pending = []
    if pending:
        pg.execute_many(UPDATE_SQL, [(c["embedding"], c["id"]) for c in pending])
        total += len(pending)

    print()
    logger.info(f"Done. Embedded {total} yrkanden.")
    return total


if __name__ == "__main__":
    mode = sys.argv[1] if len(sys.argv) > 1 else None
    if mode == "documents":
        year = int(sys.argv[2]) if len(sys.argv) > 2 else None
        make_motion_embeddings(year=year)
    elif mode == "yrkanden":
        make_yrkande_embeddings()
    else:
        make_embeddings()