69 lines
1.9 KiB
Python
69 lines
1.9 KiB
Python
"""Two tools used by the RAG agent.
|
||
|
||
- :func:`search_local_kb` – performs a semantic search in the local ChromaDB vector store.
|
||
- :func:`web_search` – performs a web search via Tavily.
|
||
|
||
Both functions are decorated with :func:`langchain.tools.tool` so that they can be used by LangChain
|
||
agents.
|
||
"""
|
||
|
||
from typing import List
|
||
|
||
from langchain.tools import tool
|
||
from langchain_ollama import ChatOllama
|
||
from tavily import TavilyClient
|
||
|
||
# The LLM used for generating responses. We keep a single instance.
|
||
_llm = ChatOllama(model="llama3")
|
||
|
||
# Tavily client – the API key is read from the environment by the tavily package.
|
||
_tavily_client = TavilyClient()
|
||
|
||
|
||
@tool("search_local_kb")
|
||
def search_local_kb(query: str, top_k: int = 3, vectorstore=None) -> str:
|
||
"""Semantic search in the local knowledge base.
|
||
|
||
Parameters
|
||
----------
|
||
query: str
|
||
The user query.
|
||
top_k: int, optional
|
||
Number of documents to return.
|
||
vectorstore: Chroma, optional
|
||
The vector store instance. It is passed by the agent.
|
||
|
||
Returns
|
||
-------
|
||
str
|
||
Concatenated text of the retrieved documents.
|
||
"""
|
||
if vectorstore is None:
|
||
raise ValueError("vectorstore must be provided")
|
||
retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
|
||
docs = retriever.get_relevant_documents(query)
|
||
return "\n\n".join(doc.page_content for doc in docs)
|
||
|
||
|
||
@tool("web_search")
|
||
def web_search(query: str, max_results: int = 3) -> str:
|
||
"""Perform a web search via Tavily.
|
||
|
||
Parameters
|
||
----------
|
||
query: str
|
||
The search query.
|
||
max_results: int, optional
|
||
Number of search results to return.
|
||
|
||
Returns
|
||
-------
|
||
str
|
||
Concatenated snippets from the search results.
|
||
"""
|
||
results = _tavily_client.search(query, max_results=max_results)
|
||
snippets = [f"{res.title}\n{res.url}\n{res.content}" for res in results]
|
||
return "\n\n".join(snippets)
|
||
|
||
# End of rag_tools.py
|