Add tools.py
This commit is contained in:
@@ -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
|
||||
# ---------------------------------------------------------------------------
|
||||
Reference in New Issue
Block a user