23 lines
1.1 KiB
Python
23 lines
1.1 KiB
Python
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]
|