add tools.py
This commit is contained in:
@@ -0,0 +1,80 @@
|
||||
import os
|
||||
from langchain_community.tools.tavily import TavilySearchResults
|
||||
from langchain_core.tools import tool
|
||||
from vectorstore import create_vectorstore
|
||||
|
||||
|
||||
# Global vectorstore instance
|
||||
_vectorstore = None
|
||||
|
||||
|
||||
def get_vectorstore():
|
||||
"""Get or create the global vectorstore instance."""
|
||||
global _vectorstore
|
||||
if _vectorstore is None:
|
||||
_vectorstore = create_vectorstore()
|
||||
return _vectorstore
|
||||
|
||||
|
||||
@tool
|
||||
def search_local_kb(query: str, top_k: int = 5) -> str:
|
||||
"""Search for relevant information in the local knowledge base (ChromaDB).
|
||||
|
||||
Use this tool for questions about local documents, notes, or stored knowledge.
|
||||
|
||||
Args:
|
||||
query: The search query
|
||||
top_k: Number of top results to return (default: 5)
|
||||
|
||||
Returns:
|
||||
Search results from the local knowledge base
|
||||
"""
|
||||
vectorstore = get_vectorstore()
|
||||
retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
|
||||
docs = retriever.invoke(query)
|
||||
|
||||
if not docs:
|
||||
return "No relevant information found in local knowledge base."
|
||||
|
||||
results = []
|
||||
for i, doc in enumerate(docs, 1):
|
||||
source = doc.metadata.get("source", "unknown")
|
||||
results.append(f"[Document {i}] {doc.page_content}\n(Source: {source})")
|
||||
|
||||
return "\n\n".join(results)
|
||||
|
||||
|
||||
@tool
|
||||
def web_search(query: str) -> str:
|
||||
"""Search the web for current information using Tavily.
|
||||
|
||||
Use this tool for questions about current events, recent news, or facts
|
||||
that may not be in the local knowledge base.
|
||||
|
||||
Args:
|
||||
query: The search query
|
||||
|
||||
Returns:
|
||||
Web search results
|
||||
"""
|
||||
tavily_api_key = os.getenv("TAVILY_API_KEY")
|
||||
if not tavily_api_key:
|
||||
return "Error: TAVILY_API_KEY not set in environment"
|
||||
|
||||
search = TavilySearchResults(
|
||||
max_results=5,
|
||||
api_key=tavily_api_key
|
||||
)
|
||||
results = search.invoke(query)
|
||||
|
||||
if not results:
|
||||
return "No web search results found."
|
||||
|
||||
formatted_results = []
|
||||
for i, result in enumerate(results, 1):
|
||||
title = result.get("title", "No title")
|
||||
content = result.get("content", "No content")
|
||||
url = result.get("url", "No URL")
|
||||
formatted_results.append(f"[Result {i}] {title}\n{content}\nURL: {url}")
|
||||
|
||||
return "\n\n".join(formatted_results)
|
||||
Reference in New Issue
Block a user