64 lines
2.2 KiB
Python
64 lines
2.2 KiB
Python
"""Utilities for creating and populating a ChromaDB vector store.
|
||
|
||
This module provides two functions:
|
||
|
||
- :func:`create_vectorstore` – creates a Chroma vector store backed by a local directory.
|
||
- :func:`load_documents` – reads all ``.txt`` and ``.md`` files from a directory, splits them into chunks using
|
||
:class:`langchain_text_splitters.RecursiveCharacterTextSplitter`, and upserts the chunks into the
|
||
provided vector store.
|
||
|
||
The vector store is persistent across runs – the ``persist_directory`` argument defaults to
|
||
``"./chroma_db"``.
|
||
"""
|
||
|
||
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 embedding model used by the vector store
|
||
EMBEDDING_MODEL = "nomic-embed-text"
|
||
|
||
|
||
def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
|
||
"""Create a Chroma vector store.
|
||
|
||
Parameters
|
||
----------
|
||
persist_directory: str
|
||
Directory where the vector store will be persisted.
|
||
|
||
Returns
|
||
-------
|
||
Chroma
|
||
A Chroma vector store instance.
|
||
"""
|
||
embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
|
||
return Chroma(persist_directory=persist_directory, embedding_function=embeddings)
|
||
|
||
|
||
def _read_text_files(directory: str) -> Iterable[str]:
|
||
"""Yield the content of all ``.txt`` and ``.md`` files in *directory*.
|
||
"""
|
||
path = Path(directory)
|
||
for file_path in path.rglob("*.txt"):
|
||
yield file_path.read_text(encoding="utf-8")
|
||
for file_path in path.rglob("*.md"):
|
||
yield file_path.read_text(encoding="utf-8")
|
||
|
||
|
||
def load_documents(directory: str, vectorstore: Chroma, chunk_size: int = 1000, chunk_overlap: int = 200) -> None:
|
||
"""Load documents from *directory* into *vectorstore*.
|
||
|
||
The documents are split into chunks using :class:`RecursiveCharacterTextSplitter` and then
|
||
upserted into the vector store.
|
||
"""
|
||
splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
|
||
for text in _read_text_files(directory):
|
||
chunks = splitter.split_text(text)
|
||
vectorstore.add_texts(chunks)
|
||
|
||
# End of vectorstore.py
|