from langchain_huggingface import HuggingFaceEndpoint as HF
from subprocess import Popen, PIPE as P
from langchain_experimental.tools.python.tool import PythonREPLTool as PYT
from langchain.agents import load_tools,create_structed_chat_agent as Agent,AgentExecutor as Ex
from langchain.agents.agent_toolkits import create_retriever_tool as crt
from langchain_community.agent_toolkits import FileManagementToolkit as FMT
from langchain.tools import Tool,YoutubeSearchTool as YTS
from langchain.memory import ConversationalBufferMomory as MEM,RedisChatHistory as HIS
from langchain.schema import SystemMessage as SM,HumanMessage as HM
from langchain import hub
import chainlit as cl
import os
from langchain.retrievers import WikipediaRetriever as Wiki
def terminal(c):
    a=Popen(c,shell=True,stdin=P,stdout=P,stderr=P)
    return a.stdout.read()+a.stderr.read()
tools=FMT().get_tools()
tools.append(PYT())
tools.append(YTS())
tools.extend(load_tools(["requests"]))
tools.extend(load_tools(["llm-math","ddg-search"]))
tools.append(Tool.from_function(func=terminal,name="terminal"))
tools.append(crt(name="wiki",description="위키 백과를 검색하여 정보를 가져온다",retriever=Wiki(lang="ko",top_k_results=1)))
llm=HF(repo_id="peterpeter8585/syai4.0")
prompt=hub.pull("hwchase17/structed-chat-agent")
@cl.on_chat_start
async def on_chat_start():
    thread_id=None
    while not thread_id:
        res=await cl.AskUserMessage(content="Please enter the username and the id.ex) peterpeter9595_chat1").send()
        if res:
            thread_id=res["content"]
    history=HIS(session_id=thread_id, url=os.environ["URL"])
    memory=MEM(chat_memory=history,memory_key="history")
    m_r=memory.load_memory_variables({})["history"]
    agent=EX(agent=Agent(llm,tools,prompt,memory_key="history"),tools=tools,verbose=True,handle_parsing_errors=True)
    
    for m in m_r:
        if isinstance(m,HM):
            await cl.Message(author="User",content=f"{m.content}").send()
        else:
            await cl.Message(author="ChatBot",content=f"{m.content}").send()
    cl.user_session.set("agent",agent)
@cl.on_message
async def on_message(message:str):
    exe=cl.user_session.get("agent")
    result=exe(message)
    await cl.Message(content=result["output"]).send()