"""
Agent creation for the RAG system.

Provides a function ``create_agent`` that returns an ``AgentExecutor`` capable of
choosing between the local KB search and the Tavily web search.
"""

from typing import List

from langchain_ollama import ChatOllama
from langchain.agents import AgentExecutor, create_openai_functions_agent
from langchain.tools import Tool

# Import the tools defined in tools.py
from tools import search_local_kb, web_search

# ---------------------------------------------------------------------------
# Agent creation
# ---------------------------------------------------------------------------

def create_agent(vectorstore_instance) -> AgentExecutor:
    """Create an agent that can decide between local KB and web search.

    Parameters
    ----------
    vectorstore_instance
        Instance of the Chroma vector store to be used by the local search tool.

    Returns
    -------
    AgentExecutor
        Configured agent ready for use.
    """
    # Make the vectorstore available to the tool via the module global
    import tools
    tools.vectorstore = vectorstore_instance

    # Define the tools
    tools_list: List[Tool] = [
        Tool(
            name="search_local_kb",
            func=search_local_kb,
            description="Search the local knowledge base (ChromaDB). Use when the answer is likely contained in the local documents.",
        ),
        Tool(
            name="web_search",
            func=web_search,
            description="Search the web via Tavily. Use when the answer requires up‑to‑date information.",
        ),
    ]

    # LLM for the agent
    llm = ChatOllama(model="llama3", temperature=0)

    # System prompt guiding the agent
    system_prompt = (
        "You are an assistant that answers user questions. "
        "If the answer can be found in the local knowledge base, use the tool "
        "`search_local_kb`. If the question asks for recent or current information, "
        "use the tool `web_search`. After obtaining the information, provide a "
        "concise answer and state the source (`chromadb` or `tavily`)."
    )

    # Create the agent using the function calling approach
    agent = create_openai_functions_agent(llm=llm, tools=tools_list, system_message=system_prompt)

    # Wrap in an executor for easy use
    return AgentExecutor(agent=agent, tools=tools_list, verbose=True)

# ---------------------------------------------------------------------------
# End of module
# ---------------------------------------------------------------------------