80 lines
2.3 KiB
Python
80 lines
2.3 KiB
Python
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) |