Update rag_tools.py

rag_tools.py

@@ -1,64 +1,68 @@
-"""
-Tools for the RAG agent: local semantic search and web search via Tavily.
+"""Two tools used by the RAG agent.
+
+- :func:`search_local_kb` – performs a semantic search in the local ChromaDB vector store.
+- :func:`web_search` – performs a web search via Tavily.
+
+Both functions are decorated with :func:`langchain.tools.tool` so that they can be used by LangChain
+agents.
 """
 
 from typing import List
 
-from langchain_community.tools.tavily import TavilySearchResults
 from langchain.tools import tool
-from langchain_chroma import Chroma
+from langchain_ollama import ChatOllama
+from tavily import TavilyClient
+
+# The LLM used for generating responses.  We keep a single instance.
+_llm = ChatOllama(model="llama3")
+
+# Tavily client – the API key is read from the environment by the tavily package.
+_tavily_client = TavilyClient()
+
 
-# ---------------------------------------------------------------------------
-# Local KB search tool
-# ---------------------------------------------------------------------------
 @tool("search_local_kb")
-
-def search_local_kb(query: str, top_k: int = 3, vectorstore: Chroma = None) -> str:
-    """Perform a semantic search in the local Chroma vector store.
+def search_local_kb(query: str, top_k: int = 3, vectorstore=None) -> str:
+    """Semantic search in the local knowledge base.
 
     Parameters
     ----------
     query: str
         The user query.
     top_k: int, optional
-        Number of top results to return.
+        Number of documents to return.
     vectorstore: Chroma, optional
-        The vector store to query. If None, the function will raise an error.
+        The vector store instance.  It is passed by the agent.
 
     Returns
     -------
     str
-        Concatenated content of the top results.
+        Concatenated text of the retrieved documents.
     """
     if vectorstore is None:
-        raise ValueError("vectorstore must be provided to search_local_kb")
+        raise ValueError("vectorstore must be provided")
     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)
+    return "\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.
+def web_search(query: str, max_results: int = 3) -> str:
+    """Perform a web search via Tavily.
 
     Parameters
     ----------
     query: str
         The search query.
-    top_k: int, optional
-        Number of top results to return.
+    max_results: int, optional
+        Number of search results to return.
 
     Returns
     -------
     str
-        Formatted search results.
+        Concatenated snippets from the search results.
     """
-    tavily = TavilySearchResults(max_results=top_k)
-    results = tavily.run(query)
-    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)
+    results = _tavily_client.search(query, max_results=max_results)
+    snippets = [f"{res.title}\n{res.url}\n{res.content}" for res in results]
+    return "\n\n".join(snippets)
+
+# End of rag_tools.py