replace `github.com/google/generative-ai-go` with `github.com/googleapis/go-genai` (#138)

mineo and Kujtim Hoxha created

* replace to github.com/googleapis/go-genai

* fix history logic

* small fixes

---------

Co-authored-by: Kujtim Hoxha <kujtimii.h@gmail.com>

Change summary

go.mod                          |  12 --
go.sum                          |  22 +----
internal/llm/provider/gemini.go | 139 +++++++++++++++++-----------------
3 files changed, 76 insertions(+), 97 deletions(-)

Detailed changes

go.mod 🔗

@@ -18,7 +18,6 @@ require (
 	github.com/charmbracelet/x/ansi v0.8.0
 	github.com/fsnotify/fsnotify v1.8.0
 	github.com/go-logfmt/logfmt v0.6.0
-	github.com/google/generative-ai-go v0.19.0
 	github.com/google/uuid v1.6.0
 	github.com/lrstanley/bubblezone v0.0.0-20250315020633-c249a3fe1231
 	github.com/mark3labs/mcp-go v0.17.0
@@ -32,16 +31,14 @@ require (
 	github.com/spf13/cobra v1.9.1
 	github.com/spf13/viper v1.20.0
 	github.com/stretchr/testify v1.10.0
-	google.golang.org/api v0.215.0
 )
 
 require (
 	cloud.google.com/go v0.116.0 // indirect
-	cloud.google.com/go/ai v0.8.0 // indirect
+	github.com/google/go-cmp v0.7.0 // indirect
+	github.com/gorilla/websocket v1.5.3 // indirect
 	cloud.google.com/go/auth v0.13.0 // indirect
-	cloud.google.com/go/auth/oauth2adapt v0.2.6 // indirect
 	cloud.google.com/go/compute/metadata v0.6.0 // indirect
-	cloud.google.com/go/longrunning v0.5.7 // indirect
 	github.com/Azure/azure-sdk-for-go/sdk/azcore v1.17.0 // indirect
 	github.com/Azure/azure-sdk-for-go/sdk/internal v1.10.0 // indirect
 	github.com/AzureAD/microsoft-authentication-library-for-go v1.2.2 // indirect
@@ -111,7 +108,6 @@ require (
 	github.com/yuin/goldmark v1.7.8 // indirect
 	github.com/yuin/goldmark-emoji v1.0.5 // indirect
 	go.opentelemetry.io/auto/sdk v1.1.0 // indirect
-	go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.54.0 // indirect
 	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.54.0 // indirect
 	go.opentelemetry.io/otel v1.35.0 // indirect
 	go.opentelemetry.io/otel/metric v1.35.0 // indirect
@@ -120,13 +116,11 @@ require (
 	golang.org/x/crypto v0.37.0 // indirect
 	golang.org/x/image v0.26.0 // indirect
 	golang.org/x/net v0.39.0 // indirect
-	golang.org/x/oauth2 v0.25.0 // indirect
 	golang.org/x/sync v0.13.0 // indirect
 	golang.org/x/sys v0.32.0 // indirect
 	golang.org/x/term v0.31.0 // indirect
 	golang.org/x/text v0.24.0 // indirect
-	golang.org/x/time v0.8.0 // indirect
-	google.golang.org/genproto/googleapis/api v0.0.0-20250106144421-5f5ef82da422 // indirect
+	google.golang.org/genai v1.3.0
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20250324211829-b45e905df463 // indirect
 	google.golang.org/grpc v1.71.0 // indirect
 	google.golang.org/protobuf v1.36.6 // indirect

go.sum 🔗

@@ -1,15 +1,9 @@
 cloud.google.com/go v0.116.0 h1:B3fRrSDkLRt5qSHWe40ERJvhvnQwdZiHu0bJOpldweE=
 cloud.google.com/go v0.116.0/go.mod h1:cEPSRWPzZEswwdr9BxE6ChEn01dWlTaF05LiC2Xs70U=
-cloud.google.com/go/ai v0.8.0 h1:rXUEz8Wp2OlrM8r1bfmpF2+VKqc1VJpafE3HgzRnD/w=
-cloud.google.com/go/ai v0.8.0/go.mod h1:t3Dfk4cM61sytiggo2UyGsDVW3RF1qGZaUKDrZFyqkE=
 cloud.google.com/go/auth v0.13.0 h1:8Fu8TZy167JkW8Tj3q7dIkr2v4cndv41ouecJx0PAHs=
 cloud.google.com/go/auth v0.13.0/go.mod h1:COOjD9gwfKNKz+IIduatIhYJQIc0mG3H102r/EMxX6Q=
-cloud.google.com/go/auth/oauth2adapt v0.2.6 h1:V6a6XDu2lTwPZWOawrAa9HUK+DB2zfJyTuciBG5hFkU=
-cloud.google.com/go/auth/oauth2adapt v0.2.6/go.mod h1:AlmsELtlEBnaNTL7jCj8VQFLy6mbZv0s4Q7NGBeQ5E8=
 cloud.google.com/go/compute/metadata v0.6.0 h1:A6hENjEsCDtC1k8byVsgwvVcioamEHvZ4j01OwKxG9I=
 cloud.google.com/go/compute/metadata v0.6.0/go.mod h1:FjyFAW1MW0C203CEOMDTu3Dk1FlqW3Rga40jzHL4hfg=
-cloud.google.com/go/longrunning v0.5.7 h1:WLbHekDbjK1fVFD3ibpFFVoyizlLRl73I7YKuAKilhU=
-cloud.google.com/go/longrunning v0.5.7/go.mod h1:8GClkudohy1Fxm3owmBGid8W0pSgodEMwEAztp38Xng=
 github.com/Azure/azure-sdk-for-go/sdk/azcore v1.17.0 h1:g0EZJwz7xkXQiZAI5xi9f3WWFYBlX1CPTrR+NDToRkQ=
 github.com/Azure/azure-sdk-for-go/sdk/azcore v1.17.0/go.mod h1:XCW7KnZet0Opnr7HccfUw1PLc4CjHqpcaxW8DHklNkQ=
 github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.7.0 h1:tfLQ34V6F7tVSwoTf/4lH5sE0o6eCJuNDTmH09nDpbc=
@@ -123,8 +117,6 @@ github.com/golang-jwt/jwt/v5 v5.2.2 h1:Rl4B7itRWVtYIHFrSNd7vhTiz9UpLdi6gZhZ3wEeD
 github.com/golang-jwt/jwt/v5 v5.2.2/go.mod h1:pqrtFR0X4osieyHYxtmOUWsAWrfe1Q5UVIyoH402zdk=
 github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek=
 github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps=
-github.com/google/generative-ai-go v0.19.0 h1:R71szggh8wHMCUlEMsW2A/3T+5LdEIkiaHSYgSpUgdg=
-github.com/google/generative-ai-go v0.19.0/go.mod h1:JYolL13VG7j79kM5BtHz4qwONHkeJQzOCkKXnpqtS/E=
 github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
 github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
 github.com/google/s2a-go v0.1.8 h1:zZDs9gcbt9ZPLV0ndSyQk6Kacx2g/X+SKYovpnz3SMM=
@@ -137,6 +129,8 @@ github.com/googleapis/gax-go/v2 v2.14.1 h1:hb0FFeiPaQskmvakKu5EbCbpntQn48jyHuvrk
 github.com/googleapis/gax-go/v2 v2.14.1/go.mod h1:Hb/NubMaVM88SrNkvl8X/o8XWwDJEPqouaLeN2IUxoA=
 github.com/gorilla/css v1.0.1 h1:ntNaBIghp6JmvWnxbZKANoLyuXTPZ4cAMlo6RyhlbO8=
 github.com/gorilla/css v1.0.1/go.mod h1:BvnYkspnSzMmwRK+b8/xgNPLiIuNZr6vbZBTPQ2A3b0=
+github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
+github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/hexops/gotextdiff v1.0.3 h1:gitA9+qJrrTCsiCl7+kh75nPqQt1cx4ZkudSTLoUqJM=
 github.com/hexops/gotextdiff v1.0.3/go.mod h1:pSWU5MAI3yDq+fZBTazCSJysOMbxWL1BSow5/V2vxeg=
 github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
@@ -254,8 +248,6 @@ github.com/yuin/goldmark-emoji v1.0.5 h1:EMVWyCGPlXJfUXBXpuMu+ii3TIaxbVBnEX9uaDC
 github.com/yuin/goldmark-emoji v1.0.5/go.mod h1:tTkZEbwu5wkPmgTcitqddVxY9osFZiavD+r4AzQrh1U=
 go.opentelemetry.io/auto/sdk v1.1.0 h1:cH53jehLUN6UFLY71z+NDOiNJqDdPRaXzTel0sJySYA=
 go.opentelemetry.io/auto/sdk v1.1.0/go.mod h1:3wSPjt5PWp2RhlCcmmOial7AvC4DQqZb7a7wCow3W8A=
-go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.54.0 h1:r6I7RJCN86bpD/FQwedZ0vSixDpwuWREjW9oRMsmqDc=
-go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.54.0/go.mod h1:B9yO6b04uB80CzjedvewuqDhxJxi11s7/GtiGa8bAjI=
 go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.54.0 h1:TT4fX+nBOA/+LUkobKGW1ydGcn+G3vRw9+g5HwCphpk=
 go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.54.0/go.mod h1:L7UH0GbB0p47T4Rri3uHjbpCFYrVrwc1I25QhNPiGK8=
 go.opentelemetry.io/otel v1.35.0 h1:xKWKPxrxB6OtMCbmMY021CqC45J+3Onta9MqjhnusiQ=
@@ -295,8 +287,6 @@ golang.org/x/net v0.24.0/go.mod h1:2Q7sJY5mzlzWjKtYUEXSlBWCdyaioyXzRB2RtU8KVE8=
 golang.org/x/net v0.25.0/go.mod h1:JkAGAh7GEvH74S6FOH42FLoXpXbE/aqXSrIQjXgsiwM=
 golang.org/x/net v0.39.0 h1:ZCu7HMWDxpXpaiKdhzIfaltL9Lp31x/3fCP11bc6/fY=
 golang.org/x/net v0.39.0/go.mod h1:X7NRbYVEA+ewNkCNyJ513WmMdQ3BineSwVtN2zD/d+E=
-golang.org/x/oauth2 v0.25.0 h1:CY4y7XT9v0cRI9oupztF8AgiIu99L/ksR/Xp/6jrZ70=
-golang.org/x/oauth2 v0.25.0/go.mod h1:XYTD2NtWslqkgxebSiOHnXEap4TF09sJSc7H1sXbhtI=
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
@@ -337,17 +327,13 @@ golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
 golang.org/x/text v0.15.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
 golang.org/x/text v0.24.0 h1:dd5Bzh4yt5KYA8f9CJHCP4FB4D51c2c6JvN37xJJkJ0=
 golang.org/x/text v0.24.0/go.mod h1:L8rBsPeo2pSS+xqN0d5u2ikmjtmoJbDBT1b7nHvFCdU=
-golang.org/x/time v0.8.0 h1:9i3RxcPv3PZnitoVGMPDKZSq1xW1gK1Xy3ArNOGZfEg=
-golang.org/x/time v0.8.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
 golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc=
 golang.org/x/tools v0.6.0/go.mod h1:Xwgl3UAJ/d3gWutnCtw505GrjyAbvKui8lOU390QaIU=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
-google.golang.org/api v0.215.0 h1:jdYF4qnyczlEz2ReWIsosNLDuzXyvFHJtI5gcr0J7t0=
-google.golang.org/api v0.215.0/go.mod h1:fta3CVtuJYOEdugLNWm6WodzOS8KdFckABwN4I40hzY=
-google.golang.org/genproto/googleapis/api v0.0.0-20250106144421-5f5ef82da422 h1:GVIKPyP/kLIyVOgOnTwFOrvQaQUzOzGMCxgFUOEmm24=
-google.golang.org/genproto/googleapis/api v0.0.0-20250106144421-5f5ef82da422/go.mod h1:b6h1vNKhxaSoEI+5jc3PJUCustfli/mRab7295pY7rw=
+google.golang.org/genai v1.3.0 h1:tXhPJF30skOjnnDY7ZnjK3q7IKy4PuAlEA0fk7uEaEI=
+google.golang.org/genai v1.3.0/go.mod h1:TyfOKRz/QyCaj6f/ZDt505x+YreXnY40l2I6k8TvgqY=
 google.golang.org/genproto/googleapis/rpc v0.0.0-20250324211829-b45e905df463 h1:e0AIkUUhxyBKh6ssZNrAMeqhA7RKUj42346d1y02i2g=
 google.golang.org/genproto/googleapis/rpc v0.0.0-20250324211829-b45e905df463/go.mod h1:qQ0YXyHHx3XkvlzUtpXDkS29lDSafHMZBAZDc03LQ3A=
 google.golang.org/grpc v1.71.0 h1:kF77BGdPTQ4/JZWMlb9VpJ5pa25aqvVqogsxNHHdeBg=

internal/llm/provider/gemini.go 🔗

@@ -9,14 +9,12 @@ import (
 	"strings"
 	"time"
 
-	"github.com/google/generative-ai-go/genai"
 	"github.com/google/uuid"
 	"github.com/opencode-ai/opencode/internal/config"
 	"github.com/opencode-ai/opencode/internal/llm/tools"
 	"github.com/opencode-ai/opencode/internal/logging"
 	"github.com/opencode-ai/opencode/internal/message"
-	"google.golang.org/api/iterator"
-	"google.golang.org/api/option"
+	"google.golang.org/genai"
 )
 
 type geminiOptions struct {
@@ -39,7 +37,7 @@ func newGeminiClient(opts providerClientOptions) GeminiClient {
 		o(&geminiOpts)
 	}
 
-	client, err := genai.NewClient(context.Background(), option.WithAPIKey(opts.apiKey))
+	client, err := genai.NewClient(context.Background(), &genai.ClientConfig{APIKey: opts.apiKey, Backend: genai.BackendGeminiAPI})
 	if err != nil {
 		logging.Error("Failed to create Gemini client", "error", err)
 		return nil
@@ -57,11 +55,14 @@ func (g *geminiClient) convertMessages(messages []message.Message) []*genai.Cont
 	for _, msg := range messages {
 		switch msg.Role {
 		case message.User:
-			var parts []genai.Part
-			parts = append(parts, genai.Text(msg.Content().String()))
+			var parts []*genai.Part
+			parts = append(parts, &genai.Part{Text: msg.Content().String()})
 			for _, binaryContent := range msg.BinaryContent() {
 				imageFormat := strings.Split(binaryContent.MIMEType, "/")
-				parts = append(parts, genai.ImageData(imageFormat[1], binaryContent.Data))
+				parts = append(parts, &genai.Part{InlineData: &genai.Blob{
+					MIMEType: imageFormat[1],
+					Data:     binaryContent.Data,
+				}})
 			}
 			history = append(history, &genai.Content{
 				Parts: parts,
@@ -70,19 +71,21 @@ func (g *geminiClient) convertMessages(messages []message.Message) []*genai.Cont
 		case message.Assistant:
 			content := &genai.Content{
 				Role:  "model",
-				Parts: []genai.Part{},
+				Parts: []*genai.Part{},
 			}
 
 			if msg.Content().String() != "" {
-				content.Parts = append(content.Parts, genai.Text(msg.Content().String()))
+				content.Parts = append(content.Parts, &genai.Part{Text: msg.Content().String()})
 			}
 
 			if len(msg.ToolCalls()) > 0 {
 				for _, call := range msg.ToolCalls() {
 					args, _ := parseJsonToMap(call.Input)
-					content.Parts = append(content.Parts, genai.FunctionCall{
-						Name: call.Name,
-						Args: args,
+					content.Parts = append(content.Parts, &genai.Part{
+						FunctionCall: &genai.FunctionCall{
+							Name: call.Name,
+							Args: args,
+						},
 					})
 				}
 			}
@@ -110,10 +113,14 @@ func (g *geminiClient) convertMessages(messages []message.Message) []*genai.Cont
 				}
 
 				history = append(history, &genai.Content{
-					Parts: []genai.Part{genai.FunctionResponse{
-						Name:     toolCall.Name,
-						Response: response,
-					}},
+					Parts: []*genai.Part{
+						{
+							FunctionResponse: &genai.FunctionResponse{
+								Name:     toolCall.Name,
+								Response: response,
+							},
+						},
+					},
 					Role: "function",
 				})
 			}
@@ -157,18 +164,6 @@ func (g *geminiClient) finishReason(reason genai.FinishReason) message.FinishRea
 }
 
 func (g *geminiClient) send(ctx context.Context, messages []message.Message, tools []tools.BaseTool) (*ProviderResponse, error) {
-	model := g.client.GenerativeModel(g.providerOptions.model.APIModel)
-	model.SetMaxOutputTokens(int32(g.providerOptions.maxTokens))
-	model.SystemInstruction = &genai.Content{
-		Parts: []genai.Part{
-			genai.Text(g.providerOptions.systemMessage),
-		},
-	}
-	// Convert tools
-	if len(tools) > 0 {
-		model.Tools = g.convertTools(tools)
-	}
-
 	// Convert messages
 	geminiMessages := g.convertMessages(messages)
 
@@ -178,16 +173,26 @@ func (g *geminiClient) send(ctx context.Context, messages []message.Message, too
 		logging.Debug("Prepared messages", "messages", string(jsonData))
 	}
 
+	history := geminiMessages[:len(geminiMessages)-1] // All but last message
+	lastMsg := geminiMessages[len(geminiMessages)-1]
+	chat, _ := g.client.Chats.Create(ctx, g.providerOptions.model.APIModel, &genai.GenerateContentConfig{
+		MaxOutputTokens: int32(g.providerOptions.maxTokens),
+		SystemInstruction: &genai.Content{
+			Parts: []*genai.Part{{Text: g.providerOptions.systemMessage}},
+		},
+		Tools: g.convertTools(tools),
+	}, history)
+
 	attempts := 0
 	for {
 		attempts++
 		var toolCalls []message.ToolCall
-		chat := model.StartChat()
-		chat.History = geminiMessages[:len(geminiMessages)-1] // All but last message
-
-		lastMsg := geminiMessages[len(geminiMessages)-1]
 
-		resp, err := chat.SendMessage(ctx, lastMsg.Parts...)
+		var lastMsgParts []genai.Part
+		for _, part := range lastMsg.Parts {
+			lastMsgParts = append(lastMsgParts, *part)
+		}
+		resp, err := chat.SendMessage(ctx, lastMsgParts...)
 		// If there is an error we are going to see if we can retry the call
 		if err != nil {
 			retry, after, retryErr := g.shouldRetry(attempts, err)
@@ -210,15 +215,15 @@ func (g *geminiClient) send(ctx context.Context, messages []message.Message, too
 
 		if len(resp.Candidates) > 0 && resp.Candidates[0].Content != nil {
 			for _, part := range resp.Candidates[0].Content.Parts {
-				switch p := part.(type) {
-				case genai.Text:
-					content = string(p)
-				case genai.FunctionCall:
+				switch {
+				case part.Text != "":
+					content = string(part.Text)
+				case part.FunctionCall != nil:
 					id := "call_" + uuid.New().String()
-					args, _ := json.Marshal(p.Args)
+					args, _ := json.Marshal(part.FunctionCall.Args)
 					toolCalls = append(toolCalls, message.ToolCall{
 						ID:       id,
-						Name:     p.Name,
+						Name:     part.FunctionCall.Name,
 						Input:    string(args),
 						Type:     "function",
 						Finished: true,
@@ -244,18 +249,6 @@ func (g *geminiClient) send(ctx context.Context, messages []message.Message, too
 }
 
 func (g *geminiClient) stream(ctx context.Context, messages []message.Message, tools []tools.BaseTool) <-chan ProviderEvent {
-	model := g.client.GenerativeModel(g.providerOptions.model.APIModel)
-	model.SetMaxOutputTokens(int32(g.providerOptions.maxTokens))
-	model.SystemInstruction = &genai.Content{
-		Parts: []genai.Part{
-			genai.Text(g.providerOptions.systemMessage),
-		},
-	}
-	// Convert tools
-	if len(tools) > 0 {
-		model.Tools = g.convertTools(tools)
-	}
-
 	// Convert messages
 	geminiMessages := g.convertMessages(messages)
 
@@ -265,6 +258,16 @@ func (g *geminiClient) stream(ctx context.Context, messages []message.Message, t
 		logging.Debug("Prepared messages", "messages", string(jsonData))
 	}
 
+	history := geminiMessages[:len(geminiMessages)-1] // All but last message
+	lastMsg := geminiMessages[len(geminiMessages)-1]
+	chat, _ := g.client.Chats.Create(ctx, g.providerOptions.model.APIModel, &genai.GenerateContentConfig{
+		MaxOutputTokens: int32(g.providerOptions.maxTokens),
+		SystemInstruction: &genai.Content{
+			Parts: []*genai.Part{{Text: g.providerOptions.systemMessage}},
+		},
+		Tools: g.convertTools(tools),
+	}, history)
+
 	attempts := 0
 	eventChan := make(chan ProviderEvent)
 
@@ -273,11 +276,6 @@ func (g *geminiClient) stream(ctx context.Context, messages []message.Message, t
 
 		for {
 			attempts++
-			chat := model.StartChat()
-			chat.History = geminiMessages[:len(geminiMessages)-1]
-			lastMsg := geminiMessages[len(geminiMessages)-1]
-
-			iter := chat.SendMessageStream(ctx, lastMsg.Parts...)
 
 			currentContent := ""
 			toolCalls := []message.ToolCall{}
@@ -285,11 +283,12 @@ func (g *geminiClient) stream(ctx context.Context, messages []message.Message, t
 
 			eventChan <- ProviderEvent{Type: EventContentStart}
 
-			for {
-				resp, err := iter.Next()
-				if err == iterator.Done {
-					break
-				}
+			var lastMsgParts []genai.Part
+
+			for _, part := range lastMsg.Parts {
+				lastMsgParts = append(lastMsgParts, *part)
+			}
+			for resp, err := range chat.SendMessageStream(ctx, lastMsgParts...) {
 				if err != nil {
 					retry, after, retryErr := g.shouldRetry(attempts, err)
 					if retryErr != nil {
@@ -318,9 +317,9 @@ func (g *geminiClient) stream(ctx context.Context, messages []message.Message, t
 
 				if len(resp.Candidates) > 0 && resp.Candidates[0].Content != nil {
 					for _, part := range resp.Candidates[0].Content.Parts {
-						switch p := part.(type) {
-						case genai.Text:
-							delta := string(p)
+						switch {
+						case part.Text != "":
+							delta := string(part.Text)
 							if delta != "" {
 								eventChan <- ProviderEvent{
 									Type:    EventContentDelta,
@@ -328,12 +327,12 @@ func (g *geminiClient) stream(ctx context.Context, messages []message.Message, t
 								}
 								currentContent += delta
 							}
-						case genai.FunctionCall:
+						case part.FunctionCall != nil:
 							id := "call_" + uuid.New().String()
-							args, _ := json.Marshal(p.Args)
+							args, _ := json.Marshal(part.FunctionCall.Args)
 							newCall := message.ToolCall{
 								ID:       id,
-								Name:     p.Name,
+								Name:     part.FunctionCall.Name,
 								Input:    string(args),
 								Type:     "function",
 								Finished: true,
@@ -421,12 +420,12 @@ func (g *geminiClient) toolCalls(resp *genai.GenerateContentResponse) []message.
 
 	if len(resp.Candidates) > 0 && resp.Candidates[0].Content != nil {
 		for _, part := range resp.Candidates[0].Content.Parts {
-			if funcCall, ok := part.(genai.FunctionCall); ok {
+			if part.FunctionCall != nil {
 				id := "call_" + uuid.New().String()
-				args, _ := json.Marshal(funcCall.Args)
+				args, _ := json.Marshal(part.FunctionCall.Args)
 				toolCalls = append(toolCalls, message.ToolCall{
 					ID:    id,
-					Name:  funcCall.Name,
+					Name:  part.FunctionCall.Name,
 					Input: string(args),
 					Type:  "function",
 				})