diff --git a/tools.py b/tools.py new file mode 100644 index 0000000..b4ff16c --- /dev/null +++ b/tools.py @@ -0,0 +1,41 @@ +""" +Agent tools for local KB search and web search via Tavily. +""" + +from typing import List, Dict + +from langchain.agents import tool +from langchain_tavily import TavilySearch +from langchain.schema.document import Document + +# Local KB search tool +@tool("search_local_kb") +def search_local_kb(query: str, top_k: int = 3) -> List[Dict]: + """Semantic search in the local ChromaDB collection. + + Returns a list of dicts with keys ``text`` and ``source``. + """ + from vectorstore import create_vectorstore + + # Assume the collection is already created and persisted + store = create_vectorstore() + results = store.query(query, top_k=top_k) + return [ + {"text": doc.page_content, "source": doc.metadata.get("source", "unknown")} + for doc in results + ] + +# Web search tool using Tavily. +@tool("web_search") +def web_search(query: str) -> List[Dict]: + """Search the web with Tavily and return a list of result snippets. + + Requires environment variable TAVILY_API_KEY. + """ + tav = TavilySearch() + results = tav.search(query, max_results=5) + return [ + {"text": r["title"] + ": " + r["snippet"] if isinstance(r, dict) else str(r)} + for r in + (results if isinstance(results, list) else []) + ]