from langchain.agents import Tool
from langchain.memory import ConversationBufferMemory
from langchain.chat_models import ChatOpenAI
from langchain.agents import initialize_agent
from llama_index import GPTSimpleVectorIndex
import os
from langchain.schema import (
    SystemMessage
)
#import openai
os.environ["OPENAI_API_KEY"] = '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'
#openai.organization = "org-LWOdWb7zopsw1g1tJUgDOYUo"
#openai.api_key = "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"


class ChatBot:
    def __init__(self, memory, agent_chain):
        self.memory = memory
        self.agent = agent_chain


def create_chatbot(model_name, seed_memory=None):
    # search = GoogleSearchAPIWrapper()
    # tools = [
    #     Tool(
    #         name="Current Search",
    #         func=search.run,
    #         description="useful for all question that asks about live events",
    #     ),
    #     Tool(
    #         name="Topic Search",
    #         func=search.run,
    #         description="useful for all question that are related to a particular topic, product, concept, or service",
    #     )
    # ]
    index = GPTSimpleVectorIndex.load_from_disk('martin.json')
    query_mode ="svm"

    tools = [
        Tool(
            name="GPT Index",
            func=lambda q: str(index.query(q,vector_store_query_mode=query_mode)),
            description="useful for when you want to answer questions about Martin Seligman , personal issues , family issues , working issues , children issues , positive psychonogy related , such as Broaden-and-Build Theory,PERMA Model,Positive Psychology Interventions,Mindfulness,Meaning and Purpose,Positive Relationships,Resilience,Flow,Optimism,Character Strengths,Happiness and Subjective Well-Being. The input to this tool should be a complete english sentence.",
            return_direct=True
        ),
    ]

    # messages = [
    #     SystemMessage(content="You are Martin Seligman. You use a tone that is warm and kind.")
    # ]

    memory = seed_memory if seed_memory is not None else ConversationBufferMemory(memory_key="chat_history")
    chat = ChatOpenAI(temperature=0, model_name=model_name)
    agent_chain = initialize_agent(tools, chat, agent="conversational-react-description", verbose=True, memory=memory)

    return ChatBot(memory, agent_chain)