add vectorstore.py

vectorstore.py

@@ -0,0 +1,45 @@
+import os
+from langchain_chroma import Chroma
+from langchain_ollama import OllamaEmbeddings
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_community.document_loaders import TextLoader, DirectoryLoader
+
+
+def create_vectorstore(persist_directory: str = "./chroma_db"):
+    """Create a ChromaDB vectorstore with Ollama embeddings."""
+    embeddings = OllamaEmbeddings(model="nomic-embed-text")
+    vectorstore = Chroma(
+        persist_directory=persist_directory,
+        embedding_function=embeddings
+    )
+    return vectorstore
+
+
+def load_documents(directory: str, vectorstore: Chroma):
+    """Load documents from directory, chunk them, and add to vectorstore."""
+    # Load all .txt and .md files from directory
+    loader = DirectoryLoader(
+        directory,
+        glob="**/*.{txt,md}",
+        loader_cls=TextLoader,
+        loader_kwargs={"encoding": "utf-8"}
+    )
+    documents = loader.load()
+    
+    # Chunk documents
+    text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size=1000,
+        chunk_overlap=200,
+        length_function=len,
+        is_separator_regex=False
+    )
+    chunks = text_splitter.split_documents(documents)
+    
+    # Add to vectorstore
+    if chunks:
+        vectorstore.add_documents(chunks)
+        print(f"Loaded {len(chunks)} chunks from {len(documents)} documents")
+    else:
+        print("No documents found to load")
+    
+    return vectorstore