Writing a Chatbot Powered by Google Gemini Part II

type From struct {
	Id           int    `json:"id"`
	IsBot        bool   `json:"is_bot"`
	FirstName    string `json:"first_name"`
	LastName     string `json:"last_name"`
	LanguageCode string `json:"language_code"`
}

type Message struct {
	MessageId int  `json:"message_id"`
	From      From `json:"from"`
	Chat      struct {
		Id        int    `json:"id"`
		FirstName string `json:"first_name"`
		LastName  string `json:"last_name"`
		Type      string `json:"type"`
	} `json:"chat"`
	Date     int    `json:"date"`
	Text     string `json:"text"`
	Entities []struct {
		Offset int    `json:"offset"`
		Length int    `json:"length"`
		Type   string `json:"type"`
	} `json:"entities"`
	MediaGroupId string `json:"media_group_id"`
	Photo        []struct {
		FileId       string `json:"file_id"`
		FileUniqueId string `json:"file_unique_id"`
		FileSize     int    `json:"file_size"`
		Width        int    `json:"width"`
		Height       int    `json:"height"`
	} `json:"photo"`
	Caption string `json:"caption"`
}

type CallbackQuery struct {
	Id   string `json:"id"`
	From From   `json:"from"`
	Data string `json:"data"`
}

type Update struct {
	UpdateId      int           `json:"update_id"`
	Message       Message       `json:"message"`
	CallbackQuery CallbackQuery `json:"callback_query"`
}

func extractFrom(update Update) int {
	if update.Message.From.Id != 0 {
		return update.Message.From.Id
	}
	return update.CallbackQuery.From.Id
}

var parts []genai.Part
if update.Message.Caption != "" {
    parts = append(parts, genai.Text(update.Message.Caption))
}

var ok bool
parts, ok = appendImagePart(writer, update.Message.Chat.Id, parts, update.Message.Photo[len(update.Message.Photo)-1].FileId)
if !ok {
    return
}

var (
    model = chatClient.GenerativeModel("gemini-pro-vision")
    ctx   = request.Context()
    iter  = model.GenerateContentStream(ctx, parts...)
    sb    strings.Builder
)
model.SafetySettings = contentSafetySettings

generateResponseStream(iter, &sb, writer, update.Message.Chat.Id)

if update.Message.MediaGroupId != "" {
    var (
        mediaGroupKey     = "media_group:" + update.Message.MediaGroupId
        mediaGroupOnceKey = "media_group_once:" + update.Message.MediaGroupId
        once              = rdb.Incr(mediaGroupOnceKey).Val() == 1
        pipe              = rdb.Pipeline()
        values            []any
    )
    if once {
        pipe.LPush(mediaGroupKey, update.Message.Chat.Id)
        sendMessageInlineWithInlineKeyboard(writer, update.Message.Chat.Id, "press done when all photos are uploaded", "{\"inline_keyboard\":[[{\"text\":\"done\",\"callback_data\":\""+mediaGroupKey+"\"}]]}")
    }
    if update.Message.Caption != "" {
        values = append(values, "caption:"+update.Message.Caption)
    }
    values = append(values, "file_id:"+update.Message.Photo[len(update.Message.Photo)-1].FileId)
    pipe.RPush(mediaGroupKey, values...)
    pipe.Expire(mediaGroupOnceKey, timeout)
    pipe.Expire(mediaGroupKey, timeout)
    _, _ = pipe.Exec()
    return
}

var (
    parts        []genai.Part
    entries      = rdb.LRange(update.CallbackQuery.Data, 0, -1).Val()
    chatID, _    = strconv.Atoi(entries[0])
    captionCount int
    captionIdx   int
)
for i, entry := range entries[1:] {
    prefix, val, _ := strings.Cut(entry, ":")
    switch prefix {
    case "caption":
        parts = append(parts, genai.Text(val))
        captionCount++
        captionIdx = i
    case "file_id":
        var ok bool
        parts, ok = appendImagePart(writer, chatID, parts, val)
        if !ok {
            return
        }
    }
}
if captionCount == 1 {
    caption := parts[captionIdx]
    copy(parts[1:captionIdx+1], parts[:captionIdx])
    parts[0] = caption
}

var (
    model = chatClient.GenerativeModel("gemini-pro-vision")
    ctx   = request.Context()
    iter  = model.GenerateContentStream(ctx, parts...)
    sb    strings.Builder
)
model.SafetySettings = contentSafetySettings

if generateResponseStream(iter, &sb, writer, chatID) {
    answerCallbackQueryInline(writer, update.CallbackQuery.Id)
}

func transformResponseError(err error, resp *genai.GenerateContentResponse) error {
	if err != nil {
		return err
	}

	for _, c := range resp.Candidates {
		if c.FinishReason != genai.FinishReasonStop {
			return &genai.BlockedError{
				Candidate: c,
			}
		}
	}
	return nil
}

func generateResponseStream(iter *genai.GenerateContentResponseIterator, sb *strings.Builder, writer http.ResponseWriter, chatID int) bool {
	var (
		messageID int
	)
	for {
		resp, err := iter.Next()
		err = transformResponseError(err, resp)
		if err == iterator.Done {
			return true
		}

		if err != nil {
			if messageID == 0 {
				sendMessageInline(writer, chatID, err.Error())
			} else {
				_ = updateMessage(chatID, messageID, "generate response error: "+err.Error())
			}
			return false
		}

		var text = string(resp.Candidates[0].Content.Parts[0].(genai.Text))
		sb.WriteString(text)
		if messageID == 0 {
			messageID, err = sendMessage(chatID, sb.String())
			if err != nil {
				sendMessageInline(writer, chatID, err.Error())
				return false
			}
		} else {
			err = updateMessage(chatID, messageID, sb.String())
			if err != nil {
				return false
			}
		}
	}
}