diff --git a/vectorstore.py b/vectorstore.py index 3accc37..f9b0825 100644 --- a/vectorstore.py +++ b/vectorstore.py @@ -1,94 +1,61 @@ -"""Vector store utilities for ChromaDB with Ollama embeddings. +"""Vector store utilities using ChromaDB and Ollama embeddings. -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. +This module provides functions to create a persistent Chroma vector store +and to load documents from a directory into the store. """ -import os from pathlib import Path -from typing import List +from typing import Iterable -from langchain_text_splitters import RecursiveCharacterTextSplitter from langchain_chroma import Chroma from langchain_ollama import OllamaEmbeddings +from langchain_text_splitters import RecursiveCharacterTextSplitter # --------------------------------------------------------------------------- -# Vector store creation +# Create a persistent Chroma vector store. # --------------------------------------------------------------------------- def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma: - """Create or load a Chroma vector store. + """Create a Chroma vector store with Ollama embeddings. Parameters ---------- persist_directory: str - Directory where the Chroma DB files are stored. + Directory where the Chroma database will be 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") - # Create Chroma store - vectorstore = Chroma(persist_directory=persist_directory, embedding_function=embeddings) - return vectorstore + return Chroma(persist_directory=persist_directory, embedding_function=embeddings) # --------------------------------------------------------------------------- -# Document loading +# Load documents from a directory into the vector store. # --------------------------------------------------------------------------- -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 all .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. - - 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. + The documents are split into chunks using a RecursiveCharacterTextSplitter + before being added to the vector store. """ - 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) + splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) docs = [] - 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.") + for path in Path(directory).rglob("*.txt"): + docs.append(path.read_text(encoding="utf-8")) + for path in Path(directory).rglob("*.md"): + docs.append(path.read_text(encoding="utf-8")) + if not docs: + return + # Split the documents into chunks. + chunks = splitter.split_text("\n\n".join(docs)) + # Create LangChain Document objects. + from langchain.docstore.document import Document + documents = [Document(page_content=chunk) for chunk in chunks] + vectorstore.add_documents(documents) # --------------------------------------------------------------------------- -# Example usage (uncomment to run directly) -# --------------------------------------------------------------------------- -# if __name__ == "__main__": -# store = create_vectorstore() -# load_documents("documents", store) -"" \ No newline at end of file +# End of vectorstore.py +# --------------------------------------------------------------------------- \ No newline at end of file