Update rag_tools.py

rag_tools.py

@@ -1,36 +1,64 @@
-"""Tool definitions for the RAG agent.
-
-Provides two tools:
-- search_local_kb: semantic search over the local ChromaDB vector store.
-- web_search: web search via Tavily.
+"""
+Tools for the RAG agent: local semantic search and web search via Tavily.
 """
 
-from typing import List, Dict
+from typing import List
 
+from langchain_community.tools.tavily import TavilySearchResults
 from langchain.tools import tool
-from langchain_ollama import ChatOllama
-from langchain_tavily import TavilySearchResults
+from langchain_chroma import Chroma
 
-# Local search tool will be created dynamically in agent.py because it needs the vectorstore.
+# ---------------------------------------------------------------------------
+# Local KB search tool
+# ---------------------------------------------------------------------------
+@tool("search_local_kb")
 
-@tool
-def web_search(query: str) -> str:
-    """Search the web using Tavily and return a short summary.
+def search_local_kb(query: str, top_k: int = 3, vectorstore: Chroma = None) -> str:
+    """Perform a semantic search in the local Chroma vector store.
+
+    Parameters
+    ----------
+    query: str
+        The user query.
+    top_k: int, optional
+        Number of top results to return.
+    vectorstore: Chroma, optional
+        The vector store to query. If None, the function will raise an error.
+
+    Returns
+    -------
+    str
+        Concatenated content of the top results.
+    """
+    if vectorstore is None:
+        raise ValueError("vectorstore must be provided to search_local_kb")
+    retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
+    docs = retriever.get_relevant_documents(query)
+    return "\n\n---\n\n".join(doc.page_content for doc in docs)
+
+# ---------------------------------------------------------------------------
+# Web search tool using Tavily
+# ---------------------------------------------------------------------------
+@tool("web_search")
+
+def web_search(query: str, top_k: int = 3) -> str:
+    """Search the web via Tavily and return a formatted string of results.
 
     Parameters
     ----------
     query: str
         The search query.
+    top_k: int, optional
+        Number of top results to return.
 
     Returns
     -------
     str
-        A concise answer with a source tag.
+        Formatted search results.
     """
-    tavily = TavilySearchResults(max_results=3, api_key=None)  # API key is taken from env
+    tavily = TavilySearchResults(max_results=top_k)
     results = tavily.run(query)
-    # Build a simple summary from the results
-    summary = "\n".join([f"{idx+1}. {r['title']}: {r['content'][:200]}" for idx, r in enumerate(results)])
-    return f"[Web Search]\n{summary}\nSource: tavily"
-
-# The local search tool will be defined in agent.py where the vectorstore is available.
+    formatted = []
+    for i, r in enumerate(results, 1):
+        formatted.append(f"{i}. {r.get('title', 'No title')}\n{r.get('url', '')}\n{r.get('content', '')}")
+    return "\n\n---\n\n".join(formatted)