add:gemini accumlator

Author: pherzheyu

model/gemini/accumulator.go

@@ -0,0 +1,84 @@
+//
+// Tencent is pleased to support the open source community by making trpc-agent-go available.
+//
+// Copyright (C) 2025 Tencent.  All rights reserved.
+//
+// trpc-agent-go is licensed under the Apache License Version 2.0.
+//
+//
+
+// Package gemini provides Gemini-compatible model implementations.
+package gemini
+
+import (
+	"strings"
+	"time"
+	"trpc.group/trpc-go/trpc-agent-go/model"
+)
+
+// Accumulator accumulates chunks from a stream
+type Accumulator struct {
+	Model            string
+	FullText         strings.Builder
+	ReasoningContent strings.Builder
+	FinishReason     strings.Builder
+	ToolCalls        []model.ToolCall
+	Usage            model.Usage
+}
+
+// Accumulate builds up the Message incrementally from a model.Response. The Message then can be used as
+// any other Message, except with the caveat that the Message.JSON field which normally can be used to inspect
+// the JSON sent over the network may not be populated fully.
+func (a *Accumulator) Accumulate(resp *model.Response) {
+	a.Model = resp.Model
+	if len(resp.Choices) > 0 {
+		for _, choice := range resp.Choices {
+			if choice.FinishReason != nil {
+				a.FinishReason.WriteString(*choice.FinishReason)
+			}
+			if choice.Delta.Content != "" {
+				a.FullText.WriteString(choice.Delta.Content)
+			}
+			if choice.Delta.ReasoningContent != "" {
+				a.ReasoningContent.WriteString(choice.Delta.ReasoningContent)
+			}
+			if len(choice.Delta.ToolCalls) > 0 {
+				a.ToolCalls = append(a.ToolCalls, choice.Delta.ToolCalls...)
+			}
+		}
+	}
+	if resp.Usage != nil {
+		a.Usage.PromptTokens += resp.Usage.PromptTokens
+		a.Usage.CompletionTokens += resp.Usage.CompletionTokens
+		a.Usage.TotalTokens += resp.Usage.TotalTokens
+	}
+}
+
+// BuildResponse builds up the final a model.Response.
+func (a *Accumulator) BuildResponse() *model.Response {
+	now := time.Now()
+	return &model.Response{
+		Model:     a.Model,
+		Created:   now.Unix(),
+		Timestamp: now,
+		Done:      true,
+		Choices: []model.Choice{
+			{
+				Message: model.Message{
+					Content:          a.FullText.String(),
+					ReasoningContent: a.ReasoningContent.String(),
+					ToolCalls:        a.ToolCalls,
+					Role:             model.RoleAssistant,
+				},
+				FinishReason: func() *string {
+					fr := a.FinishReason.String()
+					if fr == "" {
+						return nil
+					}
+					return &fr
+				}(),
+			},
+		},
+		Usage: &a.Usage,
+	}
+}

model/gemini/accumulator_test.go

@@ -0,0 +1,101 @@
+package gemini
+
+import (
+	"github.com/stretchr/testify/assert"
+	"strings"
+	"testing"
+	"trpc.group/trpc-go/trpc-agent-go/model"
+)
+
+func TestAccumulator_BuildResponse(t *testing.T) {
+	finishReason := "FinishReason"
+	content := "Content"
+	reasoningContent := "ReasoningContent"
+	m := &model.Response{
+		Usage: &model.Usage{
+			PromptTokens:     1,
+			CompletionTokens: 1,
+			TotalTokens:      2,
+		},
+		Choices: []model.Choice{
+			{
+				Delta: model.Message{
+					Content:          content,
+					ReasoningContent: reasoningContent,
+					ToolCalls: []model.ToolCall{
+						{
+							ID: "id",
+						},
+					},
+				},
+				FinishReason: &finishReason,
+			},
+		},
+	}
+	type fields struct {
+		Model            string
+		FullText         strings.Builder
+		ReasoningContent strings.Builder
+		FinishReason     strings.Builder
+		ToolCalls        []model.ToolCall
+		Usage            model.Usage
+	}
+	tests := []struct {
+		name   string
+		fields fields
+		want   *model.Response
+	}{
+		{
+			name: "Accumulate",
+			fields: fields{
+				Model:            "gemini-pro",
+				FullText:         strings.Builder{},
+				ReasoningContent: strings.Builder{},
+				FinishReason:     strings.Builder{},
+				ToolCalls:        []model.ToolCall{},
+				Usage:            model.Usage{},
+			},
+			want: &model.Response{
+				Usage: &model.Usage{
+					PromptTokens:     1,
+					CompletionTokens: 1,
+					TotalTokens:      2,
+				},
+				Choices: []model.Choice{
+					{
+						Delta: model.Message{
+							Role:             model.RoleAssistant,
+							Content:          content,
+							ReasoningContent: reasoningContent,
+							ToolCalls: []model.ToolCall{
+								{
+									ID: "id",
+								},
+							},
+						},
+						FinishReason: &finishReason,
+					},
+				},
+				Done: true,
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			a := &Accumulator{
+				Model:            tt.fields.Model,
+				FullText:         tt.fields.FullText,
+				ReasoningContent: tt.fields.ReasoningContent,
+				FinishReason:     tt.fields.FinishReason,
+				ToolCalls:        tt.fields.ToolCalls,
+				Usage:            tt.fields.Usage,
+			}
+			a.Accumulate(m)
+			got := a.BuildResponse()
+			assert.Equal(t, got.Choices[0].Message.Content, content)
+			assert.Equal(t, got.Choices[0].Message.ReasoningContent, reasoningContent)
+			assert.Equal(t, got.Choices[0].FinishReason, &finishReason)
+			assert.Equal(t, got.Usage, m.Usage)
+		})
+	}
+}

model/gemini/gemini.go

@@ -52,10 +52,10 @@ type ChatStreamCompleteCallbackFunc func(
 	ctx context.Context,
 	chatRequest []*genai.Content,
 	generateConfig *genai.GenerateContentConfig,
-	chatResponse *genai.GenerateContentResponse,
+	chatResponse *model.Response,
 )
 
-// Model implements the model.Model interface for OpenAI API.
+// Model implements the model.Model interface for Gemini API.
 type Model struct {
 	client                     *genai.Client
 	name                       string
@@ -79,7 +79,7 @@ type Model struct {
 	maxInputTokensRatio    float64
 }
 
-// New creates a new OpenAI-like model.
+// New creates a new Gemini-like model.
 func New(ctx context.Context, name string, opts ...Option) (*Model, error) {
 	o := &options{
 		channelBufferSize: defaultChannelBufferSize,
@@ -230,23 +230,33 @@ func (m *Model) handleStreamingResponse(
 ) {
 	chatCompletion := m.client.Models.GenerateContentStream(
 		ctx, m.name, chatRequest, generateConfig)
+	acc := &Accumulator{}
 	for chunk := range chatCompletion {
 		response := m.buildResponse(chunk)
-		if response.IsPartial {
-			if m.chatChunkCallback != nil {
-				m.chatChunkCallback(ctx, chatRequest, generateConfig, chunk)
-			}
-		} else {
-			if m.chatStreamCompleteCallback != nil {
-				m.chatStreamCompleteCallback(ctx, chatRequest, generateConfig, chunk)
-			}
+		acc.Accumulate(response)
+		response.Object = model.ObjectTypeChatCompletionChunk
+		response.IsPartial = true
+		if m.chatChunkCallback != nil {
+			m.chatChunkCallback(ctx, chatRequest, generateConfig, chunk)
+		}
+		if m.chatChunkCallback != nil {
+			m.chatChunkCallback(ctx, chatRequest, generateConfig, chunk)
 		}
 		select {
 		case responseChan <- response:
 		case <-ctx.Done():
 			return
 		}
 	}
+	finalResponse := acc.BuildResponse()
+	if m.chatStreamCompleteCallback != nil {
+		m.chatStreamCompleteCallback(ctx, chatRequest, generateConfig, finalResponse)
+	}
+	select {
+	case responseChan <- finalResponse:
+	case <-ctx.Done():
+		return
+	}
 }
 
 // convertContentBlock builds a single assistant message from Gemini Candidate.
@@ -292,6 +302,8 @@ func (m *Model) convertContentBlock(candidates []*genai.Candidate) (model.Messag
 	}, finishReason
 }
 
+// buildResponse builds a partial streaming response for a chunk.
+// Returns nil if the chunk should be skipped.
 func (m *Model) buildResponse(chatCompletion *genai.GenerateContentResponse) *model.Response {
 	if chatCompletion == nil {
 		return &model.Response{}
@@ -300,7 +312,7 @@ func (m *Model) buildResponse(chatCompletion *genai.GenerateContentResponse) *mo
 		ID:        chatCompletion.ResponseID,
 		Created:   chatCompletion.CreateTime.Unix(),
 		Model:     chatCompletion.ModelVersion,
-		Timestamp: time.Now(),
+		Timestamp: chatCompletion.CreateTime,
 	}
 	message, finishReason := m.convertContentBlock(chatCompletion.Candidates)
 	response.Choices = []model.Choice{
@@ -313,9 +325,6 @@ func (m *Model) buildResponse(chatCompletion *genai.GenerateContentResponse) *mo
 	// Set finish reason.
 	if finishReason != "" {
 		response.Choices[0].FinishReason = &finishReason
-		response.Done = true
-	} else {
-		response.IsPartial = true
 	}
 	// Convert usage information.
 	response.Usage = m.completionUsageToModelUsage(chatCompletion.UsageMetadata)