import os
import re
from dotenv import load_dotenv
load_dotenv()

from langchain.agents.openai_assistant import OpenAIAssistantRunnable
from langchain.agents import AgentExecutor

from langchain.schema import HumanMessage, AIMessage

import gradio

api_key = os.getenv('OPENAI_API_KEY_DAVID')
extractor_agent = os.getenv('ASSISTANT_ID_DECASTRO')

extractor_llm = OpenAIAssistantRunnable(assistant_id=extractor_agent, api_key=api_key, as_agent=True)


def remove_citation(text):
    # Define the regex pattern to match the citation format 【number†text】
    pattern = r"【\d+†\w+】"
    # Replace the pattern with an empty string
    return re.sub(pattern, "📚", text)

def predict(message, history):
    history_langchain_format = []
    for human, ai in history:
        history_langchain_format.append(HumanMessage(content=human))
        history_langchain_format.append(AIMessage(content=ai))
    history_langchain_format.append(HumanMessage(content=message))
    gpt_response = extractor_llm.invoke({"content": message})
    output = gpt_response.return_values["output"]
    non_cited_output = remove_citation(output)
    return non_cited_output

chat = gradio.ChatInterface(predict, title="DeCastro Bot", description="DeCastro whatsapp bot.")
chat.launch(share=True)