diff --git a/rag_tools.py b/rag_tools.py index 01919c6..b26d3e2 100644 --- a/rag_tools.py +++ b/rag_tools.py @@ -1,54 +1,36 @@ -""" -Tools for the RAG agent: local semantic search and web search via Tavily. +"""Tool definitions for the RAG agent. + +Provides two tools: +- search_local_kb: semantic search over the local ChromaDB vector store. +- web_search: web search via Tavily. """ -from typing import List +from typing import List, Dict from langchain.tools import tool from langchain_ollama import ChatOllama -from langchain_chroma import Chroma from langchain_tavily import TavilySearchResults -# Local semantic search tool -@tool("search_local_kb") -def search_local_kb(query: str, top_k: int = 3) -> str: - """Search the local ChromaDB knowledge base. +# Local search tool will be created dynamically in agent.py because it needs the vectorstore. - Parameters - ---------- - query: str - The user's query. - top_k: int, optional - Number of top results to return. - - Returns - ------- - str - Concatenated content of the top results. - """ - # Load the vector store (persisted) - vectorstore = Chroma(persist_directory="./chroma_db", embedding_function=ChatOllama(model="nomic-embed-text")) - retriever = vectorstore.as_retriever(search_kwargs={"k": top_k}) - docs = retriever.invoke(query) - # docs is a list of Document objects - return "\n\n---\n\n".join([doc.page_content for doc in docs]) - -# Web search tool via Tavily -@tool("web_search") +@tool def web_search(query: str) -> str: - """Perform a web search using Tavily. + """Search the web using Tavily and return a short summary. Parameters ---------- query: str - The user's query. + The search query. Returns ------- str - Summarized search results. + A concise answer with a source tag. """ - tavily = TavilySearchResults(api_key="${TAVILY_API_KEY}") + tavily = TavilySearchResults(max_results=3, api_key=None) # API key is taken from env results = tavily.run(query) - # results is a list of dicts with keys: title, url, content - return "\n\n---\n\n".join([f"{r['title']}\n{r['url']}\n{r.get('content', '')}" for r in results]) + # Build a simple summary from the results + summary = "\n".join([f"{idx+1}. {r['title']}: {r['content'][:200]}" for idx, r in enumerate(results)]) + return f"[Web Search]\n{summary}\nSource: tavily" + +# The local search tool will be defined in agent.py where the vectorstore is available.