66 lines
2.3 KiB
Python
66 lines
2.3 KiB
Python
"""Vector store utilities using ChromaDB and Ollama embeddings.
|
|
|
|
This module provides functions to create a persistent Chroma vector store and load
|
|
text documents from a directory into it. The store is exposed via the global
|
|
``store`` variable so that other modules (e.g. tools) can access it.
|
|
"""
|
|
|
|
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_core.documents import Document
|
|
|
|
# Global store that will be initialised in ``create_vectorstore``.
|
|
store: Chroma | None = None
|
|
|
|
def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
|
|
"""Create a Chroma vector store with Ollama embeddings.
|
|
|
|
Parameters
|
|
----------
|
|
persist_directory: str
|
|
Directory where the vector data will be persisted.
|
|
"""
|
|
global store
|
|
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
|
store = Chroma(
|
|
collection_name="rag_collection",
|
|
embedding_function=embeddings,
|
|
persist_directory=persist_directory,
|
|
)
|
|
return store
|
|
|
|
def load_documents(directory: str | Path, vectorstore: Chroma) -> None:
|
|
"""Load all .txt and .md files from *directory* into *vectorstore*.
|
|
|
|
The documents are split using ``RecursiveCharacterTextSplitter`` before
|
|
being added to the collection.
|
|
"""
|
|
dir_path = Path(directory)
|
|
txt_files = list(dir_path.rglob("*.txt")) + list(dir_path.rglob("*.md"))
|
|
if not txt_files:
|
|
print(f"No .txt or .md files found in {dir_path}")
|
|
return
|
|
|
|
splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
|
|
docs: List[Document] = []
|
|
for file_path in txt_files:
|
|
text = file_path.read_text(encoding="utf-8")
|
|
docs.extend(splitter.create_documents([text], metadata={"source": str(file_path)}))
|
|
|
|
vectorstore.add_documents(docs)
|
|
# Persist the collection to disk.
|
|
vectorstore.persist()
|
|
print(f"Loaded {len(docs)} documents from {dir_path} into Chroma.")
|
|
|
|
# Helper to get the global store.
|
|
|
|
def get_vectorstore() -> Chroma:
|
|
if store is None:
|
|
raise RuntimeError("Vector store has not been initialised. Call create_vectorstore() first.")
|
|
return store
|
|
|
|
"""End of vectorstore.py""" |