diff --git a/vectorstore.py b/vectorstore.py new file mode 100644 index 0000000..2a4d9ef --- /dev/null +++ b/vectorstore.py @@ -0,0 +1,54 @@ +import os +from pathlib import Path +from typing import List + +from langchain_ollama import OllamaEmbeddings +from langchain_chroma import Chroma +from langchain_text_splitters import RecursiveCharacterTextSplitter +from langchain.docstore.document import Document + + +def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma: + """Create or load a Chroma vector store backed by Ollama embeddings. + + Parameters + ---------- + persist_directory: str + Directory where Chroma will persist its data. + + Returns + ------- + Chroma + A Chroma vector store instance. + """ + embeddings = OllamaEmbeddings(model="nomic-embed-text") + return Chroma( + persist_directory=persist_directory, + embedding_function=embeddings, + ) + + +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*. + + The function reads files, splits them into chunks using + ``RecursiveCharacterTextSplitter`` and adds the resulting + :class:`~langchain.docstore.document.Document` objects to the + vector store. + """ + splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap) + docs: List[Document] = [] + for path in Path(directory).glob("**/*"): + if path.suffix.lower() not in {".txt", ".md"}: + continue + text = path.read_text(encoding="utf-8") + docs.extend(splitter.split_text(text)) + # Convert list of strings to Document objects + documents = [Document(page_content=chunk, metadata={"source": str(p)}) for chunk in docs] + vectorstore.add_documents(documents) + +# If this module is executed directly, load the default documents folder. +if __name__ == "__main__": + vs = create_vectorstore() + load_documents("documents", vs) + print("Vector store populated.")