62 lines
2.0 KiB
Python
62 lines
2.0 KiB
Python
"""Vector store utilities using ChromaDB and Ollama embeddings.
|
|
|
|
This module provides functions to create a persistent Chroma vector store and to load documents
|
|
from a directory into the store. Documents are split into chunks using
|
|
`RecursiveCharacterTextSplitter`.
|
|
"""
|
|
|
|
import os
|
|
from pathlib import Path
|
|
from typing import List
|
|
|
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
|
from langchain_chroma import Chroma
|
|
from langchain_ollama import OllamaEmbeddings
|
|
|
|
# Default persistence directory
|
|
DEFAULT_PERSIST_DIR = "./chroma_db"
|
|
|
|
|
|
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 DB will be stored.
|
|
|
|
Returns
|
|
-------
|
|
Chroma
|
|
The Chroma vector store instance.
|
|
"""
|
|
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
|
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* into *vectorstore*.
|
|
|
|
The files are read, split into chunks with a recursive character splitter and
|
|
added to the Chroma collection.
|
|
"""
|
|
splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
|
|
docs: List[str] = []
|
|
for file_path in Path(directory).rglob("*.txt"):
|
|
docs.append(file_path.read_text(encoding="utf-8"))
|
|
for file_path in Path(directory).rglob("*.md"):
|
|
docs.append(file_path.read_text(encoding="utf-8"))
|
|
|
|
if not docs:
|
|
return
|
|
# Split documents into chunks
|
|
texts = splitter.split_text("\n\n".join(docs))
|
|
# Create a list of dicts with metadata (optional)
|
|
metadatas = [{"source": "local"} for _ in texts]
|
|
vectorstore.add_texts(texts, metadatas=metadatas)
|
|
|
|
# Persist changes
|
|
vectorstore.persist()
|
|
|
|
print(f"Loaded {len(texts)} chunks into ChromaDB.")
|