18 lines
673 B
Python
18 lines
673 B
Python
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)
|