56 lines
1.8 KiB
Python
56 lines
1.8 KiB
Python
"""
|
|
Vector store utilities for ChromaDB with Ollama embeddings.
|
|
"""
|
|
|
|
import os
|
|
from pathlib import Path
|
|
|
|
from langchain_chroma import Chroma
|
|
from langchain_ollama import OllamaEmbeddings
|
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
|
from langchain.docstore.document import Document
|
|
|
|
def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
|
|
"""Create or load a Chroma vector store.
|
|
|
|
Parameters
|
|
----------
|
|
persist_directory: str
|
|
Directory where the Chroma database is 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 text and markdown files from *directory*, chunk them, and add to *vectorstore*.
|
|
|
|
Parameters
|
|
----------
|
|
directory: str
|
|
Path to the folder containing .txt and .md files.
|
|
vectorstore: Chroma
|
|
The Chroma vector store to which documents will be added.
|
|
"""
|
|
docs = []
|
|
for file_path in Path(directory).rglob("*.*"):
|
|
if file_path.suffix.lower() in ".txt .md".split():
|
|
try:
|
|
text = file_path.read_text(encoding="utf-8")
|
|
except Exception as e:
|
|
print(f"Failed to read {file_path}: {e}")
|
|
continue
|
|
docs.append(Document(page_content=text, metadata={"source": str(file_path)}))
|
|
|
|
if not docs:
|
|
print("No documents found to load.")
|
|
return
|
|
|
|
splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
|
|
chunks = splitter.split_documents(docs)
|
|
vectorstore.add_documents(chunks)
|
|
print(f"Loaded {len(chunks)} chunks into the vector store.") |