dify_example.py

import json
import discord
from discord.ext import commands
from dotenv import load_dotenv
import os
import aiohttp

# Load environment variables
load_dotenv()

DISCORD_BOT_TOKEN = os.getenv('DISCORD_BOT_TOKEN')
DIFY_API_KEY = os.getenv('DIFY_API_KEY')
DIFY_ENDPOINT = os.getenv('DIFY_ENDPOINT')


# 파일 카테고리별 확장자 정의
DIFY_FILE_CATEGORY_EXTENSIONS = {
    "document": [  # 문서 파일 확장자
        "TXT", "MD", "MDX", "MARKDOWN", "PDF", "HTML", "XLSX", "XLS",
        "DOC", "DOCX", "CSV", "EML", "MSG", "PPTX", "PPT", "XML", "EPUB"
    ],
    "image": ["JPG", "JPEG", "PNG", "GIF", "WEBP", "SVG"],  # 이미지 파일 확장자
    "audio": ["MP3", "M4A", "WAV", "WEBM", "AMR", "MPGA"],  # 오디오 파일 확장자
    "video": ["MP4", "MOV", "MPEG", "MPGA"]  # 비디오 파일 확장자
}

# 파일 카테고리 목록
DIFY_FILE_CATEGORIES = list(DIFY_FILE_CATEGORY_EXTENSIONS.keys())

conversations_db = {}
intents = discord.Intents.default()
intents.message_content = True

bot = commands.Bot(command_prefix='!', intents=intents)


# 파일 이름을 받아 해당 파일의 카테고리를 반환하는 함수
def get_dify_file_category(file_name: str) -> str:
    extension = file_name.split(".")[-1].upper() if "." in file_name else ""
    for category, extensions in DIFY_FILE_CATEGORY_EXTENSIONS.items():
        if extension in extensions:
            return category
    return "custom"  # 정의된 카테고리에 없는 경우 "custom" 반환


@bot.event
async def on_ready():
    print(f'{bot.user} 이 디스코드에 접속했습니다!')


@bot.event
async def on_message(message: discord.Message):
    # 봇이 보낸 메시지는 무시
    if message.author == bot.user:
        return

    # 봇이 멘션되었는지 확인
    if bot.user.mentioned_in(message):
        # 멘션을 제거하고 메시지 내용만 추출
        content = message.content.replace(f'<@{bot.user.id}>', '').strip()

        # 메시지 보내는 동안 타이핑 표시
        async with message.channel.typing():
            try:
                files = [
                    {
                        "url": attachment.url,
                        "type": get_dify_file_category(attachment.filename),
                        "transfer_method": "remote_url"
                    }
                    for attachment in message.attachments
                ]
                # 메세지 작성자의 Conversation ID 참조, 첫 메세지일 경우 None
                conversation_id = conversations_db.get(message.author.id, None)

                # Dify API 호출
                url = f"{DIFY_ENDPOINT}/chat-messages"
                headers = {
                    "Authorization": f"Bearer {DIFY_API_KEY}",
                    "Content-Type": "application/json"
                }
                payload = {
                    "inputs": {},
                    "query": content,
                    "user": message.author.id,  # 메세지 작성자의 ID
                    "response_mode": "streaming",
                    "conversation_id": conversation_id,  # 메세지 작성자의 Conversation ID
                    "files": files  # 첨부 파일 목록
                }
                answer = ""
                async with aiohttp.ClientSession() as session:
                    async with session.post(url, json=payload, headers=headers) as response:

                        # aiohttp 데이터 스트림 처리 (iter_lines 처럼 동작)
                        async for chunk in response.content:
                            # 빈 데이터는 무시
                            if not chunk:
                                continue

                            decoded_chunk = chunk.decode("utf-8")
                            try:
                                # SSE 이벤트 파싱
                                data = json.loads(decoded_chunk.strip()[6:])
                            except json.decoder.JSONDecodeError:
                                continue

                            # print(f"🔥 data: {data}")
                            match data["event"]:
                                case "message":
                                    # 메세지 이벤트의 경우 답변 추가
                                    answer += data["answer"]
                                case "message_end":
                                    # 메세지 종료 이벤트의 경우 Conversation ID 저장
                                    conversations_db[message.author.id] = data["conversation_id"]
                                case "agent_message":
                                    # 에이전트 메세지 이벤트의 경우 답변 추가
                                    answer += data["answer"]
                                case "agent_thought":
                                    print(f"🔥 agent_thought: {data}")
                                    pass
                                case "message_file":
                                    # 메세지 파일 이벤트의 경우 파일 URL 첨부
                                    await message.channel.send(data['url'])
                                case "error":
                                    # 오류 이벤트의 경우 오류 메시지 출력
                                    await message.channel.send(f"Error: {data}")
                                case _:
                                    pass

                            # 답변이 1000자 이상일 경우 1000자씩 잘라서 메시지 보내기
                            if len(answer) > 1000:
                                await message.channel.send(answer)
                                answer = ""

                        # 보내지 않은 답변이 남아있을 경우 마지막 메시지로 보내기
                        if answer:
                            await message.channel.send(answer)

            except Exception as e:
                import traceback
                traceback.print_exc()
                await message.channel.send(f"Error: {e}")

    # 다른 명령어들을 처리하기 위해 이벤트를 계속 진행
    await bot.process_commands(message)


@bot.command()
async def clear(ctx: commands.Context):
    """!clear 명령어를 사용하면 대화 기록이 초기화됩니다."""

    # 대화 기록 초기화
    conversations_db.pop(ctx.author.id, None)
    await ctx.send("대화 기록이 초기화되었습니다.")


bot.run(DISCORD_BOT_TOKEN)