import os
import re
import time
from typing import Any

from dotenv import load_dotenv
from langchain.callbacks.base import BaseCallbackHandler
from langchain.chat_models import ChatOpenAI
from langchain.schema import LLMResult
from slack_bolt import App
from slack_bolt.adapter.socket_mode import SocketModeHandler

CHAT_UPDATE_INTERVAL_SEC = 1

load_dotenv()

# ボットトークンを使ってアプリを初期化します
app = App(
    signing_secret=os.environ["SLACK_SIGNING_SECRET"],
    token=os.environ["SLACK_BOT_TOKEN"],
    process_before_response=True,
)


class SlackStreamingCallbackHandler(BaseCallbackHandler):
    last_send_time = time.time()
    message = ""

    def __init__(self, channel, ts):
        self.channel = channel
        self.ts = ts

    def on_llm_new_token(self, token: str, **kwargs) -> None:
        self.message += token

        now = time.time()
        if now - self.last_send_time > CHAT_UPDATE_INTERVAL_SEC:
            self.last_send_time = now
            app.client.chat_update(
                channel=self.channel, ts=self.ts, text=f"{self.message}..."
            )

    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> Any:
        app.client.chat_update(channel=self.channel, ts=self.ts, text=self.message)


@app.event("app_mention")
def handle_mention(event, say):
    channel = event["channel"]
    thread_ts = event["ts"]
    message = re.sub("<@.*>", "", event["text"])

    result = say("\n\nTyping...", thread_ts=thread_ts)
    ts = result["ts"]

    callback = SlackStreamingCallbackHandler(channel=channel, ts=ts)
    llm = ChatOpenAI(
        model_name=os.environ["OPENAI_API_MODEL"],
        temperature=os.environ["OPENAI_API_TEMPERATURE"],
        streaming=True,
        callbacks=[callback],
    )

    llm.predict(message)


# ソケットモードハンドラーを使ってアプリを起動します
if __name__ == "__main__":
    SocketModeHandler(app, os.environ["SLACK_APP_TOKEN"]).start()

langchain.chat_models.ChatOpenAI
[((347, 360), 'dotenv.load_dotenv', 'load_dotenv', ([], {}), '()\n', (358, 360), False, 'from dotenv import load_dotenv\n'), ((392, 518), 'slack_bolt.App', 'App', ([], {'signing_secret': "os.environ['SLACK_SIGNING_SECRET']", 'token': "os.environ['SLACK_BOT_TOKEN']", 'process_before_response': '(True)'}), "(signing_secret=os.environ['SLACK_SIGNING_SECRET'], token=os.environ[\n    'SLACK_BOT_TOKEN'], process_before_response=True)\n", (395, 518), False, 'from slack_bolt import App\n'), ((610, 621), 'time.time', 'time.time', ([], {}), '()\n', (619, 621), False, 'import time\n'), ((1365, 1399), 're.sub', 're.sub', (['"""<@.*>"""', '""""""', "event['text']"], {}), "('<@.*>', '', event['text'])\n", (1371, 1399), False, 'import re\n'), ((1558, 1704), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'model_name': "os.environ['OPENAI_API_MODEL']", 'temperature': "os.environ['OPENAI_API_TEMPERATURE']", 'streaming': '(True)', 'callbacks': '[callback]'}), "(model_name=os.environ['OPENAI_API_MODEL'], temperature=os.\n    environ['OPENAI_API_TEMPERATURE'], streaming=True, callbacks=[callback])\n", (1568, 1704), False, 'from langchain.chat_models import ChatOpenAI\n'), ((837, 848), 'time.time', 'time.time', ([], {}), '()\n', (846, 848), False, 'import time\n'), ((1826, 1879), 'slack_bolt.adapter.socket_mode.SocketModeHandler', 'SocketModeHandler', (['app', "os.environ['SLACK_APP_TOKEN']"], {}), "(app, os.environ['SLACK_APP_TOKEN'])\n", (1843, 1879), False, 'from slack_bolt.adapter.socket_mode import SocketModeHandler\n')]