Update rag_tools.py

rag_tools.py

@@ -1,68 +1,91 @@
-"""Two tools used by the RAG agent.
+"""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.
+Two tools are provided:
 
-Both functions are decorated with :func:`langchain.tools.tool` so that they can be used by LangChain
-agents.
+* :func:`search_local_kb` – semantic search in the local ChromaDB vector store.
+* :func:`web_search` – real‑time web search via Tavily.
+
+Both functions are decorated with :func:`langchain.tools.tool` so that LangChain
+can expose them to the agent.
 """
 
+from __future__ import annotations
+
+import os
 from typing import List
 
-from langchain.tools import tool
+from langchain.chains import RetrievalQA
+from langchain.chroma import Chroma
 from langchain_ollama import ChatOllama
-from tavily import TavilyClient
+from langchain_tavily import TavilySearchResults
+from langchain.tools import tool
 
-# 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()
+# Load environment variables (TAVILY_API_KEY)
+from dotenv import load_dotenv
 
+load_dotenv()
 
+# ---------------------------------------------------------------------------
+# Local KB search
+# ---------------------------------------------------------------------------
 @tool("search_local_kb")
-def search_local_kb(query: str, top_k: int = 3, vectorstore=None) -> str:
-    """Semantic search in the local knowledge base.
+def search_local_kb(query: str, top_k: int = 3, vectorstore: Chroma | None = None) -> str:
+    """Return the best answer from the local ChromaDB store.
 
     Parameters
     ----------
     query: str
-        The user query.
-    top_k: int, optional
-        Number of documents to return.
-    vectorstore: Chroma, optional
-        The vector store instance.  It is passed by the agent.
+        The user question.
+    top_k: int
+        Number of documents to retrieve.
+    vectorstore: Chroma | None
+        If ``None`` the function will attempt to load the default store.
 
     Returns
     -------
     str
-        Concatenated text of the retrieved documents.
+        The answer prefixed with the source identifier.
     """
     if vectorstore is None:
-        raise ValueError("vectorstore must be provided")
+        # Load the default store
+        from vectorstore import create_vectorstore
+        vectorstore = create_vectorstore()
+
     retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
-    docs = retriever.get_relevant_documents(query)
-    return "\n\n".join(doc.page_content for doc in docs)
-
+    qa_chain = RetrievalQA.from_chain_type(
+        llm=ChatOllama(model="llama3", temperature=0),
+        chain_type="stuff",
+        retriever=retriever,
+    )
+    result = qa_chain.run(query)
+    return f"[chromadb] {result}"
 
+# ---------------------------------------------------------------------------
+# Web search via Tavily
+# ---------------------------------------------------------------------------
 @tool("web_search")
-def web_search(query: str, max_results: int = 3) -> str:
-    """Perform a web search via Tavily.
+def web_search(query: str, top_k: int = 3) -> str:
+    """Return an answer based on a Tavily web search.
 
     Parameters
     ----------
     query: str
-        The search query.
-    max_results: int, optional
-        Number of search results to return.
+        The user question.
+    top_k: int
+        Number of search results to retrieve.
 
     Returns
     -------
     str
-        Concatenated snippets from the search results.
+        The answer prefixed with the source identifier.
     """
-    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
+    api_key = os.getenv("TAVILY_API_KEY")
+    if not api_key:
+        raise RuntimeError("TAVILY_API_KEY environment variable is not set")
+    tavily = TavilySearchResults(api_key=api_key, top_k=top_k)
+    results = tavily.run(query)
+    snippets = "\n\n".join(item.get("content", "") for item in results)
+    llm = ChatOllama(model="llama3", temperature=0)
+    prompt = f"Answer the question based on the following web snippets:\n{snippets}\n\nQuestion: {query}\nAnswer:"  # noqa: E501
+    answer = llm.invoke(prompt).content
+    return f"[tavily] {answer}"