diff --git a/vectorstore.py b/vectorstore.py index 46018df..1902b7a 100644 --- a/vectorstore.py +++ b/vectorstore.py @@ -1,26 +1,83 @@ +""" +Vector store utilities for ChromaDB with Ollama embeddings. + +Functions: +- create_vectorstore(persist_directory="./chroma_db") -> chromadb.Chroma +- load_documents(directory: str, vectorstore) -> None +""" 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_chroma import Chroma from langchain.docstore.document import Document -def create_vectorstore(persist_directory: str = "./chroma_db"): - if not os.path.exists(persist_directory): - os.makedirs(persist_directory, exist_ok=True) + +def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma: + """Create or load a Chroma vector store. + + Parameters + ---------- + persist_directory : str, optional + Directory where the Chroma database is stored. Defaults to ``./chroma_db``. + + Returns + ------- + chromadb.Chroma + A Chroma instance backed by Ollama embeddings. + """ + # Ensure directory exists + Path(persist_directory).mkdir(parents=True, exist_ok=True) embeddings = OllamaEmbeddings(model="nomic-embed-text") return Chroma(persist_directory=persist_directory, embedding_function=embeddings) -def load_documents(directory: str, vectorstore) -> None: - splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) + +def _load_text_files(directory: str) -> List[Document]: + """Load all .txt and .md files from *directory* into Documents. + + Parameters + ---------- + directory : str + Path to the folder containing documents. + + Returns + ------- + list[langchain.docstore.document.Document] + List of Document objects with ``page_content`` set to file text and + ``metadata`` containing the source path. + """ docs: List[Document] = [] - for path in Path(directory).rglob("*.txt"): - text = path.read_text(encoding="utf-8") - docs.extend(splitter.split_documents([Document(page_content=text)])) - for path in Path(directory).rglob("*.md"): - text = path.read_text(encoding="utf-8") - docs.extend(splitter.split_documents([Document(page_content=text)])) - if docs: - vectorstore.add_documents(docs) + for root, _, files in os.walk(directory): + for fname in files: + if not fname.lower().endswith(('.txt', '.md')): + continue + full_path = Path(root) / fname + text = full_path.read_text(encoding="utf-8") + docs.append(Document(page_content=text, metadata={"source": str(full_path)})) + return docs + + +def load_documents(directory: str, vectorstore: Chroma) -> None: + """Chunk documents from *directory* and add them to *vectorstore*. + + Parameters + ---------- + directory : str + Folder with .txt/.md files. + vectorstore : chromadb.Chroma + The vector store instance returned by ``create_vectorstore``. + """ + docs = _load_text_files(directory) + if not docs: + return + splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) + chunks = splitter.split_documents(docs) + vectorstore.add_documents(chunks) + +# Example usage (uncomment for manual testing): +# if __name__ == "__main__": +# vs = create_vectorstore() +# load_documents("documents", vs) +"""