tgbot_misaka_5882f7/botcmd_translate.go

package main

import (
	"regexp"
	"strings"
	"time"

	"github.com/samber/lo"
	tele "gopkg.in/telebot.v3"

	"git.gensokyo.cafe/kkyy/tgbot_misaka_5882f7/openai"
	"git.gensokyo.cafe/kkyy/tgbot_misaka_5882f7/openai/prompts"
)

var (
	translateMenu = &tele.ReplyMarkup{}

	translateBtnZhTw = translateMenu.Data("繁中", "btn_tr_zhtw", "Taiwanese Chinese")
	translateBtnZhCn = translateMenu.Data("简中", "btn_tr_zhcn", "Mandarin Chinese")
	translateBtnEn   = translateMenu.Data("English", "btn_tr_en", "English")
	translateBtnJa   = translateMenu.Data("日本語", "btn_tr_ja", "Japanese")

	translateBtnRetry = translateMenu.Data("Try again", "btn_tr_retry")

	translateBtns = []*tele.Btn{
		&translateBtnZhTw, &translateBtnZhCn, &translateBtnEn, &translateBtnJa,
		&translateBtnRetry,
	}

	translateCmdRe = regexp.MustCompile(`^\s*\/tr(anslate)?(@\S*)?\s*`)
)

func init() {
	translateMenu.Inline(
		translateMenu.Row(translateBtnZhTw, translateBtnZhCn),
		translateMenu.Row(translateBtnEn, translateBtnJa),
	)
}

func handleTranslateCmd(c tele.Context) error {
	msg := c.Message()
	if msg == nil {
		return nil
	}
	if msg.ReplyTo != nil {
		msg = msg.ReplyTo
	}

	payload := strings.TrimSpace(translateCmdRe.ReplaceAllString(msg.Text, ""))
	if payload == "" {
		return c.Reply("Usage: `/tr <text>`",
			&tele.SendOptions{ParseMode: tele.ModeMarkdown},
			tele.Silent,
		)
	}

	_, err := c.Bot().Reply(msg, "Sure. To what language?", tele.Silent, translateMenu)
	return err
}

func handleTranslateBtn(c tele.Context) error {
	msg := c.Message()
	if msg == nil || msg.ReplyTo == nil {
		return nil
	}
	origMsg := msg.ReplyTo
	targetLang := c.Data()
	txt := origMsg.Text
	payload := strings.TrimSpace(translateCmdRe.ReplaceAllString(txt, ""))

	if targetLang == "" || payload == "" {
		return nil
	}

	// pretend to be typing
	if err := c.Bot().Notify(msg.Chat, tele.Typing); err != nil {
		logger.Warnf("failed to send typing action: %v", err)
	}

	ai := openai.NewClient(config.OpenAIApiKey)

	req := openai.ChatRequest{
		Model: openai.ModelGpt0305Turbo,
		Messages: []openai.ChatMessage{
			{
				Role:    openai.ChatRoleSystem,
				Content: prompts.Translate(targetLang),
			},
			{
				Role:    openai.ChatRoleUser,
				Content: payload,
			},
		},
		Temperature: lo.ToPtr(0.2),
	}
	logger.Debugf("Openai chat request: %#+v", req)

	resp, err := ai.ChatCompletionStream(req)
	if err != nil {
		logger.Errorf("failed to translate: req: %#+v, err: %v", req, err)
		_, _ = c.Bot().Reply(origMsg, stickerFromID(stickerPanic), tele.Silent)
		return err
	}

	respBuilder := strings.Builder{}
	minWait := time.After(1 * time.Second)
	for {
		var (
			nNewChunk        int
			finished         bool
			minWaitSatisfied bool
		)
	Drain:
		for {
			select {
			case chunk, ok := <-resp.Stream:
				if !ok {
					finished = true
					break Drain
				}
				nNewChunk += 1
				respBuilder.WriteString(chunk)
			default:
				if minWaitSatisfied {
					break Drain
				}
				<-minWait
				minWaitSatisfied = true
			}
		}

		if nNewChunk == 0 {
			if chunk, ok := <-resp.Stream; !ok {
				finished = true
			} else {
				respBuilder.WriteString(chunk)
			}
		}
		if finished {
			break
		}

		respoText := respBuilder.String() + assistantWritingSign
		minWait = time.After(691 * time.Millisecond) // renew the timer
		if msg, err = c.Bot().Edit(msg, respoText, tele.Silent); err != nil {
			logger.Warnf("failed to edit the temporary message: %v", err)
			break
		}
		logger.Debugf("... message edited")
	}

	respText := respBuilder.String()
	retryBtn := translateBtnRetry
	retryBtn.Data = targetLang
	respMenu := &tele.ReplyMarkup{}
	respMenu.Inline(respMenu.Row(retryBtn))
	_, err = c.Bot().Edit(msg, respText, tele.Silent, respMenu)
	return err
}