add vectorstore.py

vectorstore.py

@@ -1,45 +1,26 @@
 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
+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"):
-    """Create a ChromaDB vectorstore with Ollama embeddings."""
+    if not os.path.exists(persist_directory):
+        os.makedirs(persist_directory, exist_ok=True)
     embeddings = OllamaEmbeddings(model="nomic-embed-text")
-    vectorstore = Chroma(
-        persist_directory=persist_directory,
-        embedding_function=embeddings
-    )
-    return vectorstore
+    return Chroma(persist_directory=persist_directory, embedding_function=embeddings)
 
-
-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
+def load_documents(directory: str, vectorstore) -> None:
+    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    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)