77 lines
2.8 KiB
Python
77 lines
2.8 KiB
Python
"""
|
||
Vector store utilities for the RAG agent.
|
||
|
||
Provides functions to create a ChromaDB vector store backed by Ollama embeddings
|
||
and to load documents from a directory into the store.
|
||
"""
|
||
|
||
from pathlib import Path
|
||
from typing import List
|
||
|
||
from langchain_ollama import OllamaEmbeddings
|
||
from langchain_chroma import Chroma
|
||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||
from langchain.docstore.document import Document
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Configuration constants
|
||
# ---------------------------------------------------------------------------
|
||
DEFAULT_EMBEDDING_MODEL = "nomic-embed-text"
|
||
DEFAULT_PERSIST_DIR = "./chroma_db"
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Public API
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def create_vectorstore(persist_directory: str = DEFAULT_PERSIST_DIR) -> Chroma:
|
||
"""Create (or load) a Chroma vector store.
|
||
|
||
Parameters
|
||
----------
|
||
persist_directory: str
|
||
Directory where the Chroma database will be persisted.
|
||
|
||
Returns
|
||
-------
|
||
Chroma
|
||
An instance of the Chroma vector store.
|
||
"""
|
||
embeddings = OllamaEmbeddings(model=DEFAULT_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*, chunk them and add to *vectorstore*.
|
||
|
||
The function is idempotent – if the same files are loaded again, duplicates will
|
||
not be created because Chroma will deduplicate based on the content hash.
|
||
|
||
Parameters
|
||
----------
|
||
directory: str
|
||
Path to the folder containing the documents.
|
||
vectorstore: Chroma
|
||
The vector store to populate.
|
||
chunk_size: int, optional
|
||
Maximum number of characters per chunk.
|
||
chunk_overlap: int, optional
|
||
Number of characters to overlap between consecutive chunks.
|
||
"""
|
||
splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
|
||
docs: List[Document] = []
|
||
|
||
for file_path in Path(directory).glob("**/*"):
|
||
if file_path.suffix.lower() not in {".txt", ".md"}:
|
||
continue
|
||
text = file_path.read_text(encoding="utf-8")
|
||
docs.extend(splitter.split_text(text))
|
||
|
||
# Convert list of strings to list of Documents
|
||
documents = [Document(page_content=chunk) for chunk in docs]
|
||
|
||
if documents:
|
||
vectorstore.add_documents(documents)
|
||
vectorstore.persist()
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# End of module
|
||
# --------------------------------------------------------------------------- |