diff --git a/handlers/dify.py b/handlers/dify.py index 9bbf642..c78d8d5 100644 --- a/handlers/dify.py +++ b/handlers/dify.py @@ -1,61 +1,41 @@ -from os import environ +import json +import time +import re from telebot import TeleBot from telebot.types import Message -from expiringdict import ExpiringDict from . import * + # TODO: update requirements.txt and setup tools # pip install dify-client from dify_client import ChatClient -from telegramify_markdown import convert from telegramify_markdown.customize import markdown_symbol # If you want, Customizing the head level 1 symbol markdown_symbol.head_level_1 = "📌" markdown_symbol.link = "🔗" # If you want, Customizing the link symbol -DIFY_API_KEY = environ.get("DIFY_API_KEY") - -if DIFY_API_KEY: - client = ChatClient(api_key=DIFY_API_KEY) - -# Global history cache -dify_player_dict = ExpiringDict(max_len=1000, max_age_seconds=600) -dify_player_c = ExpiringDict( - max_len=1000, max_age_seconds=600 -) # History cache is supported by dify cloud conversation_id. - def dify_handler(message: Message, bot: TeleBot) -> None: - """dify : /dify """ + """dify : /dify API_Key """ m = message.text.strip() - c = None - player_message = [] - # restart will lose all TODO - if str(message.from_user.id) not in dify_player_dict: - dify_player_dict[str(message.from_user.id)] = ( - player_message # for the imuutable list - ) - else: - player_message = dify_player_dict[str(message.from_user.id)] - # get c from dify_player_c - c = dify_player_c.get(str(message.from_user.id), None) - if m.strip() == "clear": + if re.match(r"^app-\w+$", m, re.IGNORECASE): bot.reply_to( message, - "just clear your dify messages history", + "Thanks!\nFor conversation, please make a space between your API_Key and your question.", ) - player_message.clear() - c = None return - - if m[:4].lower() == "new ": - m = m[4:].strip() - player_message.clear() - c = None + if re.match(r"^app-[a-zA-Z0-9]+ .*$", m, re.IGNORECASE): + Dify_API_KEY = m.split(" ", 1)[0] + m = m.split(" ", 1)[1] + else: + bot.reply_to(message, "Please provide a valid API key.") + return + client = ChatClient(api_key=Dify_API_KEY) + # Init client with API key m = enrich_text_with_urls(m) @@ -63,49 +43,45 @@ def dify_handler(message: Message, bot: TeleBot) -> None: # show something, make it more responsible reply_id = bot_reply_first(message, who, bot) - player_message.append({"role": "user", "content": m}) - # keep the last 5, every has two ask and answer. - if len(player_message) > 10: - player_message = player_message[2:] - - dify_reply_text = "" try: r = client.create_chat_message( inputs={}, query=m, user=str(message.from_user.id), - response_mode="blocking", - conversation_id=c, + response_mode="streaming", ) - j = r.json() - - content = j.get("answer", None) - # get c by j.get then save c to dify_player_c - dify_player_c[str(message.from_user.id)] = j.get("conversation_id", None) - if not content: - dify_reply_text = f"{who} did not answer." - player_message.pop() - else: - dify_reply_text = content - player_message.append( - { - "role": "assistant", - "content": dify_reply_text, - } - ) + s = "" + start = time.time() + overall_start = time.time() + for chunk in r.iter_lines(decode_unicode=True): + chunk = chunk.split("data:", 1)[-1] + if chunk.strip(): + chunk = json.loads(chunk.strip()) + answer_chunk = chunk.get("answer", "") + s += answer_chunk + if time.time() - start > 1.5: + start = time.time() + bot_reply_markdown(reply_id, who, s, bot, split_text=False) + if time.time() - overall_start > 120: # Timeout + s += "\n\nTimeout" + break + # maybe not complete + try: + bot_reply_markdown(reply_id, who, s, bot) + except: + pass except Exception as e: print(e) bot.reply_to(message, "answer wrong maybe up to the max token") # pop my user - player_message.pop() return # reply back as Markdown and fallback to plain text if failed. - bot_reply_markdown(reply_id, who, dify_reply_text, bot) + bot_reply_markdown(reply_id, who, s, bot) -if DIFY_API_KEY: +if True: def register(bot: TeleBot) -> None: bot.register_message_handler(dify_handler, commands=["dify"], pass_bot=True)