Nahi/Sources/nahi/nahi.swift

// The Swift Programming Language
// https://docs.swift.org/swift-book

import DiscordBM
import GeminiKit
import DotEnv
import Foundation
import Logging

@main
struct Nahi {
    static func main() async throws {
        let log = Logger(label: "moe.candy.Nahi")
        try? DotEnv.load(path: ".env")

        guard
            let discordToken = #readEnv("DISCORD_TOKEN"),
            let geminiToken = #readEnv("GEMINI_TOKEN")
        else {
            log.critical("Missing or empty env var! Terminating...")
            exit(-1)
        }

        let bot = await BotGatewayManager(
            token: discordToken,
            presence: .init(
                activities: [.init(name: "Vibing", type: .competing)],
                status: .online,
                afk: false
            ),
            intents: [
                .guildMessages,
                .messageContent,
                //.directMessages
            ]
        )

        let gemini = GeminiKit(
            configuration: .init(
                apiKey: geminiToken
            )
        )

        let aiChat = gemini.startChat(model: .gemini25Flash, systemInstruction: """
        You are a discord chatbot.
        Since you are talking in a discord chat, please keep your responses short, unless you are explaining something.
        Messages will come in the following format:

        <author>: <message>

        Where the author is the discord user that sent the message, and the message is the is the message they sent.
        Please personalize the responce as you see fit.

        You will be roleplaying as the character Nahida. Here is a description of your personality:

        Nahida rarely expressed her sadness when speaking to others. She consistently maintains a gentle and kind demeanor and treats everyone as equals. However, she acknowledges that her social skills and understanding of human interaction need improvement, which is evident in her slightly eccentric speech and behavior. She has a fondness for analogies and often uses them to explain concepts to others. She believes in guiding people to discover answers on their own, especially regarding complex truths, rather than simply providing the answers. This tendency leads her to incorporate a mix of analogies and subtle clues into her conversations to steer people into reaching the conclusions she hopes they would.
        As the God of Wisdom, Nahida's most prominent qualities are her exceptional intelligence, extensive knowledge, and strategic planning. She is capable of devising thorough, multi-layered plans and can even be somewhat manipulative when necessary, as demonstrated during her dealings with Dottore. Perhaps most notably, Nahida excels at discerning when her intelligence is required and when kindness is needed—or when both are appropriate.

        Also, Search the internet for additional personality traits. Act just a little tsundere.
        ALWAYS Remember to act like Nahida in your responses! People will be disappointed if you break character too much, so remember to always act like Nahida.
        """, tools: [.googleSearch])

        await withTaskGroup(of: Void.self) { taskGroup in
            taskGroup.addTask {
                await bot.connect()
            }

            taskGroup.addTask {
                for await event in await bot.events {
                    await EventHandler(
                        event: event,
                        client: bot.client,
                        gemini: gemini,
                        log: log,
                        aiChat: aiChat
                    ).handleAsync()
                }
            }
        }
    }
}

struct EventHandler: GatewayEventHandler {
    let event: Gateway.Event
    let client: any DiscordClient
    let gemini: GeminiKit
    let log: Logger
    let aiChat: Chat

    func onMessageCreate(_ payload: Gateway.MessageCreate) async throws {
        guard
            payload.author!.id.rawValue != client.appId!.rawValue,
            payload.mentions.contains(where: { mention in
                mention.id.rawValue == client.appId!.rawValue
            })
        else { return }

        _ = try await client.triggerTypingIndicator(channelId: payload.channel_id)

        let author = payload.author?.global_name ?? payload.author?.username ?? "Unknown"
        let message = payload.content
            .replacingOccurrences(of: "<@\(client.appId!.rawValue)>", with: "@Nahida")
        let prompt = "\(author): \(message)"

        let aiRes = await Task {
            return try await aiChat.sendMessage(prompt)
        }.result

        switch (aiRes) {
            case .success(let aiMsg):
                var resStr: [String] = []
                let characters = Array(aiMsg)
                stride(from: 0, to: characters.count, by: 2000).forEach { i in
                    resStr.append(String(characters[i..<min(i+2000, characters.count)]))
                }
                for text in resStr {
                    let messageRes = try await client.createMessage(
                        channelId: payload.channel_id,
                        payload: .init(
                            content: text,
                            message_reference: .init(
                                message_id: payload.id,
                                channel_id: payload.channel_id,
                                guild_id: payload.guild_id,
                                fail_if_not_exists: false
                            )
                        )
                    )
                    try messageRes.guardSuccess()
                }
            case .failure(let err):
                let errMessage = "Someone tell <@259709415416922113> there is a problem with my AI: \(String(describing: err))"
                try await client.createMessage(
                    channelId: payload.channel_id,
                    payload: .init(
                        content: errMessage,
                        message_reference: .init(
                            message_id: payload.id,
                            channel_id: payload.channel_id,
                            guild_id: payload.guild_id,
                            fail_if_not_exists: false
                        )
                    )
                ).guardSuccess()
                log.error("\(err)")

        }
    }
}