diff --git a/vectorstore.py b/vectorstore.py index f05d4bb..3accc37 100644 --- a/vectorstore.py +++ b/vectorstore.py @@ -1,51 +1,94 @@ -"""Vector store utilities for ChromaDB. +"""Vector store utilities for ChromaDB with Ollama embeddings. -This module provides functions to create a persistent ChromaDB vector store using -Ollama embeddings and to load documents from a directory into the store. +This module provides functions to create a persistent Chroma vector store and +load documents from a directory into it. Documents are split into chunks using +`RecursiveCharacterTextSplitter` and stored in the Chroma collection. """ +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_core.documents import Document +# --------------------------------------------------------------------------- +# Vector store creation +# --------------------------------------------------------------------------- def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma: - """Create a Chroma vector store with Ollama embeddings. + """Create or load a Chroma vector store. Parameters ---------- persist_directory: str - Directory where the vector store will be persisted. + Directory where the Chroma DB files are stored. Returns ------- Chroma A Chroma vector store instance. """ + # Ensure directory exists + Path(persist_directory).mkdir(parents=True, exist_ok=True) + # Use Ollama embeddings embeddings = OllamaEmbeddings(model="nomic-embed-text") - return Chroma(persist_directory=persist_directory, embedding_function=embeddings) + # Create Chroma store + vectorstore = Chroma(persist_directory=persist_directory, embedding_function=embeddings) + return vectorstore + +# --------------------------------------------------------------------------- +# Document loading +# --------------------------------------------------------------------------- + +def _load_text_files(directory: str) -> List[str]: + """Load all .txt and .md files from a directory into a list of strings.""" + texts = [] + for root, _, files in os.walk(directory): + for file in files: + if file.lower().endswith(('.txt', '.md')): + path = Path(root) / file + try: + content = path.read_text(encoding="utf-8") + texts.append(content) + except Exception as e: + print(f"Failed to read {path}: {e}") + return texts -def load_documents(directory: str, vectorstore: Chroma) -> None: - """Load .txt and .md files from *directory* into *vectorstore*. +def load_documents(directory: str, vectorstore: Chroma, chunk_size: int = 1000, chunk_overlap: int = 200) -> None: + """Load documents from a directory into the provided vector store. - The documents are split into chunks using ``RecursiveCharacterTextSplitter`` - before being added to the vector store. + Parameters + ---------- + directory: str + Path to the directory containing .txt/.md files. + vectorstore: Chroma + The vector store to add documents to. + chunk_size: int, optional + Maximum size of each chunk. + chunk_overlap: int, optional + Number of characters to overlap between chunks. """ - splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) + texts = _load_text_files(directory) + if not texts: + print("No text files found in the directory.") + return + + splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap) docs = [] - for file_path in Path(directory).glob("*"): - if file_path.suffix.lower() not in {".txt", ".md"}: - continue - with open(file_path, "r", encoding="utf-8") as f: - text = f.read() - chunks = splitter.split_text(text) - docs.extend([Document(page_content=c, metadata={"source": str(file_path)}) for c in chunks]) - if docs: - vectorstore.add_documents(docs) - print(f"Loaded {len(docs)} chunks from {directory} into ChromaDB.") - else: - print(f"No .txt/.md files found in {directory}.") + for text in texts: + docs.extend(splitter.split_text(text)) + + # Add documents to Chroma + vectorstore.add_texts(docs) + print(f"Loaded {len(docs)} chunks into the vector store.") + +# --------------------------------------------------------------------------- +# Example usage (uncomment to run directly) +# --------------------------------------------------------------------------- +# if __name__ == "__main__": +# store = create_vectorstore() +# load_documents("documents", store) +"" \ No newline at end of file