diff --git a/vectorstore.py b/vectorstore.py new file mode 100644 index 0000000..cee9c49 --- /dev/null +++ b/vectorstore.py @@ -0,0 +1,28 @@ +import os +from langchain_ollama import OllamaEmbeddings +from langchain_chroma import Chroma +from langchain.text_splitter import RecursiveCharacterTextSplitter + +class VectorStore: + def __init__(self, persist_path: str = "vectorstore"): + self.persist_path = persist_path + os.makedirs(self.persist_path, exist_ok=True) + embeddings = OllamaEmbeddings(model="nomic-embed-text") + self.store = Chroma(persist_directory=self.persist_path, embedding_function=embeddings) + + def add_documents(self, docs_dir: str): + texts = [] + for root, _, files in os.walk(docs_dir): + for fname in files: + if fname.lower().endswith((".txt", ".md")): + path = os.path.join(root, fname) + with open(path, "r", encoding="utf-8") as f: + content = f.read() + texts.append(content) + splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) + chunks = splitter.split_text("\n".join(texts)) + self.store.add_texts(chunks) + self.store.persist() + + def get_store(self): + return self.store