embedder.py

# core/embedder.py
# CRITICAL: Must set BEFORE importing chromadb to suppress telemetry
import os
os.environ["ANONYMIZED_TELEMETRY"] = "False"

import hashlib
import requests

# ═══════════════════════════════════════════════════════════════════════════
# NVIDIA Nemotron-3-Embed-1B — High-accuracy RAG embeddings via NIM API
# ═══════════════════════════════════════════════════════════════════════════

_DEFAULT_NVIDIA_EMBED_MODEL = "nvidia/nemotron-3-embed-1b"
_NVIDIA_EMBED_DIM = 1024  # nemotron-3-embed-1b output dimension


class NvidiaEmbedder:
    """
    Document Intelligence embedder using nvidia/nemotron-3-embed-1b.
    High-accuracy embeddings for Retrieval-Augmented Generation (RAG).
    Requires NVIDIA_API_KEY. Falls back to DummyEmbedder if unavailable.
    """

    def __init__(self, persist_dir=None):
        self.nvidia_key = os.getenv("NVIDIA_API_KEY", "").strip()
        self.embed_model = os.getenv("NVIDIA_EMBED_MODEL", _DEFAULT_NVIDIA_EMBED_MODEL).strip()
        self.persist_dir = persist_dir or "/data/invicta_data/vectors"
        os.makedirs(self.persist_dir, exist_ok=True)

        # In-memory store: {doc_id: {"text": str, "embedding": list[float], "metadata": dict}}
        self._store = {}
        self._index_path = os.path.join(self.persist_dir, "nvidia_embed_index.json")
        self._load_index()

        if not self.nvidia_key:
            print("⚠️ NvidiaEmbedder: No NVIDIA_API_KEY — will use LocalEmbedder fallback")
        else:
            print(f"✅ NvidiaEmbedder: {self.embed_model} ready")

    def _load_index(self):
        """Load persisted embeddings from disk."""
        try:
            if os.path.exists(self._index_path):
                import json
                with open(self._index_path, "r", encoding="utf-8") as f:
                    self._store = json.load(f)
                print(f"✅ NvidiaEmbedder: Loaded {len(self._store)} vectors from disk")
        except Exception as e:
            print(f"⚠️ NvidiaEmbedder: Could not load index: {e}")
            self._store = {}

    def _save_index(self):
        """Persist embeddings to disk."""
        try:
            import json
            with open(self._index_path, "w", encoding="utf-8") as f:
                json.dump(self._store, f)
        except Exception as e:
            print(f"⚠️ NvidiaEmbedder: Could not save index: {e}")

    def _doc_id(self, text):
        return hashlib.md5(text.strip().lower().encode()).hexdigest()

    def _embed(self, texts):
        """Call nvidia/nemotron-3-embed-1b and return list of embedding vectors."""
        headers = {
            "Authorization": f"Bearer {self.nvidia_key}",
            "Content-Type": "application/json"
        }
        payload = {
            "model": self.embed_model,
            "input": texts,
            "input_type": "passage",
            "encoding_format": "float"
        }
        r = requests.post(
            "https://integrate.api.nvidia.com/v1/embeddings",
            headers=headers,
            json=payload,
            timeout=30
        )
        r.raise_for_status()
        data = r.json()
        # OpenAI-compatible response: data[].embedding
        return [item["embedding"] for item in sorted(data["data"], key=lambda x: x["index"])]

    def _cosine_sim(self, a, b):
        """Cosine similarity between two vectors."""
        dot = sum(x * y for x, y in zip(a, b))
        mag_a = sum(x * x for x in a) ** 0.5
        mag_b = sum(x * x for x in b) ** 0.5
        if mag_a == 0 or mag_b == 0:
            return 0.0
        return dot / (mag_a * mag_b)

    def store(self, text, metadata=None):
        if not text or not text.strip():
            return
        if not self.nvidia_key:
            return
        doc_id = self._doc_id(text)
        if doc_id in self._store:
            return  # Already stored
        try:
            embeddings = self._embed([text])
            self._store[doc_id] = {
                "text": text,
                "embedding": embeddings[0],
                "metadata": metadata or {}
            }
            self._save_index()
        except Exception as e:
            print(f"⚠️ NvidiaEmbedder store failed: {e}")

    def find_similar(self, query, top_k=5):
        if not self.nvidia_key or not self._store:
            return []
        try:
            query_emb = self._embed([query])[0]
            # Score all stored docs
            scored = []
            for doc_id, entry in self._store.items():
                sim = self._cosine_sim(query_emb, entry["embedding"])
                scored.append((sim, entry["text"]))
            scored.sort(key=lambda x: x[0], reverse=True)
            # Return top_k with similarity > 0.4 threshold
            return [text for sim, text in scored[:top_k] if sim > 0.4]
        except Exception as e:
            print(f"⚠️ NvidiaEmbedder find_similar failed: {e}")
            return []

    def count(self):
        return len(self._store)


# ═══════════════════════════════════════════════════════════════════════════
# LOCAL FALLBACK — ChromaDB + default embedding function
# ═══════════════════════════════════════════════════════════════════════════

class LocalEmbedder:
    def __init__(self, persist_dir=None):
        self.persist_dir = persist_dir or "/data/invicta_data/vectors"
        os.makedirs(self.persist_dir, exist_ok=True)
        try:
            import chromadb
            from chromadb.config import Settings
            self.client = chromadb.Client(Settings(
                persist_directory=self.persist_dir,
                anonymized_telemetry=False
            ))
            self.collection = self.client.get_or_create_collection("invicta_embeddings")
        except Exception as e:
            print(f"⚠️ LocalEmbedder: chromadb init failed: {e}")
            self.collection = None
        self._embedding_function = None
        self._embedding_ok = None

    def _ensure_embedding_function(self):
        """Lazy-load the embedding model so startup isn't blocked by a download."""
        if self._embedding_ok is None:
            try:
                from chromadb.utils import embedding_functions
                self._embedding_function = embedding_functions.DefaultEmbeddingFunction()
                self._embedding_ok = True
                print("✅ LocalEmbedder: embedding function loaded")
            except Exception as e:
                print(f"⚠️ LocalEmbedder: embedding function unavailable: {e}")
                self._embedding_ok = False
        return self._embedding_ok

    def _doc_id(self, text):
        return hashlib.md5(text.strip().lower().encode()).hexdigest()

    def store(self, text, metadata=None):
        if not text or not text.strip() or not self.collection:
            return
        if not self._ensure_embedding_function():
            return
        doc_id = self._doc_id(text)
        meta = metadata or {}
        existing = self.collection.get(ids=[doc_id])
        if existing and existing.get("ids") and len(existing["ids"]) > 0:
            return
        self.collection.add(
            ids=[doc_id],
            documents=[text],
            metadatas=[meta],
            embeddings=[self._embedding_function([text])[0]]
        )

    def find_similar(self, query, top_k=5):
        if not self._ensure_embedding_function() or not self.collection:
            return []
        embedding = self._embedding_function([query])[0]
        results = self.collection.query(
            query_embeddings=[embedding],
            n_results=top_k,
            include=["documents", "distances"]
        )
        docs = []
        if results and results.get("documents"):
            for doc, dist in zip(results["documents"][0], results["distances"][0]):
                if dist < 0.8:
                    docs.append(doc)
        return docs

    def count(self):
        return self.collection.count() if self.collection else 0


# ═══════════════════════════════════════════════════════════════════════════
# FACTORY — Returns the best available embedder
# ═══════════════════════════════════════════════════════════════════════════

def get_embedder(persist_dir=None):
    """
    Returns the best available embedder:
    1. NvidiaEmbedder (nvidia/nemotron-3-embed-1b) — if NVIDIA_API_KEY is set
    2. LocalEmbedder (ChromaDB default) — if chromadb is installed
    3. DummyEmbedder — silent no-op fallback
    """
    nvidia_key = os.getenv("NVIDIA_API_KEY", "").strip()
    if nvidia_key:
        try:
            emb = NvidiaEmbedder(persist_dir=persist_dir)
            print("✅ get_embedder → NvidiaEmbedder (nemotron-3-embed-1b)")
            return emb
        except Exception as e:
            print(f"⚠️ NvidiaEmbedder failed, trying LocalEmbedder: {e}")

    try:
        import chromadb  # noqa: F401
        emb = LocalEmbedder(persist_dir=persist_dir)
        print("✅ get_embedder → LocalEmbedder (ChromaDB)")
        return emb
    except ImportError:
        pass

    print("⚠️ get_embedder → DummyEmbedder (no chromadb, no NVIDIA key)")

    class DummyEmbedder:
        def store(self, *a, **k): pass
        def find_similar(self, *a, **k): return []
        def count(self): return 0

    return DummyEmbedder()