from langchain_community.vectorstores import Chroma
from langchain_ollama import OllamaEmbeddings
from langchain.schema import Document

class ChromaStore:
    def __init__(self, collection_name="rag_collection"):
        self.client = Chroma(embedding_function=OllamaEmbeddings(model="nomic-embed-text"), collection_name=collection_name)
        # ensure collection exists
        if not self.client.collection_exists:
            self.client.create_collection()

    def add_documents(self, docs):
        # docs: list of Document
        self.client.add_documents(docs)

    def search(self, query, limit=5):
        return self.client.similarity_search(query, k=limit)