Update rag_tools.py

This commit is contained in:
2026-06-03 10:25:26 +00:00
parent afdf16482a
commit af8c03f3ae
+59 -34
View File
@@ -1,63 +1,88 @@
"""Tools used by the RAG agent. """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. 1. `search_local_kb` semantic search in the Chroma vector store.
* ``web_search`` realtime web search via Tavily. 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 typing import List
from langchain_core.tools import tool
from langchain_ollama import ChatOllama from langchain_ollama import ChatOllama
from langchain_tavily import TavilySearchResults from langchain_tavily import TavilySearchResults
from langchain_chroma import Chroma from langchain_chroma import Chroma
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.docstore.document import Document
from langchain.tools import tool # Global LLM instance for tool responses (can be reused)
# The LLM used for generating answers. Using the same model as the embeddings
# keeps the pipeline consistent.
_llm = ChatOllama(model="llama3") _llm = ChatOllama(model="llama3")
# Tavily client the API key is read from the environment variable # ---------------------------------------------------------------------------
# ``TAVILY_API_KEY`` by the TavilySearchResults class. # Local KB search tool
_tavily = TavilySearchResults() # ---------------------------------------------------------------------------
@tool("search_local_kb") @tool("search_local_kb")
def search_local_kb(query: str, top_k: int = 3, vectorstore: Chroma = None) -> str: 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: Parameters
query: User question. ----------
top_k: Number of results to return. query: str
vectorstore: The Chroma instance to query. The user's query.
top_k: int, optional
Number of top results to return.
vectorstore: Chroma
The Chroma vector store instance.
Returns: Returns
A string containing the concatenated top results. -------
str
A formatted string containing the retrieved chunks.
""" """
if vectorstore is None: 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}) retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
docs: List[Document] = retriever.invoke(query) docs = retriever.invoke(query)
return "\n\n".join(doc.page_content for doc in docs) # 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") @tool("web_search")
def web_search(query: str, top_k: int = 3) -> str: def web_search(query: str, top_k: int = 3) -> str:
"""Search the web using Tavily. """Perform a web search using Tavily.
Args: Parameters
query: User question. ----------
top_k: Number of results to return. query: str
The user's query.
top_k: int, optional
Number of results to return.
Returns: Returns
Concatenated snippets from the search results. -------
str
A formatted string containing the search results.
""" """
results = _tavily.run(query, max_results=top_k) tavily = TavilySearchResults(tavily_api_key=None, max_results=top_k)
# TavilySearchResults returns a list of dicts with keys like 'title', results = tavily.invoke(query)
# 'content', 'url'. We return the content for simplicity. if not results:
return "\n\n".join(r.get("content", "") for r in 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 *** # ---------------------------------------------------------------------------
# Exported tool names for agent
# ---------------------------------------------------------------------------
TOOLS = [search_local_kb, web_search]
""