diff --git a/tools.py b/tools.py index 50b5658..f39858e 100644 --- a/tools.py +++ b/tools.py @@ -1,54 +1,83 @@ -""" -Agent tools for local KB search and web search via Tavily. +"""Tool implementations for the RAG agent. + +This module defines two tools: + +* ``search_local_kb`` – semantic search in the local Chroma vector store. +* ``web_search`` – web search using Tavily. + +Both tools are decorated with ``@tool`` so that LangChain can call them automatically. """ +import os from typing import List -from langchain_ollama import ChatOllama -from langchain.tools import tool +from langchain_core.documents import Document from langchain_tavily import TavilySearchResults +from langchain.tools import tool + +from vectorstore import get_vectorstore # --------------------------------------------------------------------------- # Local KB search tool # --------------------------------------------------------------------------- -@tool +@tool("search_local_kb") +# The docstring becomes the tool description used by the model. +# The function signature must be type annotated. +# ``top_k`` is optional with default 3. +# The function returns a string containing the retrieved snippets and a +# source tag so that the agent can report where the information came from. + def search_local_kb(query: str, top_k: int = 3) -> str: - """Perform a semantic search in the local Chroma vector store. + """Search the local knowledge base for *query*. - The function expects a global variable ``vectorstore`` to be defined in the - module that imports this tool. This is a simple design choice that keeps - the tool stateless and easy to use from an agent. + Parameters + ---------- + query: str + The search query. + top_k: int, optional + Number of top results to return. + + Returns + ------- + str + A formatted string with the retrieved snippets and a source tag. """ - # The vectorstore is assumed to be a global variable set by the caller. - global vectorstore - if vectorstore is None: - raise RuntimeError("Vectorstore not initialized. Call create_vectorstore first.") - - retriever = vectorstore.as_retriever(search_kwargs={"k": top_k}) - docs = retriever.get_relevant_documents(query) + store = get_vectorstore() + retriever = store.as_retriever(search_kwargs={"k": top_k}) + docs: List[Document] = retriever.invoke(query) if not docs: - return "No relevant documents found." - - # Concatenate the content of the retrieved documents. - return "\n---\n".join(doc.page_content for doc in docs) + return f"No local KB results for '{query}'." + # Build a readable answer. + snippets = "\n\n".join([f"- {doc.page_content[:200]}…" for doc in docs]) + return f"[Local KB] {snippets}\nSource: chromadb" # --------------------------------------------------------------------------- -# Web search tool via Tavily +# Web search tool # --------------------------------------------------------------------------- -@tool +@tool("web_search") + def web_search(query: str) -> str: - """Search the web using Tavily and return a concise summary. + """Search the web using Tavily. - The Tavily client requires the environment variable ``TAVILY_API_KEY``. + Parameters + ---------- + query: str + The search query. + + Returns + ------- + str + A formatted string with the top results and a source tag. """ - tavily = TavilySearchResults(max_results=3) - results = tavily.run(query) - # results is a list of dicts with keys: title, url, content + # TavilySearchResults requires the API key to be set in the environment. + api_key = os.getenv("TAVILY_API_KEY") + if not api_key: + return "Tavily API key not set. Please set TAVILY_API_KEY environment variable." + tavily = TavilySearchResults(api_key=api_key, max_results=3) + results = tavily.invoke(query) if not results: - return "No web results found." - snippets = [f"{r['title']}\n{r['content'][:300]}" for r in results] - return "\n---\n".join(snippets) + return f"No web results for '{query}'." + snippets = "\n\n".join([f"- {res['title']}: {res['content'][:200]}…" for res in results]) + return f"[Web Search] {snippets}\nSource: tavily" -# --------------------------------------------------------------------------- -# End of module -# --------------------------------------------------------------------------- \ No newline at end of file +"""End of tools.py""" \ No newline at end of file