"""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