54 lines
1.8 KiB
Python
54 lines
1.8 KiB
Python
"""
|
||
Vector store utilities for the RAG agent.
|
||
|
||
Provides:
|
||
- create_vectorstore(persist_directory)
|
||
- load_documents(directory, vectorstore)
|
||
"""
|
||
import os
|
||
from pathlib import Path
|
||
from typing import List
|
||
|
||
from langchain_ollama import OllamaEmbeddings
|
||
from langchain_chroma import Chroma
|
||
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
||
|
||
|
||
def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
|
||
"""Create or load a Chroma vector store.
|
||
|
||
Parameters
|
||
----------
|
||
persist_directory: str
|
||
Directory where the Chroma DB will be stored. If it does not exist, it is created.
|
||
"""
|
||
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_documents(directory: str, vectorstore: Chroma) -> None:
|
||
"""Load all .txt/.md files from *directory* into the vector store.
|
||
|
||
The function reads each file, splits it with a RecursiveCharacterTextSplitter and adds the chunks to the collection.
|
||
Existing documents are overwritten – this is fine for an init script.
|
||
"""
|
||
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
|
||
docs: List[str] = []
|
||
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:
|
||
print(f"No documents found in {directory}")
|
||
return
|
||
|
||
# Split and add to vectorstore
|
||
chunks = text_splitter.split_documents([{"content": d} for d in docs])
|
||
vectorstore.add_texts([c["content"] for c in chunks])
|
||
print(f"Loaded {len(chunks)} chunks into ChromaDB.")
|