diff --git a/rag_tools.py b/rag_tools.py index db795e3..4d13824 100644 --- a/rag_tools.py +++ b/rag_tools.py @@ -1,68 +1,91 @@ -"""Two tools used by the RAG agent. +"""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. +Two tools are provided: -Both functions are decorated with :func:`langchain.tools.tool` so that they can be used by LangChain -agents. +* :func:`search_local_kb` – semantic search in the local ChromaDB vector store. +* :func:`web_search` – real‑time web search via Tavily. + +Both functions are decorated with :func:`langchain.tools.tool` so that LangChain +can expose them to the agent. """ +from __future__ import annotations + +import os from typing import List -from langchain.tools import tool +from langchain.chains import RetrievalQA +from langchain.chroma import Chroma from langchain_ollama import ChatOllama -from tavily import TavilyClient +from langchain_tavily import TavilySearchResults +from langchain.tools import tool -# 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() +# Load environment variables (TAVILY_API_KEY) +from dotenv import load_dotenv +load_dotenv() +# --------------------------------------------------------------------------- +# Local KB search +# --------------------------------------------------------------------------- @tool("search_local_kb") -def search_local_kb(query: str, top_k: int = 3, vectorstore=None) -> str: - """Semantic search in the local knowledge base. +def search_local_kb(query: str, top_k: int = 3, vectorstore: Chroma | None = None) -> str: + """Return the best answer from the local ChromaDB store. Parameters ---------- query: str - The user query. - top_k: int, optional - Number of documents to return. - vectorstore: Chroma, optional - The vector store instance. It is passed by the agent. + The user question. + top_k: int + Number of documents to retrieve. + vectorstore: Chroma | None + If ``None`` the function will attempt to load the default store. Returns ------- str - Concatenated text of the retrieved documents. + The answer prefixed with the source identifier. """ if vectorstore is None: - raise ValueError("vectorstore must be provided") + # Load the default store + from vectorstore import create_vectorstore + vectorstore = create_vectorstore() + retriever = vectorstore.as_retriever(search_kwargs={"k": top_k}) - docs = retriever.get_relevant_documents(query) - return "\n\n".join(doc.page_content for doc in docs) - + qa_chain = RetrievalQA.from_chain_type( + llm=ChatOllama(model="llama3", temperature=0), + chain_type="stuff", + retriever=retriever, + ) + result = qa_chain.run(query) + return f"[chromadb] {result}" +# --------------------------------------------------------------------------- +# Web search via Tavily +# --------------------------------------------------------------------------- @tool("web_search") -def web_search(query: str, max_results: int = 3) -> str: - """Perform a web search via Tavily. +def web_search(query: str, top_k: int = 3) -> str: + """Return an answer based on a Tavily web search. Parameters ---------- query: str - The search query. - max_results: int, optional - Number of search results to return. + The user question. + top_k: int + Number of search results to retrieve. Returns ------- str - Concatenated snippets from the search results. + The answer prefixed with the source identifier. """ - 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 + api_key = os.getenv("TAVILY_API_KEY") + if not api_key: + raise RuntimeError("TAVILY_API_KEY environment variable is not set") + tavily = TavilySearchResults(api_key=api_key, top_k=top_k) + results = tavily.run(query) + snippets = "\n\n".join(item.get("content", "") for item in results) + llm = ChatOllama(model="llama3", temperature=0) + prompt = f"Answer the question based on the following web snippets:\n{snippets}\n\nQuestion: {query}\nAnswer:" # noqa: E501 + answer = llm.invoke(prompt).content + return f"[tavily] {answer}"