Add vectorstore.py
This commit is contained in:
@@ -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
|
||||||
Reference in New Issue
Block a user