diff --git a/tools.py b/tools.py new file mode 100644 index 0000000..1cf7cb6 --- /dev/null +++ b/tools.py @@ -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) \ No newline at end of file