83 lines
2.9 KiB
Python
83 lines
2.9 KiB
Python
"""
|
|
Vector store utilities for ChromaDB.
|
|
"""
|
|
|
|
import os
|
|
from pathlib import Path
|
|
|
|
from langchain_chroma import Chroma
|
|
from langchain_ollama import OllamaEmbeddings
|
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
|
from langchain_community.document_loaders import TextLoader, UnstructuredMarkdownLoader
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Create a persistent Chroma vector store
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
|
|
"""Create or load a Chroma vector store.
|
|
|
|
Parameters
|
|
----------
|
|
persist_directory: str
|
|
Directory where the Chroma DB will be persisted.
|
|
|
|
Returns
|
|
-------
|
|
Chroma
|
|
A Chroma vector store instance.
|
|
"""
|
|
os.makedirs(persist_directory, exist_ok=True)
|
|
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
|
return Chroma(persist_directory=persist_directory, embedding_function=embeddings)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Load documents from a directory and add them to the vector store
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def load_documents(directory: str, vectorstore: Chroma, chunk_size: int = 500, chunk_overlap: int = 50) -> None:
|
|
"""Load `.txt` and `.md` files from *directory*, split them into chunks, and add to *vectorstore*.
|
|
|
|
Parameters
|
|
----------
|
|
directory: str
|
|
Directory containing the source documents.
|
|
vectorstore: Chroma
|
|
The vector store to which documents will be added.
|
|
chunk_size: int, optional
|
|
Maximum chunk size in characters.
|
|
chunk_overlap: int, optional
|
|
Number of overlapping characters between consecutive chunks.
|
|
"""
|
|
loader_classes = {
|
|
".txt": TextLoader,
|
|
".md": UnstructuredMarkdownLoader,
|
|
}
|
|
splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
|
|
|
|
docs = []
|
|
for root, _, files in os.walk(directory):
|
|
for file in files:
|
|
ext = Path(file).suffix.lower()
|
|
if ext not in loader_classes:
|
|
continue
|
|
loader = loader_classes[ext](os.path.join(root, file))
|
|
loaded_docs = loader.load()
|
|
docs.extend(loaded_docs)
|
|
|
|
if not docs:
|
|
return
|
|
|
|
# Split documents into smaller chunks
|
|
split_docs = splitter.split_documents(docs)
|
|
vectorstore.add_documents(split_docs)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Example usage
|
|
# ---------------------------------------------------------------------------
|
|
if __name__ == "__main__":
|
|
# This block is only executed when running the module directly.
|
|
store = create_vectorstore()
|
|
load_documents("documents", store)
|
|
print("Vector store populated.")
|