Update rag_tools.py

rag_tools.py

@@ -1,57 +1,80 @@
-"""Tools for the RAG agent: local KB search and web search via Tavily."""
+"""Tools for the RAG agent.
 
-from typing import List, Dict
+This module defines two LangChain tools:
+
+* ``search_local_kb`` – semantic search in the local ChromaDB vector store.
+* ``web_search`` – real‑time web search using Tavily.
+
+Both tools return a string containing the retrieved information.
+"""
+
+from typing import List
 
 from langchain_ollama import ChatOllama
-from langchain_chroma import Chroma
 from langchain_tavily import TavilySearchResults
 from langchain.tools import tool
 
-# --- Local KB search tool -----------------------------------------------------
+# The LLM used for summarising or formatting responses
+llm = ChatOllama(model="llama3")
 
+# Tavily client – the API key is read from the environment by the package
+# (requires a .env file or the TAVILY_API_KEY environment variable).
+search = TavilySearchResults()
+
+# ---------------------------------------------------------------------------
+# Local knowledge base search tool
+# ---------------------------------------------------------------------------
 @tool("search_local_kb")
-def search_local_kb(query: str, top_k: int = 3, vectorstore: Chroma = None) -> List[Dict]:
-    """Perform a semantic search in the local Chroma vector store.
+
+def search_local_kb(query: str, top_k: int = 3) -> str:
+    """Perform a semantic search in the local ChromaDB vector store.
 
     Parameters
     ----------
     query: str
-        The user query.
-    top_k: int
-        Number of top results to return.
-    vectorstore: Chroma
-        The vector store to search.
+        The user question.
+    top_k: int, optional
+        Number of top documents to return. Defaults to 3.
 
     Returns
     -------
-    List[Dict]
-        List of dictionaries containing ``content`` and ``metadata``.
+    str
+        A formatted string containing the retrieved passages.
     """
-    if vectorstore is None:
-        raise ValueError("vectorstore must be provided")
-    retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
-    docs = retriever.get_relevant_documents(query)
-    return [{"content": doc.page_content, "metadata": doc.metadata} for doc in docs]
+    # The vectorstore is expected to be loaded globally – the agent will
+    # provide it via the tool context.  We simply call the retriever.
+    retriever = globals().get("vectorstore_retriever")
+    if retriever is None:
+        raise RuntimeError("Vector store retriever not configured for the tool.")
 
-# --- Web search tool ---------------------------------------------------------
+    docs = retriever.get_relevant_documents(query, k=top_k)
+    # Concatenate the documents into a single string.
+    passages = "\n\n".join(doc.page_content for doc in docs)
+    return passages
 
+# ---------------------------------------------------------------------------
+# Web search tool
+# ---------------------------------------------------------------------------
 @tool("web_search")
-def web_search(query: str, top_k: int = 3) -> List[Dict]:
-    """Search the web using Tavily.
+
+def web_search(query: str) -> str:
+    """Search the web using Tavily and return the top results.
 
     Parameters
     ----------
     query: str
-        The user query.
-    top_k: int
-        Number of top results to return.
+        The user question.
 
     Returns
     -------
-    List[Dict]
-        List of dictionaries containing ``title``, ``url`` and ``content``.
+    str
+        A formatted string containing the search results.
     """
-    tavily = TavilySearchResults(max_results=top_k)
-    results = tavily.run(query)
-    # Tavily returns a list of dicts with keys: title, url, content
-    return results
+    results = search.run(query)
+    # TavilySearchResults returns a list of dicts with keys: title, url, content
+    formatted = []
+    for r in results:
+        formatted.append(f"Title: {r.get('title', 'N/A')}\nURL: {r.get('url', 'N/A')}\nSnippet: {r.get('content', 'N/A')}\n")
+    return "\n\n".join(formatted)
+
+# End of rag_tools.py