Add tools.py

This commit is contained in:
2026-06-02 07:47:17 +00:00
parent cd0729327d
commit 51cde60507
+54
View File
@@ -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
# ---------------------------------------------------------------------------