81 lines
2.5 KiB
Python
81 lines
2.5 KiB
Python
"""Tools for the RAG agent.
|
||
|
||
This module defines two LangChain tools:
|
||
|
||
* ``search_local_kb`` – semantic search in the local ChromaDB vector store.
|
||
* ``web_search`` – real‑time web search using Tavily.
|
||
|
||
Both tools return a string containing the retrieved information.
|
||
"""
|
||
|
||
from typing import List
|
||
|
||
from langchain_ollama import ChatOllama
|
||
from langchain_tavily import TavilySearchResults
|
||
from langchain.tools import tool
|
||
|
||
# The LLM used for summarising or formatting responses
|
||
llm = ChatOllama(model="llama3")
|
||
|
||
# Tavily client – the API key is read from the environment by the package
|
||
# (requires a .env file or the TAVILY_API_KEY environment variable).
|
||
search = TavilySearchResults()
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Local knowledge base search tool
|
||
# ---------------------------------------------------------------------------
|
||
@tool("search_local_kb")
|
||
|
||
def search_local_kb(query: str, top_k: int = 3) -> str:
|
||
"""Perform a semantic search in the local ChromaDB vector store.
|
||
|
||
Parameters
|
||
----------
|
||
query: str
|
||
The user question.
|
||
top_k: int, optional
|
||
Number of top documents to return. Defaults to 3.
|
||
|
||
Returns
|
||
-------
|
||
str
|
||
A formatted string containing the retrieved passages.
|
||
"""
|
||
# The vectorstore is expected to be loaded globally – the agent will
|
||
# provide it via the tool context. We simply call the retriever.
|
||
retriever = globals().get("vectorstore_retriever")
|
||
if retriever is None:
|
||
raise RuntimeError("Vector store retriever not configured for the tool.")
|
||
|
||
docs = retriever.get_relevant_documents(query, k=top_k)
|
||
# Concatenate the documents into a single string.
|
||
passages = "\n\n".join(doc.page_content for doc in docs)
|
||
return passages
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Web search tool
|
||
# ---------------------------------------------------------------------------
|
||
@tool("web_search")
|
||
|
||
def web_search(query: str) -> str:
|
||
"""Search the web using Tavily and return the top results.
|
||
|
||
Parameters
|
||
----------
|
||
query: str
|
||
The user question.
|
||
|
||
Returns
|
||
-------
|
||
str
|
||
A formatted string containing the search results.
|
||
"""
|
||
results = search.run(query)
|
||
# TavilySearchResults returns a list of dicts with keys: title, url, content
|
||
formatted = []
|
||
for r in results:
|
||
formatted.append(f"Title: {r.get('title', 'N/A')}\nURL: {r.get('url', 'N/A')}\nSnippet: {r.get('content', 'N/A')}\n")
|
||
return "\n\n".join(formatted)
|
||
|
||
# End of rag_tools.py
|