58 lines
1.7 KiB
Python
58 lines
1.7 KiB
Python
"""Tools for the RAG agent: local KB search and web search via Tavily."""
|
|
|
|
from typing import List, Dict
|
|
|
|
from langchain_ollama import ChatOllama
|
|
from langchain_chroma import Chroma
|
|
from langchain_tavily import TavilySearchResults
|
|
from langchain.tools import tool
|
|
|
|
# --- Local KB search tool -----------------------------------------------------
|
|
|
|
@tool("search_local_kb")
|
|
def search_local_kb(query: str, top_k: int = 3, vectorstore: Chroma = None) -> List[Dict]:
|
|
"""Perform a semantic search in the local Chroma vector store.
|
|
|
|
Parameters
|
|
----------
|
|
query: str
|
|
The user query.
|
|
top_k: int
|
|
Number of top results to return.
|
|
vectorstore: Chroma
|
|
The vector store to search.
|
|
|
|
Returns
|
|
-------
|
|
List[Dict]
|
|
List of dictionaries containing ``content`` and ``metadata``.
|
|
"""
|
|
if vectorstore is None:
|
|
raise ValueError("vectorstore must be provided")
|
|
retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
|
|
docs = retriever.get_relevant_documents(query)
|
|
return [{"content": doc.page_content, "metadata": doc.metadata} for doc in docs]
|
|
|
|
# --- Web search tool ---------------------------------------------------------
|
|
|
|
@tool("web_search")
|
|
def web_search(query: str, top_k: int = 3) -> List[Dict]:
|
|
"""Search the web using Tavily.
|
|
|
|
Parameters
|
|
----------
|
|
query: str
|
|
The user query.
|
|
top_k: int
|
|
Number of top results to return.
|
|
|
|
Returns
|
|
-------
|
|
List[Dict]
|
|
List of dictionaries containing ``title``, ``url`` and ``content``.
|
|
"""
|
|
tavily = TavilySearchResults(max_results=top_k)
|
|
results = tavily.run(query)
|
|
# Tavily returns a list of dicts with keys: title, url, content
|
|
return results
|