diff --git a/vectorstore.py b/vectorstore.py index 3d9c2fe..5d5903d 100644 --- a/vectorstore.py +++ b/vectorstore.py @@ -1,78 +1,81 @@ -"""Vector store utilities using ChromaDB and Ollama embeddings. - -This module provides functions to create a persistent ChromaDB vector store and -load documents from a directory into it. The documents are split into -manageable chunks using a recursive character text splitter. +""" +Vector store utilities using ChromaDB and Ollama embeddings. """ +import os from pathlib import Path -from typing import Iterable +from typing import List from langchain_chroma import Chroma from langchain_ollama import OllamaEmbeddings from langchain_text_splitters import RecursiveCharacterTextSplitter +from langchain_ollama import ChatOllama -# Default persistence directory -DEFAULT_PERSIST_DIR = "./chroma_db" +# Create the vector store with persistence -# Default chunking parameters – these can be tuned -DEFAULT_CHUNK_SIZE = 1000 -DEFAULT_CHUNK_OVERLAP = 200 - - -def create_vectorstore(persist_directory: str = DEFAULT_PERSIST_DIR) -> Chroma: +def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma: """Create or load a Chroma vector store. Parameters ---------- persist_directory: str - Directory where the ChromaDB files are stored. + Directory where the Chroma DB will be persisted. Returns ------- Chroma - An instance of the Chroma vector store backed by the given directory. + The Chroma vector store instance. """ embeddings = OllamaEmbeddings(model="nomic-embed-text") - vectorstore = Chroma(persist_directory=persist_directory, embedding_function=embeddings) - return vectorstore + return Chroma(persist_directory=persist_directory, embedding_function=embeddings) +# Load documents from a directory and add them to the vector store -def _load_text_files(directory: str | Path) -> Iterable[str]: - """Yield the contents of all .txt and .md files in *directory*. +def load_documents(directory: str, vectorstore: Chroma) -> None: + """Load .txt and .md files from *directory*, chunk them, and add to *vectorstore*. Parameters ---------- - directory: str | Path + directory: str Path to the directory containing the documents. - """ - directory = Path(directory) - for file_path in directory.rglob("*.txt"): - yield file_path.read_text(encoding="utf-8") - for file_path in directory.rglob("*.md"): - yield file_path.read_text(encoding="utf-8") - - -def load_documents(directory: str | Path, vectorstore: Chroma) -> None: - """Load documents from *directory* into *vectorstore*. - - The documents are split into chunks using a recursive character splitter - and then added to the vector store. Existing documents are not removed – - this function simply appends new data. - - Parameters - ---------- - directory: str | Path - Directory containing the source documents. vectorstore: Chroma - The vector store to populate. + The vector store to which the documents will be added. """ - splitter = RecursiveCharacterTextSplitter(chunk_size=DEFAULT_CHUNK_SIZE, chunk_overlap=DEFAULT_CHUNK_OVERLAP) - for text in _load_text_files(directory): - chunks = splitter.split_text(text) - vectorstore.add_texts(chunks) + # Ensure the directory exists + path = Path(directory) + if not path.is_dir(): + raise FileNotFoundError(f"Directory {directory} does not exist") - # Persist changes - vectorstore.persist() + # Collect all .txt and .md files + files = list(path.rglob("*.txt")) + list(path.rglob("*.md")) + if not files: + print(f"No .txt or .md files found in {directory}") + return -# End of vectorstore.py + # Read and chunk the documents + text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) + documents = [] + for file_path in files: + try: + content = file_path.read_text(encoding="utf-8") + except Exception as e: + print(f"Failed to read {file_path}: {e}") + continue + # Split into chunks + chunks = text_splitter.split_text(content) + for i, chunk in enumerate(chunks): + documents.append({ + "page_content": chunk, + "metadata": {"source": str(file_path), "chunk_index": i}, + }) + + if documents: + vectorstore.add_documents(documents) + print(f"Added {len(documents)} chunks to the vector store from {directory}") + else: + print("No documents were processed.") + +# Example usage: +# if __name__ == "__main__": +# store = create_vectorstore() +# load_documents("./documents", store)