61 lines
2.3 KiB
Python
61 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 to load documents from a directory into the store.
|
|
"""
|
|
|
|
from pathlib import Path
|
|
from typing import Iterable
|
|
|
|
from langchain_chroma import Chroma
|
|
from langchain_ollama import OllamaEmbeddings
|
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Create a persistent Chroma vector store.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
|
|
"""Create a Chroma vector store with Ollama embeddings.
|
|
|
|
Parameters
|
|
----------
|
|
persist_directory: str
|
|
Directory where the Chroma database will be stored.
|
|
|
|
Returns
|
|
-------
|
|
Chroma
|
|
A Chroma vector store instance.
|
|
"""
|
|
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
|
return Chroma(persist_directory=persist_directory, embedding_function=embeddings)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Load documents from a directory into the vector store.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def load_documents(directory: str, vectorstore: Chroma) -> None:
|
|
"""Load all .txt and .md files from *directory* into *vectorstore*.
|
|
|
|
The documents are split into chunks using a RecursiveCharacterTextSplitter
|
|
before being added to the vector store.
|
|
"""
|
|
splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
|
|
docs = []
|
|
for path in Path(directory).rglob("*.txt"):
|
|
docs.append(path.read_text(encoding="utf-8"))
|
|
for path in Path(directory).rglob("*.md"):
|
|
docs.append(path.read_text(encoding="utf-8"))
|
|
if not docs:
|
|
return
|
|
# Split the documents into chunks.
|
|
chunks = splitter.split_text("\n\n".join(docs))
|
|
# Create LangChain Document objects.
|
|
from langchain.docstore.document import Document
|
|
documents = [Document(page_content=chunk) for chunk in chunks]
|
|
vectorstore.add_documents(documents)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# End of vectorstore.py
|
|
# --------------------------------------------------------------------------- |