# 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()