Update rag_tools.py
This commit is contained in:
+17
-35
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user