52 lines
1.8 KiB
Python
52 lines
1.8 KiB
Python
"""Vector store utilities for ChromaDB.
|
|
|
|
This module provides functions to create a persistent ChromaDB vector store using
|
|
Ollama embeddings and to load documents from a directory into the store.
|
|
"""
|
|
|
|
from pathlib import Path
|
|
|
|
from langchain_chroma import Chroma
|
|
from langchain_ollama import OllamaEmbeddings
|
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
|
from langchain_core.documents import Document
|
|
|
|
|
|
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 store will be persisted.
|
|
|
|
Returns
|
|
-------
|
|
Chroma
|
|
A 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) -> None:
|
|
"""Load .txt and .md files from *directory* into *vectorstore*.
|
|
|
|
The documents are split into chunks using ``RecursiveCharacterTextSplitter``
|
|
before being added to the vector store.
|
|
"""
|
|
splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
|
|
docs = []
|
|
for file_path in Path(directory).glob("*"):
|
|
if file_path.suffix.lower() not in {".txt", ".md"}:
|
|
continue
|
|
with open(file_path, "r", encoding="utf-8") as f:
|
|
text = f.read()
|
|
chunks = splitter.split_text(text)
|
|
docs.extend([Document(page_content=c, metadata={"source": str(file_path)}) for c in chunks])
|
|
if docs:
|
|
vectorstore.add_documents(docs)
|
|
print(f"Loaded {len(docs)} chunks from {directory} into ChromaDB.")
|
|
else:
|
|
print(f"No .txt/.md files found in {directory}.")
|