diff --git a/tools.py b/tools.py index cd9bab5..f7a5081 100644 --- a/tools.py +++ b/tools.py @@ -1,40 +1,42 @@ -""" -Agent tools for local KB search and web search via Tavily. -""" +import os +from typing import Any -from typing import List, Dict - -from langchain.agents import tool +from langchain.tools import tool from langchain_tavily import TavilySearch -from langchain.schema.document import Document -# Local KB search tool -@tool("search_local_kb") -def search_local_kb(query: str, top_k: int = 3) -> List[Dict]: - """Semantic search in the local ChromaDB collection. +from vectorstore import create_vectorstore - Returns a list of dicts with keys ``text`` and ``source``. - """ - from vectorstore import create_vectorstore - # Assume the collection is already created and persisted - store = create_vectorstore() - results = store.query(query_texts=[query], n_results=top_k) - return [ - {"text": doc.page_content, "source": doc.metadata.get("source", "unknown")} - for doc in results - ] +def _format_local_docs(results: list[Any]) -> str: + if not results: + return "Source: chromadb\nNo relevant local documents found." -# Web search tool using Tavily. -@tool("web_search") -def web_search(query: str) -> List[Dict]: - """Search the web with Tavily and return a list of result snippets. + parts = ["Source: chromadb"] + 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) - Requires environment variable TAVILY_API_KEY. - """ - tav = TavilySearch() - results = tav.search(query, max_results=5) - return [ - {"text": r["title"] + ": " + r["snippet"] if isinstance(r, dict) else str(r)} - for r in (results if isinstance(results, list) else []) - ] + +@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}"