diff --git a/controller/relay-utils.go b/controller/relay-utils.go index c22144f..4f75efb 100644 --- a/controller/relay-utils.go +++ b/controller/relay-utils.go @@ -145,47 +145,48 @@ func countTokenMessages(messages []Message, model string) (int, error) { for _, message := range messages { tokenNum += tokensPerMessage tokenNum += getTokenNum(tokenEncoder, message.Role) - var arrayContent []MediaMessage - if err := json.Unmarshal(message.Content, &arrayContent); err != nil { - - var stringContent string - if err := json.Unmarshal(message.Content, &stringContent); err != nil { - return 0, err - } else { - tokenNum += getTokenNum(tokenEncoder, stringContent) - if message.Name != nil { - tokenNum += tokensPerName - tokenNum += getTokenNum(tokenEncoder, *message.Name) - } - } - } else { - for _, m := range arrayContent { - if m.Type == "image_url" { - var imageTokenNum int - if str, ok := m.ImageUrl.(string); ok { - imageTokenNum, err = getImageToken(&MessageImageUrl{Url: str, Detail: "auto"}) - } else { - imageUrlMap := m.ImageUrl.(map[string]interface{}) - detail, ok := imageUrlMap["detail"] - if ok { - imageUrlMap["detail"] = detail.(string) - } else { - imageUrlMap["detail"] = "auto" - } - imageUrl := MessageImageUrl{ - Url: imageUrlMap["url"].(string), - Detail: imageUrlMap["detail"].(string), - } - imageTokenNum, err = getImageToken(&imageUrl) - } - if err != nil { - return 0, err - } - - tokenNum += imageTokenNum - log.Printf("image token num: %d", imageTokenNum) + if len(message.Content) > 0 { + var arrayContent []MediaMessage + if err := json.Unmarshal(message.Content, &arrayContent); err != nil { + var stringContent string + if err := json.Unmarshal(message.Content, &stringContent); err != nil { + return 0, err } else { - tokenNum += getTokenNum(tokenEncoder, m.Text) + tokenNum += getTokenNum(tokenEncoder, stringContent) + if message.Name != nil { + tokenNum += tokensPerName + tokenNum += getTokenNum(tokenEncoder, *message.Name) + } + } + } else { + for _, m := range arrayContent { + if m.Type == "image_url" { + var imageTokenNum int + if str, ok := m.ImageUrl.(string); ok { + imageTokenNum, err = getImageToken(&MessageImageUrl{Url: str, Detail: "auto"}) + } else { + imageUrlMap := m.ImageUrl.(map[string]interface{}) + detail, ok := imageUrlMap["detail"] + if ok { + imageUrlMap["detail"] = detail.(string) + } else { + imageUrlMap["detail"] = "auto" + } + imageUrl := MessageImageUrl{ + Url: imageUrlMap["url"].(string), + Detail: imageUrlMap["detail"].(string), + } + imageTokenNum, err = getImageToken(&imageUrl) + } + if err != nil { + return 0, err + } + + tokenNum += imageTokenNum + log.Printf("image token num: %d", imageTokenNum) + } else { + tokenNum += getTokenNum(tokenEncoder, m.Text) + } } } } diff --git a/controller/relay.go b/controller/relay.go index 0ee7522..c835497 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -13,9 +13,11 @@ import ( ) type Message struct { - Role string `json:"role"` - Content json.RawMessage `json:"content"` - Name *string `json:"name,omitempty"` + Role string `json:"role"` + Content json.RawMessage `json:"content"` + Name *string `json:"name,omitempty"` + ToolCalls any `json:"tool_calls,omitempty"` + ToolCallId string `json:"tool_call_id,omitempty"` } type MediaMessage struct {