Update rag_tools.py

rag_tools.py

@@ -1,63 +1,88 @@
 """Tools used by the RAG agent.
 
-Two tools are provided:
+This module defines two LangChain tools:
 
-* ``search_local_kb`` – semantic search in the local ChromaDB store.
-* ``web_search`` – real‑time web search via Tavily.
+1. `search_local_kb` – semantic search in the Chroma vector store.
+2. `web_search` – web search using Tavily.
+
+Both tools are decorated with `@tool` so that they can be used by the agent.
 """
 
 from typing import List
 
+from langchain_core.tools import tool
 from langchain_ollama import ChatOllama
 from langchain_tavily import TavilySearchResults
 from langchain_chroma import Chroma
-from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain.docstore.document import Document
 
-from langchain.tools import tool
-
-# The LLM used for generating answers.  Using the same model as the embeddings
-# keeps the pipeline consistent.
+# Global LLM instance for tool responses (can be reused)
 _llm = ChatOllama(model="llama3")
 
-# Tavily client – the API key is read from the environment variable
-# ``TAVILY_API_KEY`` by the TavilySearchResults class.
-_tavily = TavilySearchResults()
-
+# ---------------------------------------------------------------------------
+# Local KB search tool
+# ---------------------------------------------------------------------------
 
 @tool("search_local_kb")
+
 def search_local_kb(query: str, top_k: int = 3, vectorstore: Chroma = None) -> str:
-    """Semantic search in the local ChromaDB vector store.
+    """Search the local Chroma vector store for relevant chunks.
 
-    Args:
-        query: User question.
-        top_k: Number of results to return.
-        vectorstore: The Chroma instance to query.
+    Parameters
+    ----------
+    query: str
+        The user's query.
+    top_k: int, optional
+        Number of top results to return.
+    vectorstore: Chroma
+        The Chroma vector store instance.
 
-    Returns:
-        A string containing the concatenated top results.
+    Returns
+    -------
+    str
+        A formatted string containing the retrieved chunks.
     """
     if vectorstore is None:
-        raise ValueError("vectorstore must be provided")
+        raise ValueError("Vectorstore must be provided to search_local_kb tool.")
     retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
-    docs: List[Document] = retriever.invoke(query)
-    return "\n\n".join(doc.page_content for doc in docs)
+    docs = retriever.invoke(query)
+    # docs is a list of Document objects
+    if not docs:
+        return "No relevant information found in the local knowledge base."
+    # Concatenate the content of the top documents
+    snippets = [f"{i+1}. {doc.page_content[:200]}..." for i, doc in enumerate(docs)]
+    return "\n".join(snippets)
 
+# ---------------------------------------------------------------------------
+# Web search tool
+# ---------------------------------------------------------------------------
 
 @tool("web_search")
+
 def web_search(query: str, top_k: int = 3) -> str:
-    """Search the web using Tavily.
+    """Perform a web search using Tavily.
 
-    Args:
-        query: User question.
-        top_k: Number of results to return.
+    Parameters
+    ----------
+    query: str
+        The user's query.
+    top_k: int, optional
+        Number of results to return.
 
-    Returns:
-        Concatenated snippets from the search results.
+    Returns
+    -------
+    str
+        A formatted string containing the search results.
     """
-    results = _tavily.run(query, max_results=top_k)
-    # TavilySearchResults returns a list of dicts with keys like 'title',
-    # 'content', 'url'.  We return the content for simplicity.
-    return "\n\n".join(r.get("content", "") for r in results)
+    tavily = TavilySearchResults(tavily_api_key=None, max_results=top_k)
+    results = tavily.invoke(query)
+    if not results:
+        return "No results found on the web."
+    snippets = [f"{i+1}. {res['title']} – {res['url']}" for i, res in enumerate(results)]
+    return "\n".join(snippets)
 
-*** End of File ***
+# ---------------------------------------------------------------------------
+# Exported tool names for agent
+# ---------------------------------------------------------------------------
+
+TOOLS = [search_local_kb, web_search]
+""