82 lines
2.5 KiB
Python
82 lines
2.5 KiB
Python
"""
|
|
Vector store utilities using ChromaDB and Ollama embeddings.
|
|
"""
|
|
|
|
import os
|
|
from pathlib import Path
|
|
from typing import List
|
|
|
|
from langchain_chroma import Chroma
|
|
from langchain_ollama import OllamaEmbeddings
|
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
|
from langchain_ollama import ChatOllama
|
|
|
|
# Create the vector store with persistence
|
|
|
|
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 will be persisted.
|
|
|
|
Returns
|
|
-------
|
|
Chroma
|
|
The Chroma vector store instance.
|
|
"""
|
|
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
|
return Chroma(persist_directory=persist_directory, embedding_function=embeddings)
|
|
|
|
# Load documents from a directory and add them to the vector store
|
|
|
|
def load_documents(directory: str, vectorstore: Chroma) -> None:
|
|
"""Load .txt and .md files from *directory*, chunk them, and add to *vectorstore*.
|
|
|
|
Parameters
|
|
----------
|
|
directory: str
|
|
Path to the directory containing the documents.
|
|
vectorstore: Chroma
|
|
The vector store to which the documents will be added.
|
|
"""
|
|
# Ensure the directory exists
|
|
path = Path(directory)
|
|
if not path.is_dir():
|
|
raise FileNotFoundError(f"Directory {directory} does not exist")
|
|
|
|
# Collect all .txt and .md files
|
|
files = list(path.rglob("*.txt")) + list(path.rglob("*.md"))
|
|
if not files:
|
|
print(f"No .txt or .md files found in {directory}")
|
|
return
|
|
|
|
# Read and chunk the documents
|
|
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
|
|
documents = []
|
|
for file_path in files:
|
|
try:
|
|
content = file_path.read_text(encoding="utf-8")
|
|
except Exception as e:
|
|
print(f"Failed to read {file_path}: {e}")
|
|
continue
|
|
# Split into chunks
|
|
chunks = text_splitter.split_text(content)
|
|
for i, chunk in enumerate(chunks):
|
|
documents.append({
|
|
"page_content": chunk,
|
|
"metadata": {"source": str(file_path), "chunk_index": i},
|
|
})
|
|
|
|
if documents:
|
|
vectorstore.add_documents(documents)
|
|
print(f"Added {len(documents)} chunks to the vector store from {directory}")
|
|
else:
|
|
print("No documents were processed.")
|
|
|
|
# Example usage:
|
|
# if __name__ == "__main__":
|
|
# store = create_vectorstore()
|
|
# load_documents("./documents", store)
|