55 lines
1.5 KiB
Python
55 lines
1.5 KiB
Python
"""
|
|
Tools for the RAG agent: local semantic search and web search via Tavily.
|
|
"""
|
|
|
|
from typing import List
|
|
|
|
from langchain.tools import tool
|
|
from langchain_ollama import ChatOllama
|
|
from langchain_chroma import Chroma
|
|
from langchain_tavily import TavilySearchResults
|
|
|
|
# Local semantic search tool
|
|
@tool("search_local_kb")
|
|
def search_local_kb(query: str, top_k: int = 3) -> str:
|
|
"""Search the local ChromaDB knowledge base.
|
|
|
|
Parameters
|
|
----------
|
|
query: str
|
|
The user's query.
|
|
top_k: int, optional
|
|
Number of top results to return.
|
|
|
|
Returns
|
|
-------
|
|
str
|
|
Concatenated content of the top results.
|
|
"""
|
|
# Load the vector store (persisted)
|
|
vectorstore = Chroma(persist_directory="./chroma_db", embedding_function=ChatOllama(model="nomic-embed-text"))
|
|
retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
|
|
docs = retriever.invoke(query)
|
|
# docs is a list of Document objects
|
|
return "\n\n---\n\n".join([doc.page_content for doc in docs])
|
|
|
|
# Web search tool via Tavily
|
|
@tool("web_search")
|
|
def web_search(query: str) -> str:
|
|
"""Perform a web search using Tavily.
|
|
|
|
Parameters
|
|
----------
|
|
query: str
|
|
The user's query.
|
|
|
|
Returns
|
|
-------
|
|
str
|
|
Summarized search results.
|
|
"""
|
|
tavily = TavilySearchResults(api_key="${TAVILY_API_KEY}")
|
|
results = tavily.run(query)
|
|
# results is a list of dicts with keys: title, url, content
|
|
return "\n\n---\n\n".join([f"{r['title']}\n{r['url']}\n{r.get('content', '')}" for r in results])
|