Update rag_tools.py

This commit is contained in:
2026-06-02 07:19:00 +00:00
parent 316cdff63f
commit 442321f34d
+17 -35
View File
@@ -1,54 +1,36 @@
"""
Tools for the RAG agent: local semantic search and web search via Tavily.
"""Tool definitions for the RAG agent.
Provides two tools:
- search_local_kb: semantic search over the local ChromaDB vector store.
- web_search: web search via Tavily.
"""
from typing import List
from typing import List, Dict
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.
# Local search tool will be created dynamically in agent.py because it needs the vectorstore.
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")
@tool
def web_search(query: str) -> str:
"""Perform a web search using Tavily.
"""Search the web using Tavily and return a short summary.
Parameters
----------
query: str
The user's query.
The search query.
Returns
-------
str
Summarized search results.
A concise answer with a source tag.
"""
tavily = TavilySearchResults(api_key="${TAVILY_API_KEY}")
tavily = TavilySearchResults(max_results=3, api_key=None) # API key is taken from env
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])
# Build a simple summary from the results
summary = "\n".join([f"{idx+1}. {r['title']}: {r['content'][:200]}" for idx, r in enumerate(results)])
return f"[Web Search]\n{summary}\nSource: tavily"
# The local search tool will be defined in agent.py where the vectorstore is available.