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

tools.py

@@ -1,40 +1,42 @@
-"""
-Agent tools for local KB search and web search via Tavily.
-"""
+import os
+from typing import Any
 
-from typing import List, Dict
-
-from langchain.agents import tool
+from langchain.tools import tool
 from langchain_tavily import TavilySearch
-from langchain.schema.document import Document
 
-# Local KB search tool
-@tool("search_local_kb")
-def search_local_kb(query: str, top_k: int = 3) -> List[Dict]:
-    """Semantic search in the local ChromaDB collection.
+from vectorstore import create_vectorstore
 
-    Returns a list of dicts with keys ``text`` and ``source``.
-    """
-    from vectorstore import create_vectorstore
 
-    # Assume the collection is already created and persisted
-    store = create_vectorstore()
-    results = store.query(query_texts=[query], n_results=top_k)
-    return [
-        {"text": doc.page_content, "source": doc.metadata.get("source", "unknown")}
-        for doc in results
-    ]
+def _format_local_docs(results: list[Any]) -> str:
+    if not results:
+        return "Source: chromadb\nNo relevant local documents found."
 
-# Web search tool using Tavily.
-@tool("web_search")
-def web_search(query: str) -> List[Dict]:
-    """Search the web with Tavily and return a list of result snippets.
+    parts = ["Source: chromadb"]
+    for index, doc in enumerate(results, start=1):
+        source = doc.metadata.get("source", "unknown")
+        parts.append(f"{index}. ({source}) {doc.page_content}")
+    return "\n".join(parts)
 
-    Requires environment variable TAVILY_API_KEY.
-    """
-    tav = TavilySearch()
-    results = tav.search(query, max_results=5)
-    return [
-        {"text": r["title"] + ": " + r["snippet"] if isinstance(r, dict) else str(r)}
-        for r in (results if isinstance(results, list) else [])
-    ]
+
+@tool
+def search_local_kb(query: str, top_k: int = 3) -> str:
+    """Search the local ChromaDB knowledge base for internal course notes and local documents."""
+    vectorstore = create_vectorstore()
+    retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
+    results = retriever.invoke(query)
+    return _format_local_docs(results)
+
+
+@tool
+def web_search(query: str) -> str:
+    """Search the web with Tavily for fresh facts, current news, or anything not covered by local documents."""
+    api_key = os.getenv("TAVILY_API_KEY")
+    if not api_key:
+        return "Source: tavily\nTAVILY_API_KEY is not set, so web search is unavailable."
+
+    search_tool = TavilySearch(
+        max_results=5,
+        topic="general",
+    )
+    result = search_tool.invoke({"query": query})
+    return f"Source: tavily\n{result}"