src/rag/embedder.py

"""Generate embeddings for cascade chains."""

from __future__ import annotations

from datetime import date

from src.llm.client import load_config
from src.models.schemas import CascadeChain

VALID_ENCODING_MODES = {"trigger_only", "full_chain"}
DEFAULT_ENCODING_MODE = "trigger_only"


def trigger_to_embedding_text(
    country: str,
    event_date: date | str,
    severity: str | None,
    summary: str,
    location: str | None = None,
) -> str:
    """Build the symmetric trigger-side text used on both index and query sides.

    Keeping one constructor prevents format drift between the document and the
    query — the similarity score is only meaningful if both sides share an
    identical template.
    """
    loc = f"{location}, {country}" if location else country
    sev = severity or "unknown"
    return f"Flood in {loc}. Date: {event_date}. Severity: {sev}. {summary}"


def chain_to_embedding_text(chain: CascadeChain) -> str:
    """Full-chain text: trigger + every cascade node + cross-domain transitions.

    Used for ``full_chain`` encoding mode, and always used as the ChromaDB
    ``document`` payload regardless of mode so that the stored record stays
    human-inspectable.
    """
    parts = [
        trigger_to_embedding_text(
            chain.trigger_country,
            chain.trigger_date,
            chain.trigger_severity,
            chain.trigger_summary,
        )
    ]

    if chain.cascade_events:
        domains = sorted({n.domain for n in chain.cascade_events})
        parts.append(f"Affected domains: {', '.join(domains)}.")

        for node in chain.cascade_events:
            desc = f"{node.domain}: {node.description} (severity={node.severity}"
            if node.time_offset_hours is not None:
                desc += f", +{node.time_offset_hours}h"
            desc += ")"
            parts.append(desc)

        cross_domain = []
        for node in chain.cascade_events:
            for pid in node.parent_ids:
                parent = next((n for n in chain.cascade_events if n.id == pid), None)
                if parent and parent.domain != node.domain:
                    cross_domain.append(f"{parent.domain} → {node.domain}")
        if cross_domain:
            parts.append(f"Cross-domain cascades: {'; '.join(set(cross_domain))}.")

    return " ".join(parts)


class Embedder:
    """Generate embeddings using sentence-transformers (default) or other backends."""

    def __init__(self, config: dict | None = None):
        if config is None:
            config = load_config()
        emb_cfg = config["embedding"]
        self.backend = emb_cfg.get("backend", "sentence-transformers")
        self.model_name = emb_cfg.get("model", "all-MiniLM-L6-v2")
        mode = emb_cfg.get("encoding_mode", DEFAULT_ENCODING_MODE)
        if mode not in VALID_ENCODING_MODES:
            raise ValueError(
                f"Invalid embedding.encoding_mode: {mode!r}. "
                f"Expected one of {sorted(VALID_ENCODING_MODES)}."
            )
        self.encoding_mode = mode
        self._model = None

    def _load_model(self):
        if self._model is None:
            if self.backend == "sentence-transformers":
                import torch
                from sentence_transformers import SentenceTransformer
                device = "cuda" if torch.cuda.is_available() else "cpu"
                self._model = SentenceTransformer(self.model_name, device=device)
            else:
                raise ValueError(f"Unknown embedding backend: {self.backend}")

    def embed_text(self, text: str) -> list[float]:
        """Embed a single text string."""
        self._load_model()
        embedding = self._model.encode(text, normalize_embeddings=True)
        return embedding.tolist()

    def embed_chain(self, chain: CascadeChain) -> list[float]:
        """Embed a cascade chain according to the configured encoding mode."""
        if self.encoding_mode == "trigger_only":
            text = trigger_to_embedding_text(
                chain.trigger_country,
                chain.trigger_date,
                chain.trigger_severity,
                chain.trigger_summary,
            )
        else:
            text = chain_to_embedding_text(chain)
        return self.embed_text(text)

    def embed_texts(self, texts: list[str]) -> list[list[float]]:
        """Embed multiple texts at once (more efficient for batches)."""
        self._load_model()
        embeddings = self._model.encode(texts, normalize_embeddings=True)
        return embeddings.tolist()