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.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"):
+    base_path = Path(directory)
+    if not base_path.exists():
+        raise FileNotFoundError(f"Directory not found: {directory}")
+
+    raw_docs: list[Document] = []
+    for pattern in ("*.txt", "*.md"):
+        for path in sorted(base_path.rglob(pattern)):
             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)}))
+            raw_docs.append(
+                Document(
+                    page_content=content,
+                    metadata={"source": str(path)},
+                )
+            )
 
-    if docs:
-        # Chroma expects a list of documents via add_documents
-        vectorstore.add_documents(docs)
+    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)