"""RAG agent implementation.

This module exposes two factory functions:

* ``create_agent`` – returns a LangChain agent that can use the two tools
  defined in :mod:`tools`.
* ``create_agent_executor`` – returns an executor that can be used directly
  from the command line.

The agent uses a simple system prompt that instructs it to use the knowledge
base for every query.  The tools are automatically added to the agent.
"""

from typing import Any, Dict

from langchain.agents import AgentExecutor, create_openai_tools_agent
from langchain.chat_models import ChatOpenAI
from langchain.tools import BaseTool

# Import the tools – they expose ``search_knowledge_base`` and
# ``add_to_knowledge_base`` as LangChain tools.
from tools import search_knowledge_base, add_to_knowledge_base

# Create the OpenAI chat model – for local usage we can use Ollama via
# ``ChatOpenAI`` with a custom endpoint.  For the purposes of this
# implementation we assume the user has an OpenAI-compatible endpoint.
# If Ollama is used, replace the model name with ``llama3``.
chat_model = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0)

# List of tools the agent can use
TOOLS: list[BaseTool] = [search_knowledge_base, add_to_knowledge_base]

SYSTEM_PROMPT = (
    "You are an assistant that has access to a knowledge base.  Use the "
    "provided tools to search and add information.  If you need to "
    "retrieve information, call the search_knowledge_base tool.  If you "
    "need to store new data, call add_to_knowledge_base.  Do not "
    "make up facts."
)


def create_agent() -> AgentExecutor:
    """Create a LangChain agent that can perform RAG.

    Returns
    -------
    AgentExecutor
        The configured agent.
    """
    agent = create_openai_tools_agent(
        llm=chat_model,
        tools=TOOLS,
        system_message=SYSTEM_PROMPT,
    )
    executor = AgentExecutor(agent=agent, tools=TOOLS, verbose=True)
    return executor


def create_agent_executor() -> AgentExecutor:
    """Convenience wrapper that returns the same executor.

    The function name is kept for backward compatibility with older
    examples that expected ``create_agent_executor``.
    """
    return create_agent()

# If this file is executed directly, run a simple interactive loop.
if __name__ == "__main__":
    executor = create_agent()
    print("RAG agent ready. Type /quit to exit.")
    while True:
        user_input = input("User: ")
        if user_input.strip().lower() == "/quit":
            break
        response = executor.invoke({"input": user_input})
        print("Agent:", response["output"])