from chromadb import Client as ChromaClient
from langchain_ollama import OllamaEmbeddings

class ChromaStore:
    def __init__(self, collection_name="rag_collection"):
        self.client = ChromaClient()
        self.collection_name = collection_name
        self._ensure_collection()

    def _ensure_collection(self):
        if self.collection_name not in [c.name for c in self.client.get_collections()]:
            self.client.create_collection(name=self.collection_name, metadata={})

    def add_documents(self, documents, titles):
        embeddings = OllamaEmbeddings(model="nomic-embed-text")
        vectors = embeddings.embed_documents(documents)
        payload = [{"title": t} for t in titles]
        self.client.upsert(collection_name=self.collection_name, documents=documents, ids=[str(i) for i in range(len(documents))], metadatas=payload)

    def search(self, query, limit=5):
        results = self.client.query(collection_name=self.collection_name, query_text=query, n_results=limit, include_metadata=True)
        return [r['metadata'] for r in results]