src/segmentation/embedder.py

"""
Article embedding generator — T1.6 (Người B)

Calls embedding service provided by Người A (T6.2).
Writes embeddings back to Neo4j Article nodes.

Interface contract with Người A (T6.2)
---------------------------------------
API format:
  POST {EMBED_SERVICE_URL}/embed
  Body: {"texts": ["text1", "text2", ...]}
  Response: {"embeddings": [[float, ...], ...]}  # 1024-dim each

Interface contract with cross_reference / application layer (Người C)
----------------------------------------------------------------------
After T1.6 completes:
  - Article.embedding property exists (1024-dim float array)
  - Neo4j vector index "article_embeddings" is created and populated
  - Người C queries this index via:
      CALL db.index.vector.queryNodes("article_embeddings", 20, $query_vector)
"""
from __future__ import annotations

import logging
import os
import time
from typing import TYPE_CHECKING, Optional

if TYPE_CHECKING:
    from neo4j import Driver

logger = logging.getLogger(__name__)

# Batch sizes from spec
EMBED_BATCH_SIZE = 512    # articles per embedding API call
NEO4J_BATCH_SIZE = 1_000  # articles per Neo4j write transaction

EMBED_DIM = 1024          # harrier-0.6b (updated) output dimension
VECTOR_INDEX_NAME = "article_embeddings"  # used by Người C's queries

# Retry configuration for embedding service
EMBED_MAX_RETRIES = 3
EMBED_RETRY_DELAY = 2.0  # seconds


class ArticleEmbedder:
    """
    Generates embeddings for all Article nodes and stores them in Neo4j.

    Usage
    -----
        embedder = ArticleEmbedder(
            driver=neo4j_driver,
            embed_service_url=os.getenv("EMBED_SERVICE_URL"),
        )
        stats = embedder.embed_all()
        # {"total": N, "embedded": N, "errors": N}
    """

    def __init__(
        self,
        driver: "Driver",
        embed_service_url: Optional[str] = None,
        *,
        embed_batch_size: int = EMBED_BATCH_SIZE,
        neo4j_batch_size: int = NEO4J_BATCH_SIZE,
    ) -> None:
        self._driver = driver
        self._url = embed_service_url or os.getenv("EMBED_SERVICE_URL", "http://localhost:8001")
        self._embed_batch = embed_batch_size
        self._neo4j_batch = neo4j_batch_size

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def embed_all(self, *, overwrite: bool = False) -> dict[str, int]:
        """
        Embed all Article nodes that don't yet have an embedding.

        Parameters
        ----------
        overwrite : bool
            If True, re-embed all articles (even those with existing embeddings).
            Default False for idempotent re-runs.

        Returns
        -------
        dict with keys: total, embedded, errors
        """

        stats = {"total": 0, "embedded": 0, "errors": 0}
        
        where_clause = "WHERE a.embedding IS NULL" if not overwrite else ""
        query = f"""
        MATCH (a:Article) {where_clause}
        OPTIONAL MATCH (d:Document)-[:HAS_ARTICLE]->(a)
        OPTIONAL MATCH (d2:Document)-[:HAS_CHAPTER]->(ch:Chapter)-[:HAS_ARTICLE]->(a)
        WITH a, coalesce(d.title, d2.title, "Văn bản") AS doc_title, coalesce(ch.title, "") AS ch_title
        RETURN a.uid AS uid,
               doc_title + " - " + ch_title + " - " + coalesce(a.title, "") + " - " + coalesce(a.clean_text, "") AS rich_text
        """
        
        with self._driver.session() as session:
            records = session.run(query).data()
            
        stats["total"] = len(records)
        
        for i in range(0, len(records), self._embed_batch):
            batch = records[i : i + self._embed_batch]
            uids = [r["uid"] for r in batch]
            texts = [r["rich_text"] for r in batch]
            
            try:
                embeddings = self._call_embed_service(texts)
                update_query = """
                UNWIND $batch AS row
                MATCH (a:Article {uid: row.uid})
                SET a.embedding = row.embedding
                """
                batch_data = [{"uid": uid, "embedding": emb} for uid, emb in zip(uids, embeddings)]
                with self._driver.session() as session:
                    session.run(update_query, batch=batch_data)
                stats["embedded"] += len(batch)
            except Exception as exc:
                logger.error("Failed to embed batch: %s", exc)
                stats["errors"] += len(batch)
                
        try:
            self._ensure_vector_index()
        except Exception as exc:
            logger.error("Failed to ensure vector index: %s", exc)
            
        return stats

    def embed_article(self, uid: str, text: str) -> Optional[list[float]]:
        """
        Embed a single article and write to Neo4j. Returns the embedding vector.
        Useful for incremental updates or testing.
        """

        try:
            emb = self._call_embed_service([text])[0]
            query = "MATCH (a:Article {uid: $uid}) SET a.embedding = $embedding"
            with self._driver.session() as session:
                session.run(query, uid=uid, embedding=emb)
            return emb
        except Exception as exc:
            logger.error("Failed to embed article %s: %s", uid, exc)
            return None

    def verify_embeddings(self) -> dict[str, int]:
        """
        Check that all Article nodes have a 1024-dim embedding.

        Returns: {"total_articles": N, "with_embedding": N, "missing": N, "wrong_dim": N}
        """

        query = f"""
        MATCH (a:Article)
        RETURN
          count(a) AS total_articles,
          count(a.embedding) AS with_embedding,
          count(CASE WHEN size(a.embedding) <> {EMBED_DIM} THEN 1 END) AS wrong_dim
        """
        with self._driver.session() as session:
            res = session.run(query).single()
            return dict(res) if res else {}

    # ------------------------------------------------------------------
    # Private helpers
    # ------------------------------------------------------------------

    def _call_embed_service(self, texts: list[str]) -> list[list[float]]:
        """
        Call Người A's embedding API with retry logic.
        """
        import requests

        last_error = None
        for attempt in range(EMBED_MAX_RETRIES):
            try:
                resp = requests.post(
                    self._url + "/embed",
                    json={"texts": texts},
                    timeout=60,
                )
                resp.raise_for_status()
                data = resp.json()
                embeddings = data.get("embeddings", [])
                if len(embeddings) != len(texts):
                    raise RuntimeError(f"Expected {len(texts)} embeddings, got {len(embeddings)}")
                if embeddings and len(embeddings[0]) != EMBED_DIM:
                    raise RuntimeError(f"Expected {EMBED_DIM} dims, got {len(embeddings[0])}")
                return embeddings

            except Exception as e:
                last_error = e
                if attempt < EMBED_MAX_RETRIES - 1:
                    delay = EMBED_RETRY_DELAY * (2 ** attempt)
                    logger.warning(
                        f"Embedding service failed (attempt {attempt + 1}/{EMBED_MAX_RETRIES}), "
                        f"retrying in {delay}s: {e}"
                    )
                    time.sleep(delay)

        raise last_error  # type: ignore[misc]

    def _ensure_vector_index(self) -> None:
        """
        Create the Neo4j vector index if it doesn't exist.
        Safe to call multiple times (IF NOT EXISTS).

        Cypher:
        CREATE VECTOR INDEX article_embeddings IF NOT EXISTS
        FOR (a:Article) ON (a.embedding)
        OPTIONS {indexConfig: {
          `vector.dimensions`: 1024,
          `vector.similarity_function`: 'cosine'
        }}
        """

        query = f"""
        CREATE VECTOR INDEX {VECTOR_INDEX_NAME} IF NOT EXISTS
        FOR (a:Article) ON (a.embedding)
        OPTIONS {{indexConfig: {{
          `vector.dimensions`: {EMBED_DIM},
          `vector.similarity_function`: 'cosine'
        }}}}
        """
        with self._driver.session() as session:
            session.run(query)