From 2e60c30dbfc91bde846c536ce8a5685e795c9b0c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=A0=D0=B8=D0=BD=D0=B0=D1=80=20=D0=9C=D0=B8=D1=80=D0=B7?= =?UTF-8?q?=D0=B0=D0=B3=D0=B8=D1=82=D0=BE=D0=B2?= Date: Tue, 16 Jun 2026 13:13:03 +0000 Subject: [PATCH] Build ChromaDB + Tavily RAG agent with Ollama embeddings, local/web tools, create_agent routing, and CLI ingest flow.: update vectorstore.py --- vectorstore.py | 81 +++++++++++++++++++++++++++----------------------- 1 file changed, 43 insertions(+), 38 deletions(-) diff --git a/vectorstore.py b/vectorstore.py index 839c3e9..8e35510 100644 --- a/vectorstore.py +++ b/vectorstore.py @@ -1,50 +1,55 @@ -""" -RAG vector store using ChromaDB and Ollama embeddings. -""" - from pathlib import Path -from typing import List +from uuid import uuid4 -import chromadb -from langchain.embeddings.ollama import OllamaEmbeddings -from langchain.text_splitter import RecursiveCharacterTextSplitter -from langchain.schema.document import Document -from langchain.vectorstores import Chroma +from langchain_chroma import Chroma +from langchain_core.documents import Document +from langchain_ollama import OllamaEmbeddings +from langchain_text_splitters import RecursiveCharacterTextSplitter CHROMA_DIR = "./chroma_db" +COLLECTION_NAME = "local_kb" EMBED_MODEL = "nomic-embed-text" +OLLAMA_BASE_URL = "http://127.0.0.1:11434" def create_vectorstore(persist_directory: str = CHROMA_DIR) -> Chroma: - """Create or load a Chroma vector store. - - Parameters - ---------- - persist_directory : str, optional - Directory where the Chroma database is stored. If it does not exist, - it will be created automatically by Chroma. - """ - embeddings = OllamaEmbeddings(model=EMBED_MODEL) - # Chroma can use a local directory for persistence - client = chromadb.PersistentClient(path=persist_directory) - collection = client.get_or_create_collection(name="documents", embedding_function=embeddings) - return Chroma(collection=collection, embedding_function=embeddings) + embeddings = OllamaEmbeddings( + model=EMBED_MODEL, + base_url=OLLAMA_BASE_URL, + ) + return Chroma( + collection_name=COLLECTION_NAME, + embedding_function=embeddings, + persist_directory=persist_directory, + ) -def load_documents(directory: str, vectorstore: Chroma) -> None: - """Load all .txt and .md files from *directory*, chunk them and add to the vector store. +def load_documents(directory: str, vectorstore: Chroma) -> int: + splitter = RecursiveCharacterTextSplitter( + chunk_size=1000, + chunk_overlap=200, + add_start_index=True, + ) - The function does not return anything; it mutates the provided collection. - """ - text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) - docs: List[Document] = [] - for path in Path(directory).rglob("*.txt"): - content = path.read_text(encoding="utf-8") - docs.extend(text_splitter.create_documents([content], metadata={"source": str(path)})) - for path in Path(directory).rglob("*.md"): - content = path.read_text(encoding="utf-8") - docs.extend(text_splitter.create_documents([content], metadata={"source": str(path)})) + base_path = Path(directory) + if not base_path.exists(): + raise FileNotFoundError(f"Directory not found: {directory}") - if docs: - # Chroma expects a list of documents via add_documents - vectorstore.add_documents(docs) + raw_docs: list[Document] = [] + for pattern in ("*.txt", "*.md"): + for path in sorted(base_path.rglob(pattern)): + content = path.read_text(encoding="utf-8") + raw_docs.append( + Document( + page_content=content, + metadata={"source": str(path)}, + ) + ) + + if not raw_docs: + return 0 + + chunks = splitter.split_documents(raw_docs) + ids = [str(uuid4()) for _ in chunks] + vectorstore.add_documents(documents=chunks, ids=ids) + return len(chunks)