65 lines
2.0 KiB
Python
65 lines
2.0 KiB
Python
"""
|
|
Tools for the RAG agent: local semantic search and web search via Tavily.
|
|
"""
|
|
|
|
from typing import List
|
|
|
|
from langchain_community.tools.tavily import TavilySearchResults
|
|
from langchain.tools import tool
|
|
from langchain_chroma import Chroma
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Local KB search tool
|
|
# ---------------------------------------------------------------------------
|
|
@tool("search_local_kb")
|
|
|
|
def search_local_kb(query: str, top_k: int = 3, vectorstore: Chroma = None) -> str:
|
|
"""Perform a semantic search in the local Chroma vector store.
|
|
|
|
Parameters
|
|
----------
|
|
query: str
|
|
The user query.
|
|
top_k: int, optional
|
|
Number of top results to return.
|
|
vectorstore: Chroma, optional
|
|
The vector store to query. If None, the function will raise an error.
|
|
|
|
Returns
|
|
-------
|
|
str
|
|
Concatenated content of the top results.
|
|
"""
|
|
if vectorstore is None:
|
|
raise ValueError("vectorstore must be provided to search_local_kb")
|
|
retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
|
|
docs = retriever.get_relevant_documents(query)
|
|
return "\n\n---\n\n".join(doc.page_content for doc in docs)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Web search tool using Tavily
|
|
# ---------------------------------------------------------------------------
|
|
@tool("web_search")
|
|
|
|
def web_search(query: str, top_k: int = 3) -> str:
|
|
"""Search the web via Tavily and return a formatted string of results.
|
|
|
|
Parameters
|
|
----------
|
|
query: str
|
|
The search query.
|
|
top_k: int, optional
|
|
Number of top results to return.
|
|
|
|
Returns
|
|
-------
|
|
str
|
|
Formatted search results.
|
|
"""
|
|
tavily = TavilySearchResults(max_results=top_k)
|
|
results = tavily.run(query)
|
|
formatted = []
|
|
for i, r in enumerate(results, 1):
|
|
formatted.append(f"{i}. {r.get('title', 'No title')}\n{r.get('url', '')}\n{r.get('content', '')}")
|
|
return "\n\n---\n\n".join(formatted)
|