from langchain_ollama import Ollama
from langchain.agents import initialize_agent, AgentType
from langchain.tools import Tool

from .tools import search_knowledge_base, add_to_knowledge_base
from .config import LLM_MODEL


def create_agent():
    """
    Create an RAG-enabled agent that can search and add to a knowledge base.

    Returns
    -------
    AgentExecutor
        The configured agent.
    """
    llm = Ollama(model=LLM_MODEL)

    tools = [
        Tool(
            name="search_knowledge_base",
            func=search_knowledge_base,
            description="Search the knowledge base for relevant documents."
        ),
        Tool(
            name="add_to_knowledge_base",
            func=add_to_knowledge_base,
            description="Add a new document to the knowledge base."
        ),
    ]

    system_prompt = (
        "You are an AI assistant that can search and add information to a knowledge base. "
        "Use the provided tools to answer user queries."
    )

    agent = initialize_agent(
        tools=tools,
        llm=llm,
        agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
        verbose=True,
        system_message=system_prompt,
    )
    return agent