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