"""
Tools for the RAG agent: local semantic search and web search via Tavily.
"""

from typing import List

from langchain.tools import tool
from langchain_ollama import ChatOllama
from langchain_chroma import Chroma
from langchain_tavily import TavilySearchResults

# Local semantic search tool
@tool("search_local_kb")
def search_local_kb(query: str, top_k: int = 3) -> str:
    """Search the local ChromaDB knowledge base.

    Parameters
    ----------
    query: str
        The user's query.
    top_k: int, optional
        Number of top results to return.

    Returns
    -------
    str
        Concatenated content of the top results.
    """
    # Load the vector store (persisted)
    vectorstore = Chroma(persist_directory="./chroma_db", embedding_function=ChatOllama(model="nomic-embed-text"))
    retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
    docs = retriever.invoke(query)
    # docs is a list of Document objects
    return "\n\n---\n\n".join([doc.page_content for doc in docs])

# Web search tool via Tavily
@tool("web_search")
def web_search(query: str) -> str:
    """Perform a web search using Tavily.

    Parameters
    ----------
    query: str
        The user's query.

    Returns
    -------
    str
        Summarized search results.
    """
    tavily = TavilySearchResults(api_key="${TAVILY_API_KEY}")
    results = tavily.run(query)
    # results is a list of dicts with keys: title, url, content
    return "\n\n---\n\n".join([f"{r['title']}\n{r['url']}\n{r.get('content', '')}" for r in results])