package main import ( "crypto/md5" "encoding/base64" "strconv" "strings" "time" "unicode" "git.gensokyo.cafe/kkyy/tgbot_misaka_5882f7/openai" "git.gensokyo.cafe/kkyy/tgbot_misaka_5882f7/openai/prompts" "github.com/samber/lo" tele "gopkg.in/telebot.v3" ) 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): // It has only two participants: the assistant and the user. // Or, it has only one participant, which is the user. In this case, it has to be in a private chat, or the assistant must be mentioned if it's in a group chat. // No bot commands, images, stickers, or other message types involved. participants := map[string]struct{}{} hasAssistant := false hasUser := false assistantMentioned := false isPrivateChat := true thread, err := getCachedThread(msg) if err != nil { return nil } for _, msg := range thread { for _, ent := range msg.Entities { if ent.Type == tele.EntityCommand { return nil } } from := msg.Sender chat := msg.Chat if from == nil || chat == nil { return nil } participant := strconv.FormatInt(chat.ID, 10) + ":" + strconv.FormatInt(from.ID, 10) participants[participant] = struct{}{} if chat.Type != tele.ChatPrivate { isPrivateChat = false } if !lo.Contains([]tele.ChatType{tele.ChatPrivate, tele.ChatGroup, tele.ChatSuperGroup}, chat.Type) { return nil } if msg.Text == "" { return nil } if from.ID == botUsr.ID { hasAssistant = true if strings.HasSuffix(msg.Text, assistantWritingSign) { return nil } } else { if from.IsBot { return nil } hasUser = true // Only @bot_user_name in the beginning will be counted as "mention" mentionStr := "@" + botUsr.Username if strings.HasPrefix(strings.TrimSpace(msg.Text), mentionStr) { assistantMentioned = true } } } if len(participants) > 2 { return nil } if hasAssistant && hasUser { return thread } if hasUser && (isPrivateChat || assistantMentioned) { return thread } return nil } type assistantStreamedResponseCb func(text string, finished bool) error // TODO interrupt response with context. func assistantStreamedResponse(request openai.ChatRequest, cb assistantStreamedResponseCb) error { logger.Debugw("Openai chat request", "req", request) ai := openai.NewClient(config.OpenAIApiKey) var ( resp *openai.ChatResponseStream err error ) tries := 2 for { if tries--; tries < 0 { return err } if resp, err = ai.ChatCompletionStream(request); err == nil { break } logger.Warnw("assistant: failed to get response", "error", err) } var nErrs, nUpdates int go func() { respBuilder := strings.Builder{} minWait := time.After(time.Second) for { nUpdates += 1 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 } select { case <-minWait: case <-resp.Done: } minWaitSatisfied = true } } if nNewChunk == 0 { if chunk, ok := <-resp.Stream; !ok { finished = true } else { respBuilder.WriteString(chunk) } } if finished { break } // Renew the timer. // The Telegram API rate limit for group messages is 20 per minute. So we cannot update messages too frequently. minWaitDurSecs := lo.Min([]int{nUpdates, 4}) + nErrs*3 minWait = time.After(time.Duration(minWaitDurSecs) * time.Second) // Send the partial message respText := respBuilder.String() + assistantWritingSign if err := cb(respText, false); err != nil { logger.Warnw("failed to send partial update", "error", err) nErrs += 1 if nErrs > 5 { logger.Warnw("too many errors, aborting") return } continue } logger.Debugf("... message edited") } respText := respBuilder.String() if err = cb(respText, true); err != nil { logger.Warnw("assistant: failed to send message", "error", err) } }() return nil } func handleAssistantConversation(c tele.Context, thread []*tele.Message) error { me := c.Bot().Me lastMsg := thread[len(thread)-1] if err := cacheMessage(lastMsg); err != nil { logger.Warnw("failed to cache message", "error", err) } nBytes := 0 // Used to estimated number of tokens. For now we treat 3 bytes as 1 token. nBytesMax := (4096 - 512) * 3 // Leave some space for the response sysMsg := prompts.Assistant() chatReqMsgs := []openai.ChatMessage{ { Role: openai.ChatRoleSystem, Content: sysMsg, }, } nBytes += len(sysMsg) convMsgs := []openai.ChatMessage{} for l := len(thread) - 1; l >= 0; l-- { text := assistantRemoveMention(thread[l].Text, me.Username) textLen := len(text) if textLen+nBytes > nBytesMax { break } nBytes += textLen role := openai.ChatRoleUser from := thread[l].Sender if from != nil && from.ID == me.ID { role = openai.ChatRoleAssistant } convMsgs = append(convMsgs, openai.ChatMessage{ Role: role, Content: text, }) } if len(convMsgs) == 0 { // It turns out that this will never happen because Telegram splits messages when they exceed a certain length. return c.Reply("Your message is too long (Sorry!)") } for l := len(convMsgs) - 1; l >= 0; l-- { chatReqMsgs = append(chatReqMsgs, convMsgs[l]) } req := openai.ChatRequest{ Model: openai.ModelGpt0305Turbo, Messages: chatReqMsgs, Temperature: lo.ToPtr(0.42), MaxTokens: 2048, User: assistantHashUserId(lastMsg.Sender.ID), } typingNotifyCh := setTyping(c) var replyMsg *tele.Message reqErr := assistantStreamedResponse(req, func(text string, finished bool) error { var err error if replyMsg == nil { <-typingNotifyCh replyMsg, err = c.Bot().Reply(lastMsg, text, tele.Silent) } else { replyMsg, err = c.Bot().Edit(replyMsg, text) } if finished && err == nil { replyMsg.ReplyTo = lastMsg // nasty bug if err := cacheMessage(replyMsg); err != nil { logger.Warnw("failed to cache message", "error", err) } } return err }) if reqErr != nil { logger.Errorw("assistant: failed to complete conversation", "error", reqErr) return c.Reply("Sorry, there's a technical issue. 😵💫 Please try again later.", tele.Silent) } return nil } func assistantRemoveMention(msg, name string) string { mentionStr := "@" + name orig := msg msg = strings.TrimLeftFunc(msg, unicode.IsSpace) if strings.HasPrefix(msg, mentionStr) { return msg[len(mentionStr):] } return orig } func assistantHashUserId(uid int64) string { seasoned := []byte("RdnuRPqp66vtbc28QRO0ecKSLKXifz7G9UbXLoyCMpw" + strconv.FormatInt(uid, 10)) hashed := md5.Sum(seasoned) // Don't judge me return base64.URLEncoding.EncodeToString(hashed[:])[:22] }