from langchain_ollama import ChatOllama
from langchain.agents import create_react_agent, AgentExecutor
from langchain_core.prompts import PromptTemplate
from tools import search_knowledge_base, add_to_knowledge_base

LLM_MODEL = "llama3"

SYSTEM_PROMPT = """You are a helpful AI assistant with access to a knowledge base.
Always use the knowledge base tools to search for relevant information before answering questions.
When you receive new information that should be remembered, add it to the knowledge base.

You have access to the following tools:

{tools}

Use the following format:

Question: the input question you must answer
Thought: you should always think about what to do
Action: the action to take, should be one of [{tool_names}]
Action Input: the input to the action
Observation: the result of the action
... (this Thought/Action/Action Input/Observation can repeat N times)
Thought: I now know the final answer
Final Answer: the final answer to the original input question

Begin!

Question: {input}
Thought:{agent_scratchpad}"""


def create_rag_agent() -> AgentExecutor:
    llm = ChatOllama(model=LLM_MODEL, temperature=0)
    tools = [search_knowledge_base, add_to_knowledge_base]
    prompt = PromptTemplate.from_template(SYSTEM_PROMPT)
    agent = create_react_agent(llm=llm, tools=tools, prompt=prompt)
    agent_executor = AgentExecutor(
        agent=agent,
        tools=tools,
        verbose=True,
        handle_parsing_errors=True,
        max_iterations=10,
    )
    return agent_executor


def run_agent(query: str) -> str:
    agent = create_rag_agent()
    result = agent.invoke({"input": query})
    return result.get("output", "")