Update rag_tools.py
This commit is contained in:
+59
-34
@@ -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`` – real‑time 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]
|
||||||
|
""
|
||||||
Reference in New Issue
Block a user