SearchEngine.py

from chromadb.utils import embedding_functions
from chromadb.config import Settings
import chromadb

class searchengine:
    def __init__(self, model_name,collection_name):
        self.sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(
            model_name = model_name
            )
        self.chroma_client = chromadb.Client(
            )
        self.collection = self.chroma_client.get_or_create_collection(name=collection_name)
        
    def add(self, text , metadata,id):
        self.collection.add(
            documents = [text],
            metadatas = [metadata],
            ids = [id]
        )
    
    def add_list(self, texts:list, metadatas:list, ids:list):
        self.collection.add(
            documents = texts,
            metadatas = metadatas,
            ids = ids
        )
        
    def query(self, query, number=2):
        results = self.collection.query(
            query_texts = [query],
            n_results = number
        )
        return results

    def count(self):
        return self.collection.count()
    
    def peak(self):
        result = self.collection.peek()