Update tools.py

tools.py

@@ -1,54 +1,83 @@
-"""
+"""Tool implementations for the RAG agent.
-Agent tools for local KB search and web search via Tavily.
+
+This module defines two tools:
+
+* ``search_local_kb`` – semantic search in the local Chroma vector store.
+* ``web_search`` – web search using Tavily.
+
+Both tools are decorated with ``@tool`` so that LangChain can call them automatically.
 """
 
+import os
 from typing import List
 
-from langchain_ollama import ChatOllama
+from langchain_core.documents import Document
-from langchain.tools import tool
 from langchain_tavily import TavilySearchResults
+from langchain.tools import tool
+
+from vectorstore import get_vectorstore
 
 # ---------------------------------------------------------------------------
 # Local KB search tool
 # ---------------------------------------------------------------------------
-@tool
+@tool("search_local_kb")
+# The docstring becomes the tool description used by the model.
+# The function signature must be type annotated.
+# ``top_k`` is optional with default 3.
+# The function returns a string containing the retrieved snippets and a
+# source tag so that the agent can report where the information came from.
+
 def search_local_kb(query: str, top_k: int = 3) -> str:
-    """Perform a semantic search in the local Chroma vector store.
+    """Search the local knowledge base for *query*.
 
-    The function expects a global variable ``vectorstore`` to be defined in the
+    Parameters
-    module that imports this tool.  This is a simple design choice that keeps
+    ----------
-    the tool stateless and easy to use from an agent.
+    query: str
+        The search query.
+    top_k: int, optional
+        Number of top results to return.
+
+    Returns
+    -------
+    str
+        A formatted string with the retrieved snippets and a source tag.
     """
-    # The vectorstore is assumed to be a global variable set by the caller.
+    store = get_vectorstore()
-    global vectorstore
+    retriever = store.as_retriever(search_kwargs={"k": top_k})
-    if vectorstore is None:
+    docs: List[Document] = retriever.invoke(query)
-        raise RuntimeError("Vectorstore not initialized. Call create_vectorstore first.")
-
-    retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
-    docs = retriever.get_relevant_documents(query)
     if not docs:
-        return "No relevant documents found."
+        return f"No local KB results for '{query}'."
-
+    # Build a readable answer.
-    # Concatenate the content of the retrieved documents.
+    snippets = "\n\n".join([f"- {doc.page_content[:200]}…" for doc in docs])
-    return "\n---\n".join(doc.page_content for doc in docs)
+    return f"[Local KB] {snippets}\nSource: chromadb"
 
 # ---------------------------------------------------------------------------
-# Web search tool via Tavily
+# Web search tool
 # ---------------------------------------------------------------------------
-@tool
+@tool("web_search")
+
 def web_search(query: str) -> str:
-    """Search the web using Tavily and return a concise summary.
+    """Search the web using Tavily.
 
-    The Tavily client requires the environment variable ``TAVILY_API_KEY``.
+    Parameters
+    ----------
+    query: str
+        The search query.
+
+    Returns
+    -------
+    str
+        A formatted string with the top results and a source tag.
     """
-    tavily = TavilySearchResults(max_results=3)
+    # TavilySearchResults requires the API key to be set in the environment.
-    results = tavily.run(query)
+    api_key = os.getenv("TAVILY_API_KEY")
-    # results is a list of dicts with keys: title, url, content
+    if not api_key:
+        return "Tavily API key not set. Please set TAVILY_API_KEY environment variable."
+    tavily = TavilySearchResults(api_key=api_key, max_results=3)
+    results = tavily.invoke(query)
     if not results:
-        return "No web results found."
+        return f"No web results for '{query}'."
-    snippets = [f"{r['title']}\n{r['content'][:300]}" for r in results]
+    snippets = "\n\n".join([f"- {res['title']}: {res['content'][:200]}…" for res in results])
-    return "\n---\n".join(snippets)
+    return f"[Web Search] {snippets}\nSource: tavily"
 
-# ---------------------------------------------------------------------------
+"""End of tools.py"""
-# End of module
-# ---------------------------------------------------------------------------