diff --git a/providers/google/google.go b/providers/google/google.go index 5a41b16fbf749588a6bdcfa984b1b712ecb60195..097e8c52c0f7f8a0b93a69ba1618b4d151d4c6f6 100644 --- a/providers/google/google.go +++ b/providers/google/google.go @@ -842,9 +842,14 @@ func (g *languageModel) Stream(ctx context.Context, call fantasy.Call) (fantasy. finishReason = fantasy.FinishReasonStop } + var finalUsage fantasy.Usage + if usage != nil { + finalUsage = *usage + } + yield(fantasy.StreamPart{ Type: fantasy.StreamPartTypeFinish, - Usage: *usage, + Usage: finalUsage, FinishReason: finishReason, }) }, nil @@ -1088,9 +1093,14 @@ func (g *languageModel) streamObjectWithJSONMode(ctx context.Context, call fanta finishReason = fantasy.FinishReasonStop } + var finalUsage fantasy.Usage + if usage != nil { + finalUsage = *usage + } + yield(fantasy.ObjectStreamPart{ Type: fantasy.ObjectStreamPartTypeFinish, - Usage: *usage, + Usage: finalUsage, FinishReason: finishReason, }) } else if streamErr == nil && lastParsedObject == nil {