utils/embeddings.py

import os
from typing import List
from sentence_transformers import SentenceTransformer
import chromadb
from chromadb.utils import embedding_functions

DEFAULT_MODEL = os.getenv("EMB_MODEL", "sentence-transformers/all-MiniLM-L6-v2")

def get_embedder():
    # Use a light SentenceTransformer and pass to Chroma
    model = SentenceTransformer(DEFAULT_MODEL)
    ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=DEFAULT_MODEL)
    return model, ef

def get_chroma(collection_name: str = "devices", persist_dir: str = ".chroma"):
    _, ef = get_embedder()
    client = chromadb.PersistentClient(path=persist_dir)
    try:
        col = client.get_collection(collection_name, embedding_function=ef)
    except Exception:
        col = client.create_collection(collection_name, embedding_function=ef)
    return client, col