diff --git a/rag_tools.py b/rag_tools.py index 658cb10..db795e3 100644 --- a/rag_tools.py +++ b/rag_tools.py @@ -1,64 +1,68 @@ -""" -Tools for the RAG agent: local semantic search and web search via Tavily. +"""Two tools used by the RAG agent. + +- :func:`search_local_kb` – performs a semantic search in the local ChromaDB vector store. +- :func:`web_search` – performs a web search via Tavily. + +Both functions are decorated with :func:`langchain.tools.tool` so that they can be used by LangChain +agents. """ from typing import List -from langchain_community.tools.tavily import TavilySearchResults from langchain.tools import tool -from langchain_chroma import Chroma +from langchain_ollama import ChatOllama +from tavily import TavilyClient + +# The LLM used for generating responses. We keep a single instance. +_llm = ChatOllama(model="llama3") + +# Tavily client – the API key is read from the environment by the tavily package. +_tavily_client = TavilyClient() + -# --------------------------------------------------------------------------- -# Local KB search tool -# --------------------------------------------------------------------------- @tool("search_local_kb") - -def search_local_kb(query: str, top_k: int = 3, vectorstore: Chroma = None) -> str: - """Perform a semantic search in the local Chroma vector store. +def search_local_kb(query: str, top_k: int = 3, vectorstore=None) -> str: + """Semantic search in the local knowledge base. Parameters ---------- query: str The user query. top_k: int, optional - Number of top results to return. + Number of documents to return. vectorstore: Chroma, optional - The vector store to query. If None, the function will raise an error. + The vector store instance. It is passed by the agent. Returns ------- str - Concatenated content of the top results. + Concatenated text of the retrieved documents. """ if vectorstore is None: - raise ValueError("vectorstore must be provided to search_local_kb") + raise ValueError("vectorstore must be provided") retriever = vectorstore.as_retriever(search_kwargs={"k": top_k}) docs = retriever.get_relevant_documents(query) - return "\n\n---\n\n".join(doc.page_content for doc in docs) + return "\n\n".join(doc.page_content for doc in docs) + -# --------------------------------------------------------------------------- -# Web search tool using Tavily -# --------------------------------------------------------------------------- @tool("web_search") - -def web_search(query: str, top_k: int = 3) -> str: - """Search the web via Tavily and return a formatted string of results. +def web_search(query: str, max_results: int = 3) -> str: + """Perform a web search via Tavily. Parameters ---------- query: str The search query. - top_k: int, optional - Number of top results to return. + max_results: int, optional + Number of search results to return. Returns ------- str - Formatted search results. + Concatenated snippets from the search results. """ - tavily = TavilySearchResults(max_results=top_k) - results = tavily.run(query) - formatted = [] - for i, r in enumerate(results, 1): - formatted.append(f"{i}. {r.get('title', 'No title')}\n{r.get('url', '')}\n{r.get('content', '')}") - return "\n\n---\n\n".join(formatted) + results = _tavily_client.search(query, max_results=max_results) + snippets = [f"{res.title}\n{res.url}\n{res.content}" for res in results] + return "\n\n".join(snippets) + +# End of rag_tools.py