88 lines
2.7 KiB
Python
88 lines
2.7 KiB
Python
"""Tools used by the RAG agent.
|
||
|
||
This module defines two LangChain tools:
|
||
|
||
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
|
||
|
||
# Global LLM instance for tool responses (can be reused)
|
||
_llm = ChatOllama(model="llama3")
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Local KB search tool
|
||
# ---------------------------------------------------------------------------
|
||
|
||
@tool("search_local_kb")
|
||
|
||
def search_local_kb(query: str, top_k: int = 3, vectorstore: Chroma = None) -> str:
|
||
"""Search the local Chroma vector store for relevant chunks.
|
||
|
||
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
|
||
-------
|
||
str
|
||
A formatted string containing the retrieved chunks.
|
||
"""
|
||
if vectorstore is None:
|
||
raise ValueError("Vectorstore must be provided to search_local_kb tool.")
|
||
retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
|
||
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:
|
||
"""Perform a web search using Tavily.
|
||
|
||
Parameters
|
||
----------
|
||
query: str
|
||
The user's query.
|
||
top_k: int, optional
|
||
Number of results to return.
|
||
|
||
Returns
|
||
-------
|
||
str
|
||
A formatted string containing the search 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)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Exported tool names for agent
|
||
# ---------------------------------------------------------------------------
|
||
|
||
TOOLS = [search_local_kb, web_search]
|
||
"" |