94 lines
3.3 KiB
Python
94 lines
3.3 KiB
Python
"""Vector store utilities for ChromaDB with Ollama embeddings.
|
|
|
|
This module provides functions to create a persistent Chroma vector store and
|
|
load documents from a directory into it. Documents are split into chunks using
|
|
`RecursiveCharacterTextSplitter` and stored in the Chroma collection.
|
|
"""
|
|
|
|
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
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Vector store creation
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
|
|
"""Create or load a Chroma vector store.
|
|
|
|
Parameters
|
|
----------
|
|
persist_directory: str
|
|
Directory where the Chroma DB files are stored.
|
|
|
|
Returns
|
|
-------
|
|
Chroma
|
|
A Chroma vector store instance.
|
|
"""
|
|
# Ensure directory exists
|
|
Path(persist_directory).mkdir(parents=True, exist_ok=True)
|
|
# Use Ollama embeddings
|
|
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
|
# Create Chroma store
|
|
vectorstore = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
|
|
return vectorstore
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Document loading
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _load_text_files(directory: str) -> List[str]:
|
|
"""Load all .txt and .md files from a directory into a list of strings."""
|
|
texts = []
|
|
for root, _, files in os.walk(directory):
|
|
for file in files:
|
|
if file.lower().endswith(('.txt', '.md')):
|
|
path = Path(root) / file
|
|
try:
|
|
content = path.read_text(encoding="utf-8")
|
|
texts.append(content)
|
|
except Exception as e:
|
|
print(f"Failed to read {path}: {e}")
|
|
return texts
|
|
|
|
|
|
def load_documents(directory: str, vectorstore: Chroma, chunk_size: int = 1000, chunk_overlap: int = 200) -> None:
|
|
"""Load documents from a directory into the provided vector store.
|
|
|
|
Parameters
|
|
----------
|
|
directory: str
|
|
Path to the directory containing .txt/.md files.
|
|
vectorstore: Chroma
|
|
The vector store to add documents to.
|
|
chunk_size: int, optional
|
|
Maximum size of each chunk.
|
|
chunk_overlap: int, optional
|
|
Number of characters to overlap between chunks.
|
|
"""
|
|
texts = _load_text_files(directory)
|
|
if not texts:
|
|
print("No text files found in the directory.")
|
|
return
|
|
|
|
splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
|
|
docs = []
|
|
for text in texts:
|
|
docs.extend(splitter.split_text(text))
|
|
|
|
# Add documents to Chroma
|
|
vectorstore.add_texts(docs)
|
|
print(f"Loaded {len(docs)} chunks into the vector store.")
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Example usage (uncomment to run directly)
|
|
# ---------------------------------------------------------------------------
|
|
# if __name__ == "__main__":
|
|
# store = create_vectorstore()
|
|
# load_documents("documents", store)
|
|
"" |