diff --git a/rag_tools.py b/rag_tools.py index 43d457d..01919c6 100644 --- a/rag_tools.py +++ b/rag_tools.py @@ -1,80 +1,54 @@ -"""Tools for the RAG agent. - -This module defines two LangChain tools: - -* ``search_local_kb`` – semantic search in the local ChromaDB vector store. -* ``web_search`` – real‑time web search using Tavily. - -Both tools return a string containing the retrieved information. +""" +Tools for the RAG agent: local semantic search and web search via Tavily. """ from typing import List -from langchain_ollama import ChatOllama -from langchain_tavily import TavilySearchResults from langchain.tools import tool +from langchain_ollama import ChatOllama +from langchain_chroma import Chroma +from langchain_tavily import TavilySearchResults -# The LLM used for summarising or formatting responses -llm = ChatOllama(model="llama3") - -# Tavily client – the API key is read from the environment by the package -# (requires a .env file or the TAVILY_API_KEY environment variable). -search = TavilySearchResults() - -# --------------------------------------------------------------------------- -# Local knowledge base search tool -# --------------------------------------------------------------------------- +# Local semantic search tool @tool("search_local_kb") - def search_local_kb(query: str, top_k: int = 3) -> str: - """Perform a semantic search in the local ChromaDB vector store. + """Search the local ChromaDB knowledge base. Parameters ---------- query: str - The user question. + The user's query. top_k: int, optional - Number of top documents to return. Defaults to 3. + Number of top results to return. Returns ------- str - A formatted string containing the retrieved passages. + Concatenated content of the top results. """ - # The vectorstore is expected to be loaded globally – the agent will - # provide it via the tool context. We simply call the retriever. - retriever = globals().get("vectorstore_retriever") - if retriever is None: - raise RuntimeError("Vector store retriever not configured for the tool.") + # 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]) - docs = retriever.get_relevant_documents(query, k=top_k) - # Concatenate the documents into a single string. - passages = "\n\n".join(doc.page_content for doc in docs) - return passages - -# --------------------------------------------------------------------------- -# Web search tool -# --------------------------------------------------------------------------- +# Web search tool via Tavily @tool("web_search") - def web_search(query: str) -> str: - """Search the web using Tavily and return the top results. + """Perform a web search using Tavily. Parameters ---------- query: str - The user question. + The user's query. Returns ------- str - A formatted string containing the search results. + Summarized search results. """ - results = search.run(query) - # TavilySearchResults returns a list of dicts with keys: title, url, content - formatted = [] - for r in results: - formatted.append(f"Title: {r.get('title', 'N/A')}\nURL: {r.get('url', 'N/A')}\nSnippet: {r.get('content', 'N/A')}\n") - return "\n\n".join(formatted) - -# End of rag_tools.py + tavily = TavilySearchResults(api_key="${TAVILY_API_KEY}") + 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])