from langchain.agents import load_tools
from langchain.agents import initialize_agent
from langchain.chat_models import ChatOpenAI

from virl.config import cfg
from virl.utils.common_utils import print_prompt, print_answer, parse_answer_to_json
from .gpt_chat import GPTChat
from .azure_gpt import AzureGPTChat


__all__ = {
    'GPT': GPTChat,
    'AzureGPT': AzureGPTChat,
}


def build_chatbot(name):
    return __all__[name](cfg)


class UnifiedChat(object):
    chatbots = None

    def __init__(self):
        UnifiedChat.chatbots = {
            name: build_chatbot(name) for name in cfg.LLM.NAMES
        }

    @classmethod
    def ask(cls, question, **kwargs):
        print_prompt(question)
        chatbot = kwargs.get('chatbot', cfg.LLM.DEFAULT)
        answer = cls.chatbots[chatbot].ask(question, **kwargs)
        print_answer(answer)

        return answer

    @classmethod
    def search(cls, question, json=False):
        llm = ChatOpenAI(model='gpt-3.5-turbo', temperature=0)
        tools = load_tools(["serpapi"], llm=llm)
        agent = initialize_agent(tools, llm, verbose=True)
        answer = agent.run(question)

        if json:
            answer = parse_answer_to_json(answer)

        return answer

langchain.agents.initialize_agent langchain.agents.load_tools langchain.chat_models.ChatOpenAI
[((679, 701), 'virl.utils.common_utils.print_prompt', 'print_prompt', (['question'], {}), '(question)\n', (691, 701), False, 'from virl.utils.common_utils import print_prompt, print_answer, parse_answer_to_json\n'), ((830, 850), 'virl.utils.common_utils.print_answer', 'print_answer', (['answer'], {}), '(answer)\n', (842, 850), False, 'from virl.utils.common_utils import print_prompt, print_answer, parse_answer_to_json\n'), ((949, 997), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'model': '"""gpt-3.5-turbo"""', 'temperature': '(0)'}), "(model='gpt-3.5-turbo', temperature=0)\n", (959, 997), False, 'from langchain.chat_models import ChatOpenAI\n'), ((1014, 1046), 'langchain.agents.load_tools', 'load_tools', (["['serpapi']"], {'llm': 'llm'}), "(['serpapi'], llm=llm)\n", (1024, 1046), False, 'from langchain.agents import load_tools\n'), ((1063, 1105), 'langchain.agents.initialize_agent', 'initialize_agent', (['tools', 'llm'], {'verbose': '(True)'}), '(tools, llm, verbose=True)\n', (1079, 1105), False, 'from langchain.agents import initialize_agent\n'), ((1182, 1210), 'virl.utils.common_utils.parse_answer_to_json', 'parse_answer_to_json', (['answer'], {}), '(answer)\n', (1202, 1210), False, 'from virl.utils.common_utils import print_prompt, print_answer, parse_answer_to_json\n')]