54 lines
1.7 KiB
Python
54 lines
1.7 KiB
Python
"""
|
|
Vector store utilities for ChromaDB.
|
|
"""
|
|
|
|
from pathlib import Path
|
|
|
|
from langchain_ollama import OllamaEmbeddings
|
|
from langchain_chroma import Chroma
|
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
|
|
|
|
|
def create_vectorstore(persist_directory: str = "./chroma_db"):
|
|
"""Create or load a Chroma vector store.
|
|
|
|
Parameters
|
|
----------
|
|
persist_directory: str
|
|
Directory where the Chroma database is persisted.
|
|
|
|
Returns
|
|
-------
|
|
Chroma
|
|
The Chroma vector store instance.
|
|
"""
|
|
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
|
return Chroma(persist_directory=persist_directory, embedding_function=embeddings)
|
|
|
|
|
|
def load_documents(directory: str, vectorstore: Chroma, chunk_size: int = 1000, chunk_overlap: int = 200):
|
|
"""Read all .txt and .md files from a directory, split them into chunks and add to the vectorstore.
|
|
|
|
Parameters
|
|
----------
|
|
directory: str
|
|
Path to the folder containing documents.
|
|
vectorstore: Chroma
|
|
The vector store to which documents will be added.
|
|
chunk_size: int
|
|
Maximum number of characters per chunk.
|
|
chunk_overlap: int
|
|
Number of overlapping characters between consecutive chunks.
|
|
"""
|
|
splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
|
|
docs = []
|
|
for file_path in Path(directory).rglob("*.txt"):
|
|
docs.append(file_path.read_text(encoding="utf-8"))
|
|
for file_path in Path(directory).rglob("*.md"):
|
|
docs.append(file_path.read_text(encoding="utf-8"))
|
|
if not docs:
|
|
return
|
|
texts = splitter.split_text("\n\n".join(docs))
|
|
vectorstore.add_texts(texts)
|
|
vectorstore.persist()
|