83 lines
2.7 KiB
Python
83 lines
2.7 KiB
Python
"""Tool implementations for the RAG agent.
|
||
|
||
This module defines two tools:
|
||
|
||
* ``search_local_kb`` – semantic search in the local Chroma vector store.
|
||
* ``web_search`` – web search using Tavily.
|
||
|
||
Both tools are decorated with ``@tool`` so that LangChain can call them automatically.
|
||
"""
|
||
|
||
import os
|
||
from typing import List
|
||
|
||
from langchain_core.documents import Document
|
||
from langchain_tavily import TavilySearchResults
|
||
from langchain.tools import tool
|
||
|
||
from vectorstore import get_vectorstore
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Local KB search tool
|
||
# ---------------------------------------------------------------------------
|
||
@tool("search_local_kb")
|
||
# The docstring becomes the tool description used by the model.
|
||
# The function signature must be type annotated.
|
||
# ``top_k`` is optional with default 3.
|
||
# The function returns a string containing the retrieved snippets and a
|
||
# source tag so that the agent can report where the information came from.
|
||
|
||
def search_local_kb(query: str, top_k: int = 3) -> str:
|
||
"""Search the local knowledge base for *query*.
|
||
|
||
Parameters
|
||
----------
|
||
query: str
|
||
The search query.
|
||
top_k: int, optional
|
||
Number of top results to return.
|
||
|
||
Returns
|
||
-------
|
||
str
|
||
A formatted string with the retrieved snippets and a source tag.
|
||
"""
|
||
store = get_vectorstore()
|
||
retriever = store.as_retriever(search_kwargs={"k": top_k})
|
||
docs: List[Document] = retriever.invoke(query)
|
||
if not docs:
|
||
return f"No local KB results for '{query}'."
|
||
# Build a readable answer.
|
||
snippets = "\n\n".join([f"- {doc.page_content[:200]}…" for doc in docs])
|
||
return f"[Local KB] {snippets}\nSource: chromadb"
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Web search tool
|
||
# ---------------------------------------------------------------------------
|
||
@tool("web_search")
|
||
|
||
def web_search(query: str) -> str:
|
||
"""Search the web using Tavily.
|
||
|
||
Parameters
|
||
----------
|
||
query: str
|
||
The search query.
|
||
|
||
Returns
|
||
-------
|
||
str
|
||
A formatted string with the top results and a source tag.
|
||
"""
|
||
# TavilySearchResults requires the API key to be set in the environment.
|
||
api_key = os.getenv("TAVILY_API_KEY")
|
||
if not api_key:
|
||
return "Tavily API key not set. Please set TAVILY_API_KEY environment variable."
|
||
tavily = TavilySearchResults(api_key=api_key, max_results=3)
|
||
results = tavily.invoke(query)
|
||
if not results:
|
||
return f"No web results for '{query}'."
|
||
snippets = "\n\n".join([f"- {res['title']}: {res['content'][:200]}…" for res in results])
|
||
return f"[Web Search] {snippets}\nSource: tavily"
|
||
|
||
"""End of tools.py""" |