diff --git a/rag_tools.py b/rag_tools.py index c8c463b..1d678f2 100644 --- a/rag_tools.py +++ b/rag_tools.py @@ -1,63 +1,88 @@ """Tools used by the RAG agent. -Two tools are provided: +This module defines two LangChain tools: -* ``search_local_kb`` – semantic search in the local ChromaDB store. -* ``web_search`` – real‑time web search via Tavily. +1. `search_local_kb` – semantic search in the Chroma vector store. +2. `web_search` – web search using Tavily. + +Both tools are decorated with `@tool` so that they can be used by the agent. """ from typing import List +from langchain_core.tools import tool from langchain_ollama import ChatOllama from langchain_tavily import TavilySearchResults from langchain_chroma import Chroma -from langchain_text_splitters import RecursiveCharacterTextSplitter -from langchain.docstore.document import Document -from langchain.tools import tool - -# The LLM used for generating answers. Using the same model as the embeddings -# keeps the pipeline consistent. +# Global LLM instance for tool responses (can be reused) _llm = ChatOllama(model="llama3") -# Tavily client – the API key is read from the environment variable -# ``TAVILY_API_KEY`` by the TavilySearchResults class. -_tavily = TavilySearchResults() - +# --------------------------------------------------------------------------- +# Local KB search tool +# --------------------------------------------------------------------------- @tool("search_local_kb") + def search_local_kb(query: str, top_k: int = 3, vectorstore: Chroma = None) -> str: - """Semantic search in the local ChromaDB vector store. + """Search the local Chroma vector store for relevant chunks. - Args: - query: User question. - top_k: Number of results to return. - vectorstore: The Chroma instance to query. + Parameters + ---------- + query: str + The user's query. + top_k: int, optional + Number of top results to return. + vectorstore: Chroma + The Chroma vector store instance. - Returns: - A string containing the concatenated top results. + Returns + ------- + str + A formatted string containing the retrieved chunks. """ if vectorstore is None: - raise ValueError("vectorstore must be provided") + raise ValueError("Vectorstore must be provided to search_local_kb tool.") retriever = vectorstore.as_retriever(search_kwargs={"k": top_k}) - docs: List[Document] = retriever.invoke(query) - return "\n\n".join(doc.page_content for doc in docs) + docs = retriever.invoke(query) + # docs is a list of Document objects + if not docs: + return "No relevant information found in the local knowledge base." + # Concatenate the content of the top documents + snippets = [f"{i+1}. {doc.page_content[:200]}..." for i, doc in enumerate(docs)] + return "\n".join(snippets) +# --------------------------------------------------------------------------- +# Web search tool +# --------------------------------------------------------------------------- @tool("web_search") + def web_search(query: str, top_k: int = 3) -> str: - """Search the web using Tavily. + """Perform a web search using Tavily. - Args: - query: User question. - top_k: Number of results to return. + Parameters + ---------- + query: str + The user's query. + top_k: int, optional + Number of results to return. - Returns: - Concatenated snippets from the search results. + Returns + ------- + str + A formatted string containing the search results. """ - results = _tavily.run(query, max_results=top_k) - # TavilySearchResults returns a list of dicts with keys like 'title', - # 'content', 'url'. We return the content for simplicity. - return "\n\n".join(r.get("content", "") for r in results) + tavily = TavilySearchResults(tavily_api_key=None, max_results=top_k) + results = tavily.invoke(query) + if not results: + return "No results found on the web." + snippets = [f"{i+1}. {res['title']} – {res['url']}" for i, res in enumerate(results)] + return "\n".join(snippets) -*** End of File *** \ No newline at end of file +# --------------------------------------------------------------------------- +# Exported tool names for agent +# --------------------------------------------------------------------------- + +TOOLS = [search_local_kb, web_search] +"" \ No newline at end of file