File size: 2,631 Bytes
3d7a63c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 | # core/embedder.py
# CRITICAL: Must set BEFORE importing chromadb to suppress telemetry
import os
os.environ["ANONYMIZED_TELEMETRY"] = "False"
import hashlib
import chromadb
from chromadb.config import Settings
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)
self.client = chromadb.Client(Settings(
persist_directory=self.persist_dir,
anonymized_telemetry=False
))
self.collection = self.client.get_or_create_collection("invicta_embeddings")
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("✅ Embedding function loaded")
except Exception as e:
print(f"⚠️ 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():
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():
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() |