Build ChromaDB + Tavily RAG agent with Ollama embeddings, local/web tools, create_agent routing, and CLI ingest flow.: update vectorstore.py

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_chroma import Chroma
-from langchain.embeddings.ollama import OllamaEmbeddings
+from langchain_core.documents import Document
-from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_ollama import OllamaEmbeddings
-from langchain.schema.document import Document
+from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain.vectorstores import Chroma
 
 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.
+    embeddings = OllamaEmbeddings(
-
+        model=EMBED_MODEL,
-    Parameters
+        base_url=OLLAMA_BASE_URL,
-    ----------
+    )
-    persist_directory : str, optional
+    return Chroma(
-        Directory where the Chroma database is stored. If it does not exist,
+        collection_name=COLLECTION_NAME,
-        it will be created automatically by Chroma.
+        embedding_function=embeddings,
-    """
+        persist_directory=persist_directory,
-    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)
 
 
-def load_documents(directory: str, vectorstore: Chroma) -> None:
+def load_documents(directory: str, vectorstore: Chroma) -> int:
-    """Load all .txt and .md files from *directory*, chunk them and add to the vector store.
+    splitter = RecursiveCharacterTextSplitter(
+        chunk_size=1000,
+        chunk_overlap=200,
+        add_start_index=True,
+    )
 
-    The function does not return anything; it mutates the provided collection.
+    base_path = Path(directory)
-    """
+    if not base_path.exists():
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+        raise FileNotFoundError(f"Directory not found: {directory}")
-    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)}))
 
-    if docs:
+    raw_docs: list[Document] = []
-        # Chroma expects a list of documents via add_documents
+    for pattern in ("*.txt", "*.md"):
-        vectorstore.add_documents(docs)
+        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)