diff --git a/tools.py b/tools.py new file mode 100644 index 0000000..7289501 --- /dev/null +++ b/tools.py @@ -0,0 +1,42 @@ +import os +from typing import Any + +from langchain.tools import tool +from langchain_tavily import TavilySearch + +from vectorstore import create_vectorstore + + +def _format_local_docs(results: list[Any]) -> str: + if not results: + return "Source: qdrant\nNo relevant local documents found." + + parts = ["Source: qdrant"] + for index, doc in enumerate(results, start=1): + source = doc.metadata.get("source", "unknown") + parts.append(f"{index}. ({source}) {doc.page_content}") + return "\n".join(parts) + + +@tool +def search_local_kb(query: str, top_k: int = 3) -> str: + """Search the local ChromaDB knowledge base for internal course notes and local documents.""" + vectorstore = create_vectorstore() + retriever = vectorstore.as_retriever(search_kwargs={"k": top_k}) + results = retriever.invoke(query) + return _format_local_docs(results) + + +@tool +def web_search(query: str) -> str: + """Search the web with Tavily for fresh facts, current news, or anything not covered by local documents.""" + api_key = os.getenv("TAVILY_API_KEY") + if not api_key: + return "Source: tavily\nTAVILY_API_KEY is not set, so web search is unavailable." + + search_tool = TavilySearch( + max_results=5, + topic="general", + ) + result = search_tool.invoke({"query": query}) + return f"Source: tavily\n{result}"