"""
 /*************************************************************************
 * 
 * CONFIDENTIAL
 * __________________
 * 
 *  Copyright (2023-2024) AI Labs, IronOne Technologies, LLC
 *  All Rights Reserved
 *
 *  Author  : Theekshana Samaradiwakara
 *  Description :Python Backend API to chat with private data  
 *  CreatedDate : 17/10/2024
 *  LastModifiedDate : 
 *************************************************************************/
 """
import os
import time
import logging
logger = logging.getLogger(__name__)

from langchain.agents import AgentExecutor, LLMSingleActionAgent

from reggpt.output_parsers.qa_agent import CustomOutputParser
from reggpt.output_parsers.output_parser import qa_agent_output_parser
from reggpt.tools.qa_agent import qa_agent_tools
from reggpt.configs.model import QA_AGENT_MODEL_TYPE
from reggpt.chains.qa_agent import get_qa_agent_chain


output_parser = CustomOutputParser()
qa_agent_chain = get_qa_agent_chain(QA_AGENT_MODEL_TYPE)

qa_agent_tools_names = [tool.name for tool in qa_agent_tools]
qa_agent = LLMSingleActionAgent(
    llm_chain=qa_agent_chain,
    output_parser=output_parser,
    stop=["\nObservation:"],
    allowed_tools=qa_agent_tools_names,
)

qa_agent_executor = AgentExecutor.from_agent_and_tools(
    agent=qa_agent, tools=qa_agent_tools, verbose=True
)

def run_qa_agent(query):
    try:
        logger.info(f"run_qa_agent : Question: {query}")

        # Get the answer from the chain
        start = time.time()
        res = qa_agent_executor.run(query)
        end = time.time()

        # log the result
       
        logger.info(f"Answer (took {round(end - start, 2)} s.) \n: {res}")

        return qa_agent_output_parser(query,res)

    except Exception as e:
        logger.exception(e)
        raise e