29 lines
1.1 KiB
Python
29 lines
1.1 KiB
Python
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
|