diff --git a/tools.py b/tools.py new file mode 100644 index 0000000..98341cc --- /dev/null +++ b/tools.py @@ -0,0 +1,58 @@ +"""Agent tools: local KB search (ChromaDB) and web search (Tavily).""" +import os +from langchain.tools import tool +from vectorstore import similarity_search + +@tool +def search_local_kb(query: str, top_k: int = 5) -> str: + """Search the local ChromaDB knowledge base for relevant information. + + Use this when the question may be answered from locally stored documents. + + Args: + query: natural language search query + top_k: number of results to return (default 5) + + Returns: + numbered list of relevant passages, or a message if nothing found + """ + docs = similarity_search(query, k=top_k) + if not docs: + return "No relevant documents found in local knowledge base." + results = "\n\n".join( + f"{i + 1}. {doc.page_content}" for i, doc in enumerate(docs) + ) + return f"[Source: Local KB]\n{results}" + +@tool +def web_search(query: str) -> str: + """Search the web for current information using Tavily. + + Use this when the question requires up-to-date or general knowledge + not available in the local knowledge base. + + Args: + query: search query string + + Returns: + web search results with titles, URLs and excerpts + """ + try: + from tavily import TavilyClient + api_key = os.getenv("TAVILY_API_KEY", "") + if not api_key: + return "[Source: Web] Tavily API key not set. Add TAVILY_API_KEY to .env" + client = TavilyClient(api_key=api_key) + response = client.search(query, max_results=5) + items = response.get("results", []) + if not items: + return "[Source: Web] No results found." + lines = [] + for i, r in enumerate(items, 1): + title = r.get("title", "No title") + url = r.get("url", "") + snippet = r.get("content", "")[:300] + lines.append(f"{i}. {title}\n URL: {url}\n {snippet}") + return "[Source: Web]\n" + "\n\n".join(lines) + except Exception as e: + return f"[Source: Web] Search error: {e}"