80 lines
2.8 KiB
Python
80 lines
2.8 KiB
Python
"""Utilities for creating and populating a Chroma vector store.
|
||
|
||
This module provides two helper functions:
|
||
|
||
* ``create_vectorstore`` – creates a Chroma collection backed by Ollama embeddings.
|
||
* ``load_documents`` – reads ``.txt``/``.md`` files, splits them into chunks and adds them to the collection.
|
||
|
||
The vector store is persisted in ``./chroma_db`` by default.
|
||
"""
|
||
|
||
from pathlib import Path
|
||
from typing import List
|
||
|
||
from langchain_chroma import Chroma
|
||
from langchain_ollama import OllamaEmbeddings
|
||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Configuration
|
||
# ---------------------------------------------------------------------------
|
||
DEFAULT_PERSIST_DIR = "./chroma_db"
|
||
EMBEDDING_MODEL = "nomic-embed-text"
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Public API
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def create_vectorstore(persist_directory: str = DEFAULT_PERSIST_DIR) -> Chroma:
|
||
"""Create a Chroma vector store with Ollama embeddings.
|
||
|
||
Parameters
|
||
----------
|
||
persist_directory: str
|
||
Directory where the vector store will be persisted.
|
||
|
||
Returns
|
||
-------
|
||
Chroma
|
||
A Chroma collection ready for adding documents.
|
||
"""
|
||
embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
|
||
return Chroma(persist_directory=persist_directory, embedding_function=embeddings)
|
||
|
||
|
||
def load_documents(directory: str, vectorstore: Chroma, chunk_size: int = 1000, chunk_overlap: int = 200) -> None:
|
||
"""Load all ``.txt`` and ``.md`` files from *directory*, split them into chunks and add to *vectorstore*.
|
||
|
||
Parameters
|
||
----------
|
||
directory: str
|
||
Path to the folder containing documents.
|
||
vectorstore: Chroma
|
||
The Chroma collection to populate.
|
||
chunk_size: int
|
||
Number of characters per chunk.
|
||
chunk_overlap: int
|
||
Number of characters that overlap between consecutive chunks.
|
||
"""
|
||
splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
|
||
docs: List[str] = []
|
||
|
||
for path in Path(directory).glob("**/*"):
|
||
if path.is_file() and path.suffix.lower() in {".txt", ".md"}:
|
||
text = path.read_text(encoding="utf-8")
|
||
docs.extend(splitter.split_text(text))
|
||
|
||
if docs:
|
||
vectorstore.add_texts(docs)
|
||
else:
|
||
print("[vectorstore] No documents found in", directory)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Example usage (uncomment to run as a script)
|
||
# ---------------------------------------------------------------------------
|
||
# if __name__ == "__main__":
|
||
# store = create_vectorstore()
|
||
# load_documents("documents", store)
|
||
# print("Vector store populated.")
|
||
"""
|