from typing import Dict
import langchain_community
from langchain.agents import AgentExecutor
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import JsonOutputParser, StrOutputParser
from langchain_community.chat_models import ChatOllama

from agents.functions_agent.base import create_functions_agent
from functions import tools, get_openai_tools
from config import config

from prompts.prompt import (
    rag_agent_prompt, 
    doc_grader_agent_prompt, 
    router_agent_prompt, 
    hallucination_grader_agent_prompt, 
    rephrase_agent_prompt,
    output_agent_prompt,
)

tools_dict = get_openai_tools()

def get_agents(llm: langchain_community.chat_models.ChatOllama) -> Dict[str, AgentExecutor]:
    """Returns all available agents."""
    functions_agent = create_functions_agent(llm=llm, prompt=rag_agent_prompt)
    functions_agent_executor = AgentExecutor(agent=functions_agent, tools=tools, verbose=True, return_intermediate_steps=True)
    
    rephrase_agent = rephrase_agent_prompt | llm | StrOutputParser()
    router_agent = router_agent_prompt | llm | JsonOutputParser()
    hallucination_grader_agent = hallucination_grader_agent_prompt | llm | JsonOutputParser()
    doc_grader_agent = doc_grader_agent_prompt | llm | JsonOutputParser()
    output_agent = output_agent_prompt | llm | StrOutputParser()
    
    return {
        "function_agent": functions_agent_executor,
        "doc_grader_agent": doc_grader_agent,
        "hallucination_grader_agent": hallucination_grader_agent,
        "router_agent": router_agent,
        "rephrase_agent": rephrase_agent,
        "output_agent": output_agent,
    }
    

if __name__ == "__main__":
    from langchain.memory import ChatMessageHistory
    history = ChatMessageHistory()
    
    llm = ChatOllama(model = config.ollama_model, temperature = 0.55)
    function_agent = get_agents(llm)["function_agent"]
    
    while True:
        try:
            inp = input("User:")
            if inp == "/bye":
                break

            response = function_agent.invoke({"input": inp, "chat_history": history, "tools" : tools_dict})
            response['output'] = response['output'].replace("<|im_end|>", "")
            history.add_user_message(inp)
            history.add_ai_message(response['output'])

            print(response['output'])
        except Exception as e:
            print(e)
            continue