diff --git a/load_documents.py b/load_documents.py index 1fe884a..e22f98d 100644 --- a/load_documents.py +++ b/load_documents.py @@ -1,13 +1,12 @@ -# Utility script to load documents from a directory into the Chroma vector store. -# -# The script walks through the specified directory, reads all .txt files, splits them into chunks using -# `RecursiveCharacterTextSplitter`, embeds the chunks with `OllamaEmbeddings`, and stores them in the -# local Chroma collection via the helper functions defined in :mod:`main`. +#!/usr/bin/env python +# Load documents from a directory into the Qdrant vector store. # # Usage: # python load_documents.py /path/to/docs # -# The script prints the number of documents added. +# The script walks through the specified directory, reads all .txt files, splits them into chunks using +# `RecursiveCharacterTextSplitter`, embeds the chunks with `OllamaEmbeddings`, and stores them in the +# local Qdrant collection via the helper functions defined in :mod:`main`. import os import sys @@ -48,8 +47,7 @@ def load_documents_from_dir(directory: str) -> int: documents = chunk_document(content, title) ids = [str(uuid4()) for _ in documents] embeddings = embedding.embed_documents([doc.page_content for doc in documents]) - collection = vector_store.get_collection(name="rag_memory") - collection.add(ids=ids, documents=[doc.page_content for doc in documents], embeddings=embeddings, metadatas=[doc.metadata for doc in documents]) + vector_store.add_documents(documents, ids=ids, embeddings=embeddings) total_chunks += len(documents) return total_chunks