45 lines
1.4 KiB
Python
45 lines
1.4 KiB
Python
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
|
|
|
|
|
|
def create_vectorstore(persist_directory: str = "./chroma_db"):
|
|
"""Create a ChromaDB vectorstore with Ollama embeddings."""
|
|
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
|
vectorstore = Chroma(
|
|
persist_directory=persist_directory,
|
|
embedding_function=embeddings
|
|
)
|
|
return vectorstore
|
|
|
|
|
|
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 |