84 lines
2.6 KiB
Python
84 lines
2.6 KiB
Python
"""
|
|
Vector store utilities for ChromaDB with Ollama embeddings.
|
|
|
|
Functions:
|
|
- create_vectorstore(persist_directory="./chroma_db") -> chromadb.Chroma
|
|
- load_documents(directory: str, vectorstore) -> None
|
|
"""
|
|
import os
|
|
from pathlib import Path
|
|
from typing import List
|
|
|
|
from langchain_ollama import OllamaEmbeddings
|
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
|
from langchain_chroma import Chroma
|
|
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, optional
|
|
Directory where the Chroma database is stored. Defaults to ``./chroma_db``.
|
|
|
|
Returns
|
|
-------
|
|
chromadb.Chroma
|
|
A Chroma instance backed by Ollama embeddings.
|
|
"""
|
|
# Ensure directory exists
|
|
Path(persist_directory).mkdir(parents=True, exist_ok=True)
|
|
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
|
return Chroma(persist_directory=persist_directory, embedding_function=embeddings)
|
|
|
|
|
|
def _load_text_files(directory: str) -> List[Document]:
|
|
"""Load all .txt and .md files from *directory* into Documents.
|
|
|
|
Parameters
|
|
----------
|
|
directory : str
|
|
Path to the folder containing documents.
|
|
|
|
Returns
|
|
-------
|
|
list[langchain.docstore.document.Document]
|
|
List of Document objects with ``page_content`` set to file text and
|
|
``metadata`` containing the source path.
|
|
"""
|
|
docs: List[Document] = []
|
|
for root, _, files in os.walk(directory):
|
|
for fname in files:
|
|
if not fname.lower().endswith(('.txt', '.md')):
|
|
continue
|
|
full_path = Path(root) / fname
|
|
text = full_path.read_text(encoding="utf-8")
|
|
docs.append(Document(page_content=text, metadata={"source": str(full_path)}))
|
|
return docs
|
|
|
|
|
|
def load_documents(directory: str, vectorstore: Chroma) -> None:
|
|
"""Chunk documents from *directory* and add them to *vectorstore*.
|
|
|
|
Parameters
|
|
----------
|
|
directory : str
|
|
Folder with .txt/.md files.
|
|
vectorstore : chromadb.Chroma
|
|
The vector store instance returned by ``create_vectorstore``.
|
|
"""
|
|
docs = _load_text_files(directory)
|
|
if not docs:
|
|
return
|
|
splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
|
|
chunks = splitter.split_documents(docs)
|
|
vectorstore.add_documents(chunks)
|
|
|
|
# Example usage (uncomment for manual testing):
|
|
# if __name__ == "__main__":
|
|
# vs = create_vectorstore()
|
|
# load_documents("documents", vs)
|
|
"""
|