plenum/scripts/make_embeddings.py

"""
Chunks and embeds all talks that are not yet represented in the 'chunks' table.

Replaces scripts/make_arango_embeddings.py.

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

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

Usage:
  python scripts/make_embeddings.py
"""

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

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

os.chdir("/home/lasse/riksdagen")
sys.path.append("/home/lasse/riksdagen")

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 = 384
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
        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(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 chunks to embed.")
        return 0

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

    def _rows(chunks: List[Dict]) -> List[tuple]:
        return [
            (c["id"], c[parent_key], c["chunk_index"], c["text"], c["embedding"])
            for c in 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"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} chunks into PostgreSQL.")
    return total_inserted


def make_embeddings() -> int:
    """
    Find talks with no chunks and generate + insert embeddings.
    Returns the total number of chunk rows inserted.
    """
    # Find talks that have no rows in the 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.anforandetext
        FROM talks t
        LEFT JOIN chunks c ON c.talk_id = t.id
        WHERE c.id IS NULL
          AND t.anforandetext IS NOT NULL
          AND t.anforandetext != ''
        """
    )
    logger.info(f"Found {len(missing)} talks without chunks")

    all_batches = _chunk_docs(missing, text_key="anforandetext", parent_key="talk_id")

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


def make_motion_embeddings(year: int | None = None) -> int:
    """
    Find motions with no 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.dok_id AS id, m.text
        FROM motions m
        LEFT JOIN motion_chunks c ON c.motion_id = m.dok_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)} motions without chunks")

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

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


def make_yrkande_embeddings(limit: int | None = None) -> int:
    """
    Embed motion_yrkanden rows that have no embedding yet. Each yrkande (lydelse)
    is a short, self-contained proposal, so it is embedded whole (no chunking).
    Rows already exist (populated from the forslag JSONB), so this UPDATEs them.
    Returns the number of yrkanden embedded.
    """
    sql = (
        "SELECT id, lydelse FROM motion_yrkanden "
        "WHERE embedding IS NULL AND lydelse IS NOT NULL AND lydelse != ''"
    )
    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["lydelse"]} 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 motion_yrkanden 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 == "motions":
        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()