diff --git a/vectorstore.py b/vectorstore.py index ef1f919..eba9eb1 100644 --- a/vectorstore.py +++ b/vectorstore.py @@ -1,30 +1,25 @@ from pathlib import Path from uuid import uuid4 -from langchain_qdrant import Qdrant -from langchain_core.documents import Document +import chromadb +from chromadb.config import Settings from langchain_ollama import OllamaEmbeddings +from langchain_core.documents import Document from langchain_text_splitters import RecursiveCharacterTextSplitter -QDRANT_DIR = "./qdrant_db" +CHROMA_DIR = "./chroma_db" COLLECTION_NAME = "local_kb" EMBED_MODEL = "nomic-embed-text" OLLAMA_BASE_URL = "http://127.0.0.1:11434" -def create_vectorstore(persist_directory: str = QDRANT_DIR) -> Qdrant: - embeddings = OllamaEmbeddings( - model=EMBED_MODEL, - base_url=OLLAMA_BASE_URL, - ) - return Qdrant( - collection_name=COLLECTION_NAME, - embedding_function=embeddings, - persist_directory=persist_directory, - ) +def create_vectorstore(persist_directory: str = CHROMA_DIR): + client = chromadb.Client(Settings(persist_directory=persist_directory)) + collection = client.get_or_create_collection(name=COLLECTION_NAME) + return collection -def load_documents(directory: str, vectorstore: Qdrant) -> int: +def load_documents(directory: str, vectorstore: chromadb.Collection) -> int: splitter = RecursiveCharacterTextSplitter( chunk_size=1000, chunk_overlap=200, @@ -51,5 +46,7 @@ def load_documents(directory: str, vectorstore: Qdrant) -> int: chunks = splitter.split_documents(raw_docs) ids = [str(uuid4()) for _ in chunks] - vectorstore.add_documents(documents=chunks, ids=ids) + embeddings = OllamaEmbeddings(model=EMBED_MODEL, base_url=OLLAMA_BASE_URL) + vectors = embeddings.embed_documents([doc.page_content for doc in chunks]) + vectorstore.add(ids=ids, documents=[doc.page_content for doc in chunks], embeddings=vectors, metadatas=[doc.metadata for doc in chunks]) return len(chunks)