vector_search/index_faiss.py

"""
In-Memory FAISS HNSW Vector Indexing and Search.

Features:
- In-Memory HNSW (IndexHNSWFlat) with cosine similarity (METRIC_INNER_PRODUCT on normalized vectors).
- M=32, efConstruction=200, efSearch=64.
- Single combined index per strategy spanning all configured languages (config.LANGUAGES).
- Metadata-based language pre-filtering.
- Corpus centroid computation and persistence for pre-retrieval off-topic guardrails.
"""

import json
import logging
import os
import sys
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple
import faiss
import numpy as np

# Ensure project root is in sys.path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
import config
from doc_chunking.metadata import Chunk, filter_chunks_by_language
from doc_chunking.passage_native import process_corpus_passage_native
from doc_chunking.sentence_window import process_longdocs_sentence_window
from doc_chunking.semantic import process_longdocs_semantic
from vector_search.embed import get_embedder

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


class StrategyVectorIndex:
    """
    Manages an in-memory FAISS HNSW index and aligned metadata for a specific chunking strategy.
    """
    def __init__(
        self,
        strategy_name: str,
        dim: int = config.EMBEDDING_DIM,
        m: int = config.HNSW_M,
        ef_construction: int = config.HNSW_EF_CONSTRUCTION,
        ef_search: int = config.HNSW_EF_SEARCH,
    ):
        self.strategy_name = strategy_name
        self.dim = dim
        self.m = m
        self.ef_construction = ef_construction
        self.ef_search = ef_search
        
        # FAISS HNSW with inner product (cosine similarity on unit-normalized vectors)
        self.index = faiss.IndexHNSWFlat(self.dim, self.m, faiss.METRIC_INNER_PRODUCT)
        self.index.hnsw.efConstruction = self.ef_construction
        self.index.hnsw.efSearch = self.ef_search
        
        self.chunks: List[Chunk] = []
        # Precomputed index mappings for ultra-fast language filtering
        self.lang_to_indices: Dict[str, List[int]] = {}

    @property
    def size(self) -> int:
        """Returns the number of indexed chunks."""
        return len(self.chunks)

    def add_chunks(self, chunks: List[Chunk], vectors: np.ndarray):
        """Add chunks and precomputed vectors to index."""
        if len(chunks) == 0:
            return
        if len(chunks) != len(vectors):
            raise ValueError(f"Chunk count ({len(chunks)}) != vector count ({len(vectors)})")
            
        start_idx = len(self.chunks)
        self.chunks.extend(chunks)
        
        # Build language index mapping
        for idx_offset, c in enumerate(chunks):
            global_idx = start_idx + idx_offset
            lang = c.source_lang.lower()
            if lang not in self.lang_to_indices:
                self.lang_to_indices[lang] = []
            self.lang_to_indices[lang].append(global_idx)
            
        self.index.add(np.ascontiguousarray(vectors, dtype=np.float32))
        logger.info(
            f"Added {len(chunks)} chunks to '{self.strategy_name}' index. "
            f"Total index size: {self.index.ntotal}."
        )

    def search(
        self, query_vec: np.ndarray, target_lang: Optional[str] = None, top_k: int = 15
    ) -> List[Dict[str, Any]]:
        """
        Search the index for query vector with optional language filtering.
        """
        if self.index.ntotal == 0:
            return []
            
        # Ensure query_vec is 2D (1, dim)
        if query_vec.ndim == 1:
            query_vec = np.expand_dims(query_vec, axis=0)
            
        # Query FAISS HNSW (optimized candidate search for sub-1ms traversal)
        search_k = min(self.index.ntotal, max(400, top_k * 25) if target_lang else max(60, top_k * 3))
        scores, indices = self.index.search(query_vec, search_k)

        
        results: List[Dict[str, Any]] = []
        target_lang_clean = target_lang.lower().strip() if target_lang else None
        
        for score, idx in zip(scores[0], indices[0]):
            if idx < 0 or idx >= len(self.chunks):
                continue
            chunk = self.chunks[idx]
            
            # Apply language metadata pre-filter
            if target_lang_clean and chunk.source_lang.lower() != target_lang_clean:
                continue
                
            results.append({
                "chunk_id": chunk.chunk_id,
                "text": chunk.text,
                "score": float(score),
                "source_lang": chunk.source_lang,
                "chunk_strategy": chunk.chunk_strategy,
                "source_query_ids": chunk.source_query_ids,
                "doc_id": chunk.doc_id,
                "context_window": chunk.context_window,
                "metadata": chunk.metadata,
            })
            
            if len(results) >= top_k:
                break
                
        return results

    def save(self, directory: Path):
        """Save FAISS index and chunk metadata to disk."""
        directory.mkdir(parents=True, exist_ok=True)
        index_file = directory / f"{self.strategy_name}.faiss"
        meta_file = directory / f"{self.strategy_name}_meta.json"
        
        faiss.write_index(self.index, str(index_file))
        
        # Serialize chunk metadata
        meta_data = {
            "strategy_name": self.strategy_name,
            "dim": self.dim,
            "m": self.m,
            "ef_construction": self.ef_construction,
            "ef_search": self.ef_search,
            "chunks": [c.model_dump() for c in self.chunks],
        }
        with open(meta_file, "w", encoding="utf-8") as f:
            json.dump(meta_data, f, ensure_ascii=False)
        logger.info(f"Saved index and metadata for '{self.strategy_name}' to {directory}")

    @classmethod
    def load(cls, directory: Path, strategy_name: str) -> "StrategyVectorIndex":
        """Load FAISS index and chunk metadata from disk."""
        index_file = directory / f"{strategy_name}.faiss"
        meta_file = directory / f"{strategy_name}_meta.json"
        
        if not index_file.exists() or not meta_file.exists():
            raise FileNotFoundError(f"Index or metadata file missing in {directory} for '{strategy_name}'")
            
        with open(meta_file, "r", encoding="utf-8") as f:
            meta_data = json.load(f)
            
        inst = cls(
            strategy_name=meta_data["strategy_name"],
            dim=meta_data.get("dim", config.EMBEDDING_DIM),
            m=meta_data.get("m", config.HNSW_M),
            ef_construction=meta_data.get("ef_construction", config.HNSW_EF_CONSTRUCTION),
            ef_search=meta_data.get("ef_search", config.HNSW_EF_SEARCH),
        )
        inst.index = faiss.read_index(str(index_file))
        inst.index.hnsw.efSearch = inst.ef_search
        
        # Reconstruct chunks and language mapping
        inst.chunks = [Chunk(**c) for c in meta_data["chunks"]]
        inst.lang_to_indices = {}
        for idx, c in enumerate(inst.chunks):
            lang = c.source_lang.lower()
            if lang not in inst.lang_to_indices:
                inst.lang_to_indices[lang] = []
            inst.lang_to_indices[lang].append(idx)
            
        logger.info(f"Loaded index '{strategy_name}' ({inst.index.ntotal} vectors) from {directory}")
        return inst


class IndexManager:
    """
    Manages all strategy indexes and off-topic centroid models for the pipeline.
    """
    def __init__(self, index_dir: Path = config.INDEX_DIR):
        self.index_dir = index_dir
        self.indexes: Dict[str, StrategyVectorIndex] = {}
        self.centroids: Dict[str, np.ndarray] = {}  # lang -> centroid vector
        self.global_centroid: Optional[np.ndarray] = None
        self.embedder = get_embedder()

    @staticmethod
    def _iter_jsonl(path: Path, limit: Optional[int] = None):
        """Yield JSONL records without loading an entire corpus into memory."""
        yielded = 0
        with open(path, "r", encoding="utf-8") as handle:
            for line_number, line in enumerate(handle, start=1):
                if limit is not None and yielded >= limit:
                    break
                if not line.strip():
                    continue
                try:
                    record = json.loads(line)
                except json.JSONDecodeError as exc:
                    logger.warning("Skipping malformed JSONL record %s:%s: %s", path, line_number, exc)
                    continue
                yielded += 1
                yield record

    def _add_embedded_chunks(
        self,
        index: StrategyVectorIndex,
        chunks: List[Chunk],
        centroid_sums: Dict[str, np.ndarray],
        centroid_counts: Dict[str, int],
    ) -> int:
        """Embed and add one bounded batch, updating centroid statistics online."""
        if not chunks:
            return 0
        texts = [chunk.embed_text for chunk in chunks]
        vectors = self.embedder.encode_passages(
            texts,
            batch_size=min(config.INDEX_BUILD_BATCH_SIZE, 128),
        )
        if len(vectors) != len(chunks):
            raise RuntimeError(
                f"Embedder returned {len(vectors)} vectors for {len(chunks)} chunks"
            )
        index.add_chunks(chunks, vectors)
        for chunk, vector in zip(chunks, vectors):
            lang = chunk.source_lang.lower()
            if lang not in centroid_sums:
                centroid_sums[lang] = np.zeros(index.dim, dtype=np.float64)
                centroid_counts[lang] = 0
            centroid_sums[lang] += vector.astype(np.float64, copy=False)
            centroid_counts[lang] += 1
        return len(chunks)

    def build_all_indexes(self, max_passages_per_lang: Optional[int] = None):
        """
        Build complete combined indexes for all configured languages.

        ``None`` means no per-language cap. Records are streamed from JSONL and
        embedded in bounded batches, so the full corpus is indexed without the
        previous implicit 700-record truncation or an all-corpus memory spike.
        """
        if max_passages_per_lang is not None and max_passages_per_lang <= 0:
            raise ValueError("max_passages_per_lang must be positive or None")
        batch_size = max(1, int(getattr(config, "INDEX_BUILD_BATCH_SIZE", 512)))
        logger.info(
            "Building complete indexes for languages=%s (passage_limit=%s, batch_size=%s)",
            config.LANGUAGES,
            max_passages_per_lang if max_passages_per_lang is not None else "unlimited",
            batch_size,
        )

        self.indexes = {}
        self.centroids = {}
        self.global_centroid = None
        centroid_sums: Dict[str, np.ndarray] = {}
        centroid_counts: Dict[str, int] = {}
        passage_counts: Dict[str, int] = {}
        longdoc_counts: Dict[str, int] = {}

        # Passage-native index: stream each language corpus and add bounded batches.
        passage_index = StrategyVectorIndex("passage_native")
        passage_total = 0
        for lang in config.LANGUAGES:
            corpus_file = config.PROCESSED_DATA_DIR / f"{lang}_corpus.jsonl"
            if not corpus_file.exists():
                logger.warning("Corpus file %s not found; skipping language '%s'.", corpus_file, lang)
                continue
            batch_records: List[Dict[str, Any]] = []
            lang_total = 0
            for record in self._iter_jsonl(corpus_file, max_passages_per_lang):
                batch_records.append(record)
                if len(batch_records) >= batch_size:
                    chunks = process_corpus_passage_native(batch_records)
                    added = self._add_embedded_chunks(
                        passage_index, chunks, centroid_sums, centroid_counts
                    )
                    lang_total += added
                    logger.info("Language '%s': embedded %d chunks (lang total: %d)...", lang, added, lang_total)
                    batch_records.clear()
            if batch_records:
                chunks = process_corpus_passage_native(batch_records)
                added = self._add_embedded_chunks(
                    passage_index, chunks, centroid_sums, centroid_counts
                )
                lang_total += added
                logger.info("Language '%s': embedded final %d chunks (lang total: %d)...", lang, added, lang_total)
            passage_total += lang_total
            passage_counts[lang] = lang_total
            logger.info("Finished indexing %s passage-native chunks for '%s'.", lang_total, lang)

        if passage_total:
            passage_index.save(self.index_dir)
            self.indexes["passage_native"] = passage_index
            self._save_centroids(centroid_sums, centroid_counts)
        else:
            logger.warning("No passage corpus records were indexed.")

        # Long-document index: process one document at a time and embed bounded batches.
        longdoc_index = StrategyVectorIndex("semantic_longdoc")
        longdoc_total = 0
        for lang in config.LANGUAGES:
            longdoc_file = config.PROCESSED_DATA_DIR / f"{lang}_longdocs.jsonl"
            if not longdoc_file.exists():
                logger.info("Long-document file %s not found; skipping.", longdoc_file)
                continue
            pending_chunks: List[Chunk] = []
            lang_total = 0
            for record in self._iter_jsonl(longdoc_file):
                pending_chunks.extend(process_longdocs_sentence_window([record]))
                pending_chunks.extend(process_longdocs_semantic([record]))
                if len(pending_chunks) >= batch_size:
                    lang_total += self._add_embedded_chunks(
                        longdoc_index, pending_chunks, {}, {}
                    )
                    pending_chunks.clear()
            if pending_chunks:
                lang_total += self._add_embedded_chunks(
                    longdoc_index, pending_chunks, {}, {}
                )
            longdoc_total += lang_total
            longdoc_counts[lang] = lang_total
            logger.info("Indexed %s long-document chunks for '%s'.", lang_total, lang)

        if longdoc_total:
            longdoc_index.save(self.index_dir)
            self.indexes["semantic_longdoc"] = longdoc_index
        else:
            logger.warning("No long-document records were indexed.")

        manifest = {
            "languages": list(config.LANGUAGES),
            "full_corpus": max_passages_per_lang is None,
            "passage_counts": passage_counts,
            "longdoc_counts": longdoc_counts,
            "embedding_model": config.EMBEDDING_MODEL_NAME,
        }
        with open(self.index_dir / "index_manifest.json", "w", encoding="utf-8") as handle:
            json.dump(manifest, handle, ensure_ascii=False, indent=2)

        logger.info(
            "Completed full indexing: %s passage-native and %s long-document chunks.",
            passage_total,
            longdoc_total,
        )

    def _save_centroids(
        self, centroid_sums: Dict[str, np.ndarray], centroid_counts: Dict[str, int]
    ):
        """Normalize and persist online per-language and global centroid statistics."""
        if not centroid_counts:
            return
        centroids_dict: Dict[str, List[float]] = {}
        total_sum = np.zeros(config.EMBEDDING_DIM, dtype=np.float64)
        total_count = 0
        self.centroids = {}
        for lang in config.LANGUAGES:
            if lang not in centroid_counts:
                continue
            mean_vec = centroid_sums[lang] / centroid_counts[lang]
            norm_vec = (mean_vec / (np.linalg.norm(mean_vec) + 1e-12)).astype(np.float32)
            centroids_dict[lang] = norm_vec.tolist()
            self.centroids[lang] = norm_vec
            total_sum += centroid_sums[lang]
            total_count += centroid_counts[lang]
        global_mean = total_sum / max(total_count, 1)
        global_norm = (global_mean / (np.linalg.norm(global_mean) + 1e-12)).astype(np.float32)
        centroids_dict["global"] = global_norm.tolist()
        self.global_centroid = global_norm
        centroid_file = self.index_dir / "centroids.json"
        with open(centroid_file, "w", encoding="utf-8") as handle:
            json.dump(centroids_dict, handle, ensure_ascii=False)
        logger.info("Saved corpus centroids to %s", centroid_file)

    def _compute_and_save_centroids(self, chunks: List[Chunk], vectors: np.ndarray):
        """Compatibility helper for callers that already hold one vector batch."""
        sums: Dict[str, np.ndarray] = {}
        counts: Dict[str, int] = {}
        for chunk, vector in zip(chunks, vectors):
            lang = chunk.source_lang.lower()
            sums.setdefault(lang, np.zeros(vectors.shape[1], dtype=np.float64))
            sums[lang] += vector.astype(np.float64, copy=False)
            counts[lang] = counts.get(lang, 0) + 1
        self._save_centroids(sums, counts)

    def load_all_indexes(self):
        """Loads all existing strategy indexes and centroids from disk into memory with auto-rebuild fallback."""
        loaded_ok = True
        for strategy in ["passage_native", "semantic_longdoc"]:
            try:
                idx = StrategyVectorIndex.load(self.index_dir, strategy)
                if idx.index.ntotal == 0:
                    raise ValueError(f"Index '{strategy}' has 0 vectors.")
                self.indexes[strategy] = idx
            except Exception as e:
                logger.warning(f"Could not load index '{strategy}': {e}")
                loaded_ok = False
                
        # Load centroids
        centroid_file = self.index_dir / "centroids.json"
        if centroid_file.exists():
            try:
                with open(centroid_file, "r", encoding="utf-8") as f:
                    c_data = json.load(f)
                for k, v in c_data.items():
                    arr = np.array(v, dtype=np.float32)
                    if k == "global":
                        self.global_centroid = arr
                    else:
                        self.centroids[k] = arr
                logger.info(f"Loaded centroids for: {list(self.centroids.keys())}")
            except Exception as e:
                logger.warning(f"Failed loading centroids: {e}")
                loaded_ok = False
        else:
            loaded_ok = False

        # Rebuild stale artifacts produced by old languages or missing manifest.
        manifest_file = self.index_dir / "index_manifest.json"
        manifest_ok = False
        if manifest_file.exists():
            try:
                with open(manifest_file, "r", encoding="utf-8") as handle:
                    manifest = json.load(handle)
                manifest_ok = set(manifest.get("languages", [])) == set(config.LANGUAGES)
            except Exception as exc:
                logger.warning("Failed loading index manifest: %s", exc)
        else:
            logger.info("Index manifest is missing; treating existing indexes as stale.")

        if not loaded_ok or not manifest_ok or "passage_native" not in self.indexes or self.indexes["passage_native"].index.ntotal == 0:
            logger.info("[IndexManager] Auto-recovering: building complete indexes from corpus...")
            limit = getattr(config, "MAX_INDEX_PASSAGES_PER_LANG", None)
            self.build_all_indexes(max_passages_per_lang=limit)



_INDEX_MANAGER: Optional[IndexManager] = None


def get_index_manager() -> Any:
    """Singleton getter for IndexManager (supports FAISS and Qdrant vector backends)."""
    global _INDEX_MANAGER
    if getattr(config, "VECTOR_STORE_BACKEND", "faiss").lower() == "qdrant":
        from vector_search.index_qdrant import get_qdrant_manager
        return get_qdrant_manager()
    if _INDEX_MANAGER is None:
        _INDEX_MANAGER = IndexManager()
        _INDEX_MANAGER.load_all_indexes()
    return _INDEX_MANAGER


if __name__ == "__main__":
    manager = IndexManager()
    manager.build_all_indexes(max_passages_per_lang=None)