diff --git a/vectorstore.py b/vectorstore.py index 5d5903d..cca6069 100644 --- a/vectorstore.py +++ b/vectorstore.py @@ -1,25 +1,29 @@ -""" -Vector store utilities using ChromaDB and Ollama embeddings. +"""Vector store utilities using ChromaDB and Ollama embeddings. + +This module provides functions to create a persistent Chroma vector store and to load documents +from a directory into the store. Documents are split into chunks using +`RecursiveCharacterTextSplitter`. """ import os from pathlib import Path from typing import List +from langchain_text_splitters import RecursiveCharacterTextSplitter from langchain_chroma import Chroma from langchain_ollama import OllamaEmbeddings -from langchain_text_splitters import RecursiveCharacterTextSplitter -from langchain_ollama import ChatOllama -# Create the vector store with persistence +# Default persistence directory +DEFAULT_PERSIST_DIR = "./chroma_db" -def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma: + +def create_vectorstore(persist_directory: str = DEFAULT_PERSIST_DIR) -> Chroma: """Create or load a Chroma vector store. Parameters ---------- persist_directory: str - Directory where the Chroma DB will be persisted. + Directory where the Chroma DB will be stored. Returns ------- @@ -29,53 +33,29 @@ def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma: embeddings = OllamaEmbeddings(model="nomic-embed-text") return Chroma(persist_directory=persist_directory, embedding_function=embeddings) -# Load documents from a directory and add them to the vector store -def load_documents(directory: str, vectorstore: Chroma) -> None: - """Load .txt and .md files from *directory*, chunk them, and add to *vectorstore*. +def load_documents(directory: str, vectorstore: Chroma, chunk_size: int = 1000, chunk_overlap: int = 200) -> None: + """Load all .txt and .md files from *directory* into *vectorstore*. - Parameters - ---------- - directory: str - Path to the directory containing the documents. - vectorstore: Chroma - The vector store to which the documents will be added. + The files are read, split into chunks with a recursive character splitter and + added to the Chroma collection. """ - # Ensure the directory exists - path = Path(directory) - if not path.is_dir(): - raise FileNotFoundError(f"Directory {directory} does not exist") + splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap) + docs: List[str] = [] + for file_path in Path(directory).rglob("*.txt"): + docs.append(file_path.read_text(encoding="utf-8")) + for file_path in Path(directory).rglob("*.md"): + docs.append(file_path.read_text(encoding="utf-8")) - # 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}") + if not docs: return + # Split documents into chunks + texts = splitter.split_text("\n\n".join(docs)) + # Create a list of dicts with metadata (optional) + metadatas = [{"source": "local"} for _ in texts] + vectorstore.add_texts(texts, metadatas=metadatas) - # 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}, - }) + # Persist changes + vectorstore.persist() - 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) + print(f"Loaded {len(texts)} chunks into ChromaDB.")