diff --git a/vectorstore.py b/vectorstore.py new file mode 100644 index 0000000..d81e0b3 --- /dev/null +++ b/vectorstore.py @@ -0,0 +1,22 @@ +import os +from pathlib import Path +from langchain.vectorstores import Chroma +from langchain.embeddings.ollama import OllamaEmbeddings +from langchain.document_loaders import TextLoader, MarkdownLoader +from langchain.text_splitter import RecursiveCharacterTextSplitter + +def create_vectorstore(persist_directory: str = "./chroma_db"): + embeddings = OllamaEmbeddings(model="nomic-embed-text") + return Chroma(embedding_function=embeddings, persist_directory=persist_directory) + +def load_documents(directory: str, vectorstore): + docs = [] + for file in Path(directory).glob("**/*.*"): + if file.suffix.lower() not in {".txt", ".md"}: + continue + loader = TextLoader(str(file)) if file.suffix == ".txt" else MarkdownLoader(str(file)) + docs.extend(loader.load()) + splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50) + texts = splitter.split_documents(docs) + vectorstore.add_documents(texts) + vectorstore.persist()