diff --git a/tools.py b/tools.py new file mode 100644 index 0000000..b33c855 --- /dev/null +++ b/tools.py @@ -0,0 +1,28 @@ +import os +from langchain.tools import tool +from langchain_tavily import TavilySearchResults +from langchain_chroma import Chroma + +@tool +def search_local_kb(query: str, top_k: int = 3) -> str: + """Search the local knowledge base for relevant information.""" + vectorstore = Chroma( + collection_name="knowledge", + persist_directory="./chroma_db", + ) + docs = vectorstore.similarity_search(query, k=top_k) + if not docs: + return "No results found in local knowledge base." + return "\n".join(f\"{i+1}. {doc.page_content[:200]}...\" for i, doc in enumerate(docs)) + +@tool +def web_search(query: str) -> str: + """Search the web for up-to-date information.""" + tavily = TavilySearchResults( + api_key=os.getenv("TAVILY_API_KEY"), + max_results=3, + ) + results = tavily.run(query) + if not results: + return "No results found on the web." + return "\n".join(f\"{i+1}. {res['title']}\\n{res['url']}\\n{res['content'][:200]}...\" for i, res in enumerate(results)) \ No newline at end of file