"""Tool implementations for the RAG agent.

This module defines two tools:

* ``search_local_kb`` – semantic search in the local Chroma vector store.
* ``web_search`` – web search using Tavily.

Both tools are decorated with ``@tool`` so that LangChain can call them automatically.
"""

import os
from typing import List

from langchain_core.documents import Document
from langchain_tavily import TavilySearchResults
from langchain.tools import tool

from vectorstore import get_vectorstore

# ---------------------------------------------------------------------------
# Local KB search tool
# ---------------------------------------------------------------------------
@tool("search_local_kb")
# The docstring becomes the tool description used by the model.
# The function signature must be type annotated.
# ``top_k`` is optional with default 3.
# The function returns a string containing the retrieved snippets and a
# source tag so that the agent can report where the information came from.

def search_local_kb(query: str, top_k: int = 3) -> str:
    """Search the local knowledge base for *query*.

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

    Returns
    -------
    str
        A formatted string with the retrieved snippets and a source tag.
    """
    store = get_vectorstore()
    retriever = store.as_retriever(search_kwargs={"k": top_k})
    docs: List[Document] = retriever.invoke(query)
    if not docs:
        return f"No local KB results for '{query}'."
    # Build a readable answer.
    snippets = "\n\n".join([f"- {doc.page_content[:200]}…" for doc in docs])
    return f"[Local KB] {snippets}\nSource: chromadb"

# ---------------------------------------------------------------------------
# Web search tool
# ---------------------------------------------------------------------------
@tool("web_search")

def web_search(query: str) -> str:
    """Search the web using Tavily.

    Parameters
    ----------
    query: str
        The search query.

    Returns
    -------
    str
        A formatted string with the top results and a source tag.
    """
    # TavilySearchResults requires the API key to be set in the environment.
    api_key = os.getenv("TAVILY_API_KEY")
    if not api_key:
        return "Tavily API key not set. Please set TAVILY_API_KEY environment variable."
    tavily = TavilySearchResults(api_key=api_key, max_results=3)
    results = tavily.invoke(query)
    if not results:
        return f"No web results for '{query}'."
    snippets = "\n\n".join([f"- {res['title']}: {res['content'][:200]}…" for res in results])
    return f"[Web Search] {snippets}\nSource: tavily"

"""End of tools.py"""