politics-tracking/backend/politrack/enrich/embed.py

"""Chunking + embeddings for retrieval (RAG).

Every claim gets one chunk (its summary plus evidence titles) so agents can
retrieve facts *with citations*; documents with licensed full text are split
into ~1,200 character chunks. Embeddings are computed by an Ollama server
(model bge-m3, multilingual, 1024 dims) if POLITRACK_OLLAMA_URL is set;
otherwise only German full-text search is available.
"""

import httpx
from psycopg import Connection

from politrack.config import get_settings
from politrack.core import Enricher, RunContext, register
from politrack.core.spec import DERIVED

CHUNK_SIZE = 1200
UPSERT_CHUNK = """
INSERT INTO chunk (claim_id, document_id, chunk_index, text, lang) VALUES (%s, %s, %s, %s, 'de')
ON CONFLICT (COALESCE(document_id, 0), COALESCE(claim_id, 0), chunk_index)
DO UPDATE SET text = EXCLUDED.text, embedding = CASE WHEN chunk.text = EXCLUDED.text THEN chunk.embedding END
"""


@register
class Chunker(Enricher):
    spec = DERIVED
    name = "enrich.embed"
    description = "Build search chunks for claims/documents and embed them (Ollama, optional)"
    schedule = "*/15 * * * *"

    def run(self, ctx: RunContext) -> None:
        while (n := build_chunks(ctx.conn)):
            ctx.stats["chunks"] += n
        ctx.stats["embedded"] = embed_pending(ctx.conn)
        return None


def split_text(text: str, size: int = CHUNK_SIZE) -> list[str]:
    paragraphs, chunks, current = [p.strip() for p in text.split("\n") if p.strip()], [], ""
    for p in paragraphs:
        if current and len(current) + len(p) > size:
            chunks.append(current)
            current = ""
        current = f"{current}\n{p}".strip()
        while len(current) > size * 1.5:
            chunks.append(current[:size])
            current = current[size:]
    if current:
        chunks.append(current)
    return chunks


def build_chunks(conn: Connection, batch: int = 5000) -> int:
    claims = conn.execute(
        """
        SELECT c.id, c.claim_type, c.summary_de, c.occurred_at, p.full_name, o.name AS org_name,
               string_agg(DISTINCT s.name || ': ' || COALESCE(d.title, d.url), ' | ') AS sources
        FROM claim c
        LEFT JOIN person p ON p.id = c.person_id
        LEFT JOIN organization o ON o.id = c.organization_id
        JOIN claim_evidence e ON e.claim_id = c.id
        JOIN document d ON d.id = e.document_id
        JOIN source s ON s.id = d.source_id
        WHERE c.merged_into IS NULL AND c.review_status = 'published'
          AND NOT EXISTS (SELECT 1 FROM chunk ch WHERE ch.claim_id = c.id)
        GROUP BY c.id, p.full_name, o.name
        LIMIT %s
        """,
        (batch,),
    ).fetchall()
    for c in claims:
        who = c["full_name"] or c["org_name"] or ""
        when = c["occurred_at"].date().isoformat() if c["occurred_at"] else ""
        text = f"[{c['claim_type']}] {when} {who}\n{c['summary_de']}\nQuellen: {c['sources']}"
        conn.execute(UPSERT_CHUNK, (c["id"], None, 0, text.strip()))

    docs = conn.execute(
        """
        SELECT d.id, d.title, d.full_text FROM document d
        WHERE d.full_text IS NOT NULL AND NOT EXISTS (SELECT 1 FROM chunk ch WHERE ch.document_id = d.id)
        LIMIT %s
        """,
        (batch // 10,),
    ).fetchall()
    for d in docs:
        for i, part in enumerate(split_text(d["full_text"])):
            conn.execute(UPSERT_CHUNK, (None, d["id"], i, f"{d['title'] or ''}\n{part}".strip()))
    conn.commit()
    return len(claims) + len(docs)


def embed_texts(texts: list[str]) -> list[list[float]]:
    settings = get_settings()
    resp = httpx.post(f"{settings.ollama_url}/api/embed", json={"model": settings.embedding_model, "input": texts},
                      timeout=300)
    resp.raise_for_status()
    return resp.json()["embeddings"]


def embed_pending(conn: Connection, batch: int = 64, max_batches: int = 50) -> int:
    settings = get_settings()
    if not settings.ollama_url:
        return 0
    total = 0
    for _ in range(max_batches):
        rows = conn.execute("SELECT id, text FROM chunk WHERE embedding IS NULL ORDER BY id LIMIT %s", (batch,)).fetchall()
        if not rows:
            break
        vectors = embed_texts([r["text"] for r in rows])
        with conn.cursor() as cur:
            cur.executemany(
                "UPDATE chunk SET embedding = %s::vector, embedding_model = %s WHERE id = %s",
                [(str(v), settings.embedding_model, r["id"]) for r, v in zip(rows, vectors)],
            )
        conn.commit()
        total += len(rows)
    return total