63 lines
1.9 KiB
Python
63 lines
1.9 KiB
Python
"""Tools used by the RAG agent.
|
||
|
||
Two tools are provided:
|
||
|
||
* ``search_local_kb`` – semantic search in the local ChromaDB store.
|
||
* ``web_search`` – real‑time web search via Tavily.
|
||
"""
|
||
|
||
from typing import List
|
||
|
||
from langchain_ollama import ChatOllama
|
||
from langchain_tavily import TavilySearchResults
|
||
from langchain_chroma import Chroma
|
||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||
from langchain.docstore.document import Document
|
||
|
||
from langchain.tools import tool
|
||
|
||
# The LLM used for generating answers. Using the same model as the embeddings
|
||
# keeps the pipeline consistent.
|
||
_llm = ChatOllama(model="llama3")
|
||
|
||
# Tavily client – the API key is read from the environment variable
|
||
# ``TAVILY_API_KEY`` by the TavilySearchResults class.
|
||
_tavily = TavilySearchResults()
|
||
|
||
|
||
@tool("search_local_kb")
|
||
def search_local_kb(query: str, top_k: int = 3, vectorstore: Chroma = None) -> str:
|
||
"""Semantic search in the local ChromaDB vector store.
|
||
|
||
Args:
|
||
query: User question.
|
||
top_k: Number of results to return.
|
||
vectorstore: The Chroma instance to query.
|
||
|
||
Returns:
|
||
A string containing the concatenated top results.
|
||
"""
|
||
if vectorstore is None:
|
||
raise ValueError("vectorstore must be provided")
|
||
retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
|
||
docs: List[Document] = retriever.invoke(query)
|
||
return "\n\n".join(doc.page_content for doc in docs)
|
||
|
||
|
||
@tool("web_search")
|
||
def web_search(query: str, top_k: int = 3) -> str:
|
||
"""Search the web using Tavily.
|
||
|
||
Args:
|
||
query: User question.
|
||
top_k: Number of results to return.
|
||
|
||
Returns:
|
||
Concatenated snippets from the search results.
|
||
"""
|
||
results = _tavily.run(query, max_results=top_k)
|
||
# TavilySearchResults returns a list of dicts with keys like 'title',
|
||
# 'content', 'url'. We return the content for simplicity.
|
||
return "\n\n".join(r.get("content", "") for r in results)
|
||
|
||
*** End of File *** |