diff --git a/rag_tools.py b/rag_tools.py new file mode 100644 index 0000000..e89a058 --- /dev/null +++ b/rag_tools.py @@ -0,0 +1,57 @@ +"""Tools for the RAG agent: local KB search and web search via Tavily.""" + +from typing import List, Dict + +from langchain_ollama import ChatOllama +from langchain_chroma import Chroma +from langchain_tavily import TavilySearchResults +from langchain.tools import tool + +# --- Local KB search tool ----------------------------------------------------- + +@tool("search_local_kb") +def search_local_kb(query: str, top_k: int = 3, vectorstore: Chroma = None) -> List[Dict]: + """Perform a semantic search in the local Chroma vector store. + + Parameters + ---------- + query: str + The user query. + top_k: int + Number of top results to return. + vectorstore: Chroma + The vector store to search. + + Returns + ------- + List[Dict] + List of dictionaries containing ``content`` and ``metadata``. + """ + if vectorstore is None: + raise ValueError("vectorstore must be provided") + retriever = vectorstore.as_retriever(search_kwargs={"k": top_k}) + docs = retriever.get_relevant_documents(query) + return [{"content": doc.page_content, "metadata": doc.metadata} for doc in docs] + +# --- Web search tool --------------------------------------------------------- + +@tool("web_search") +def web_search(query: str, top_k: int = 3) -> List[Dict]: + """Search the web using Tavily. + + Parameters + ---------- + query: str + The user query. + top_k: int + Number of top results to return. + + Returns + ------- + List[Dict] + List of dictionaries containing ``title``, ``url`` and ``content``. + """ + tavily = TavilySearchResults(max_results=top_k) + results = tavily.run(query) + # Tavily returns a list of dicts with keys: title, url, content + return results