diff --git a/vectorstore.py b/vectorstore.py index de9b87f..839c3e9 100644 --- a/vectorstore.py +++ b/vectorstore.py @@ -1,35 +1,37 @@ """ -RAG vector store using Qdrant and Ollama embeddings. +RAG vector store using ChromaDB and Ollama embeddings. """ from pathlib import Path from typing import List -import chromadb # kept for compatibility if needed +import chromadb from langchain.embeddings.ollama import OllamaEmbeddings from langchain.text_splitter import RecursiveCharacterTextSplitter from langchain.schema.document import Document -from langchain.vectorstores import Qdrant +from langchain.vectorstores import Chroma CHROMA_DIR = "./chroma_db" EMBED_MODEL = "nomic-embed-text" -def create_vectorstore(persist_directory: str = CHROMA_DIR): - """Create or load a Qdrant vector store. +def create_vectorstore(persist_directory: str = CHROMA_DIR) -> Chroma: + """Create or load a Chroma vector store. Parameters ---------- - persist_directory: - Directory where the Qdrant database is stored. If it does not exist, it will be created. + persist_directory : str, optional + Directory where the Chroma database is stored. If it does not exist, + it will be created automatically by Chroma. """ embeddings = OllamaEmbeddings(model=EMBED_MODEL) - # Qdrant can use a local file store via `path` argument - client = Qdrant(persist_directory=persist_directory, embedding_function=embeddings) - return client + # Chroma can use a local directory for persistence + client = chromadb.PersistentClient(path=persist_directory) + collection = client.get_or_create_collection(name="documents", embedding_function=embeddings) + return Chroma(collection=collection, embedding_function=embeddings) -def load_documents(directory: str, vectorstore) -> None: +def load_documents(directory: str, vectorstore: Chroma) -> None: """Load all .txt and .md files from *directory*, chunk them and add to the vector store. The function does not return anything; it mutates the provided collection. @@ -44,8 +46,5 @@ def load_documents(directory: str, vectorstore) -> None: docs.extend(text_splitter.create_documents([content], metadata={"source": str(path)})) if docs: - # Qdrant expects texts and metadatas lists - vectorstore.add_texts( - texts=[doc.page_content for doc in docs], - metadatas=[doc.metadata for doc in docs], - ) + # Chroma expects a list of documents via add_documents + vectorstore.add_documents(docs)