diff --git a/tools.py b/tools.py new file mode 100644 index 0000000..0c54f66 --- /dev/null +++ b/tools.py @@ -0,0 +1,30 @@ +import os +from typing import Any + +from langchain.tools import tool +from langchain_ollama import OllamaEmbeddings +from langchain_chroma import Chroma +from langchain_tavily import TavilySearchResults + +from vectorstore import create_vectorstore + +# Global vector store instance +VECTORSTORE = create_vectorstore() + +@tool +def search_local_kb(query: str, top_k: int = 3) -> str: + """Semantic search in the local ChromaDB knowledge base.""" + retriever = VECTORSTORE.as_retriever(search_kwargs={"k": top_k}) + docs = retriever.get_relevant_documents(query) + if not docs: + return "No relevant local knowledge found." + return "\n\n".join(doc.page_content for doc in docs) + +@tool +def web_search(query: str) -> str: + """Web search using Tavily.""" + tavily = TavilySearchResults(api_key=os.getenv("TAVILY_API_KEY")) + results = tavily.run(query) + if not results: + return "No web results found." + return "\n\n".join(f"{r['title']}\n{r['content']}" for r in results)