diff --git a/tools.py b/tools.py new file mode 100644 index 0000000..50b5658 --- /dev/null +++ b/tools.py @@ -0,0 +1,54 @@ +""" +Agent tools for local KB search and web search via Tavily. +""" + +from typing import List + +from langchain_ollama import ChatOllama +from langchain.tools import tool +from langchain_tavily import TavilySearchResults + +# --------------------------------------------------------------------------- +# Local KB search tool +# --------------------------------------------------------------------------- +@tool +def search_local_kb(query: str, top_k: int = 3) -> str: + """Perform a semantic search in the local Chroma vector store. + + The function expects a global variable ``vectorstore`` to be defined in the + module that imports this tool. This is a simple design choice that keeps + the tool stateless and easy to use from an agent. + """ + # The vectorstore is assumed to be a global variable set by the caller. + global vectorstore + if vectorstore is None: + 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." + + # Concatenate the content of the retrieved documents. + return "\n---\n".join(doc.page_content for doc in docs) + +# --------------------------------------------------------------------------- +# Web search tool via Tavily +# --------------------------------------------------------------------------- +@tool +def web_search(query: str) -> str: + """Search the web using Tavily and return a concise summary. + + The Tavily client requires the environment variable ``TAVILY_API_KEY``. + """ + tavily = TavilySearchResults(max_results=3) + results = tavily.run(query) + # results is a list of dicts with keys: title, url, content + if not results: + return "No web results found." + snippets = [f"{r['title']}\n{r['content'][:300]}" for r in results] + return "\n---\n".join(snippets) + +# --------------------------------------------------------------------------- +# End of module +# --------------------------------------------------------------------------- \ No newline at end of file