79 lines
2.4 KiB
Python
79 lines
2.4 KiB
Python
"""Vector store utilities using ChromaDB and Ollama embeddings.
|
||
|
||
This module provides functions to create a persistent ChromaDB vector store and
|
||
load documents from a directory into it. The documents are split into
|
||
manageable chunks using a recursive character text splitter.
|
||
"""
|
||
|
||
from pathlib import Path
|
||
from typing import Iterable
|
||
|
||
from langchain_chroma import Chroma
|
||
from langchain_ollama import OllamaEmbeddings
|
||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||
|
||
# Default persistence directory
|
||
DEFAULT_PERSIST_DIR = "./chroma_db"
|
||
|
||
# Default chunking parameters – these can be tuned
|
||
DEFAULT_CHUNK_SIZE = 1000
|
||
DEFAULT_CHUNK_OVERLAP = 200
|
||
|
||
|
||
def create_vectorstore(persist_directory: str = DEFAULT_PERSIST_DIR) -> Chroma:
|
||
"""Create or load a Chroma vector store.
|
||
|
||
Parameters
|
||
----------
|
||
persist_directory: str
|
||
Directory where the ChromaDB files are stored.
|
||
|
||
Returns
|
||
-------
|
||
Chroma
|
||
An instance of the Chroma vector store backed by the given directory.
|
||
"""
|
||
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
||
vectorstore = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
|
||
return vectorstore
|
||
|
||
|
||
def _load_text_files(directory: str | Path) -> Iterable[str]:
|
||
"""Yield the contents of all .txt and .md files in *directory*.
|
||
|
||
Parameters
|
||
----------
|
||
directory: str | Path
|
||
Path to the directory containing the documents.
|
||
"""
|
||
directory = Path(directory)
|
||
for file_path in directory.rglob("*.txt"):
|
||
yield file_path.read_text(encoding="utf-8")
|
||
for file_path in directory.rglob("*.md"):
|
||
yield file_path.read_text(encoding="utf-8")
|
||
|
||
|
||
def load_documents(directory: str | Path, vectorstore: Chroma) -> None:
|
||
"""Load documents from *directory* into *vectorstore*.
|
||
|
||
The documents are split into chunks using a recursive character splitter
|
||
and then added to the vector store. Existing documents are not removed –
|
||
this function simply appends new data.
|
||
|
||
Parameters
|
||
----------
|
||
directory: str | Path
|
||
Directory containing the source documents.
|
||
vectorstore: Chroma
|
||
The vector store to populate.
|
||
"""
|
||
splitter = RecursiveCharacterTextSplitter(chunk_size=DEFAULT_CHUNK_SIZE, chunk_overlap=DEFAULT_CHUNK_OVERLAP)
|
||
for text in _load_text_files(directory):
|
||
chunks = splitter.split_text(text)
|
||
vectorstore.add_texts(chunks)
|
||
|
||
# Persist changes
|
||
vectorstore.persist()
|
||
|
||
# End of vectorstore.py
|