import gradio as gr
import os

from langchain import OpenAI, ConversationChain
from langchain.prompts import PromptTemplate
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores.faiss import FAISS
from langchain.docstore.document import Document
from langchain.agents import Tool
from langchain.chains.conversation.memory import ConversationBufferMemory
from langchain.utilities import GoogleSearchAPIWrapper
from langchain.agents import initialize_agent

from langchain.chains.conversation.memory import ConversationEntityMemory
from langchain.chains.conversation.prompt import ENTITY_MEMORY_CONVERSATION_TEMPLATE

from langchain.agents import ZeroShotAgent, Tool, AgentExecutor
from langchain import SerpAPIWrapper, LLMChain

# ツールの準備
search = GoogleSearchAPIWrapper()
tools = [
    Tool(
        name = "Current Search",
        func=search.run,
        description="Use this allways",
    ),
]

# メモリの準備
memory = ConversationBufferMemory(memory_key="chat_history")

# エージェントの準備
llm=OpenAI(model_name = "text-davinci-003",temperature=0)
agent_chain = initialize_agent(
    tools, 
    llm, 
    agent="zero-shot-react-description", 
    verbose=True, 
    memory=memory
)

def chat(message, site,history):
    history = history or []
    #siteの//以前を削除
    site = site.replace("https://","")
    response = ""
    try:
        response = agent_chain.run(input=message+" site:"+site)
    except KeyError:
        if not response:
            response = "not found in the site"
    history.append((message, response))

    return history, history


with gr.Blocks() as demo:
    gr.Markdown("<h3><center>WebSiteChatBotAI</center></h3>")
    gr.Markdown("<p><center>paste web site URL and input question and push Run</center></p>")
    site = gr.Textbox(placeholder="paste URL",label="WebSite")
    chatbot = gr.Chatbot()
    with gr.Row():
        inp = gr.Textbox(placeholder="Question",label =None)
        btn = gr.Button("Run").style(full_width=False)
        state = gr.State()
        agent_state = gr.State()
        btn.click(chat, [inp,site,state],[chatbot, state])
if __name__ == '__main__':
    demo.launch()