From 89286b7e734223275c15b020bed389e42ffc8200 Mon Sep 17 00:00:00 2001
From: Yiyang Kang <kkyy@kkyy.me>
Date: Mon, 20 Mar 2023 21:45:59 +0800
Subject: [PATCH] refactor: cleanup redundant code

---
 assistant.go        |  8 ++---
 bot.go              |  9 +++++
 botcmd_translate.go | 82 ++++++++++-----------------------------------
 3 files changed, 28 insertions(+), 71 deletions(-)

diff --git a/assistant.go b/assistant.go
index d5396e5..1849115 100644
--- a/assistant.go
+++ b/assistant.go
@@ -12,7 +12,7 @@ import (
 	tele "gopkg.in/telebot.v3"
 )
 
-var assistantWritingSign = "\n...📝"
+var assistantWritingSign = "\n... 📝"
 
 func matchAssistantConversation(botUsr *tele.User, msg *tele.Message) []*tele.Message {
 	// A thread must meet the following conditions to be considered a conversation with the assistant (for now):
@@ -227,11 +227,7 @@ func handleAssistantConversation(c tele.Context, thread []*tele.Message) error {
 		MaxTokens:   2048,
 	}
 
-	typingNotifyCh := make(chan struct{})
-	go func() {
-		defer close(typingNotifyCh)
-		_ = c.Bot().Notify(lastMsg.Chat, tele.Typing)
-	}()
+	typingNotifyCh := setTyping(c)
 
 	var replyMsg *tele.Message
 	reqErr := assistantStreamedResponse(req, func(text string, finished bool) error {
diff --git a/bot.go b/bot.go
index fa7413f..d805bc9 100644
--- a/bot.go
+++ b/bot.go
@@ -117,3 +117,12 @@ func stickerFromID(id string) *tele.Sticker {
 		},
 	}
 }
+
+func setTyping(c tele.Context) chan error {
+	resultCh := make(chan error, 1)
+	go func() {
+		defer close(resultCh)
+		resultCh <- c.Bot().Notify(c.Chat(), tele.Typing)
+	}()
+	return resultCh
+}
diff --git a/botcmd_translate.go b/botcmd_translate.go
index d46b5e7..0ba1fbc 100644
--- a/botcmd_translate.go
+++ b/botcmd_translate.go
@@ -3,7 +3,6 @@ package main
 import (
 	"regexp"
 	"strings"
-	"time"
 
 	"github.com/samber/lo"
 	tele "gopkg.in/telebot.v3"
@@ -72,13 +71,6 @@ func handleTranslateBtn(c tele.Context) error {
 		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{
@@ -93,67 +85,27 @@ func handleTranslateBtn(c tele.Context) error {
 		},
 		Temperature: lo.ToPtr(0.2),
 	}
-	logger.Debugf("Openai chat request: %#+v", req)
 
-	resp, err := ai.ChatCompletionStream(req)
+	actionCh := setTyping(c)
+	err := assistantStreamedResponse(req, func(text string, finished bool) error {
+		var err error
+		<-actionCh
+		if finished {
+			retryBtn := translateBtnRetry
+			retryBtn.Data = targetLang
+			respMenu := &tele.ReplyMarkup{}
+			respMenu.Inline(respMenu.Row(retryBtn))
+
+			msg, err = c.Bot().Edit(msg, text, respMenu)
+		} else {
+			msg, err = c.Bot().Edit(msg, text)
+		}
+		return err
+	})
 	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
+	return nil
 }