Skip to content

Commit

Permalink
fix: fix base64 images additional billing
Browse files Browse the repository at this point in the history
  • Loading branch information
zmh-program committed Feb 15, 2024
1 parent 1128f00 commit b95bf94
Show file tree
Hide file tree
Showing 3 changed files with 9 additions and 4 deletions.
2 changes: 1 addition & 1 deletion adapter/azure/processor.go
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ func formatMessages(props *ChatProps) interface{} {
return nil
}

props.Buffer.AddImage(obj)
props.Buffer.AddImage(obj, url)

return &MessageContent{
Type: "image_url",
Expand Down
2 changes: 1 addition & 1 deletion adapter/chatgpt/processor.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ func formatMessages(props *ChatProps) interface{} {
if err != nil {
globals.Info(fmt.Sprintf("cannot process image: %s (source: %s)", err.Error(), url))
} else {
props.Buffer.AddImage(obj)
props.Buffer.AddImage(obj, url)
}

return &MessageContent{
Expand Down
9 changes: 7 additions & 2 deletions utils/buffer.go
Original file line number Diff line number Diff line change
Expand Up @@ -58,10 +58,15 @@ func (b *Buffer) GetChunk() string {
return b.Latest
}

func (b *Buffer) AddImage(image *Image) {
func (b *Buffer) AddImage(image *Image, source string) {
b.Images = append(b.Images, *image)

b.Quota += float32(image.CountTokens(b.Model)) * b.Charge.GetInput()
if b.Charge.IsBillingType(globals.TokenBilling) {
b.Quota += float32(image.CountTokens(b.Model)) * b.Charge.GetInput()

// remove tokens from image source
b.Quota -= CountInputToken(b.Charge, b.Model, []globals.Message{{Content: source, Role: globals.User}})
}
}

func (b *Buffer) GetImages() Images {
Expand Down

0 comments on commit b95bf94

Please sign in to comment.