Skip to content

Commit 3e76cf1

Browse files
appleboyappleboy
authored
refactor: migrate to new genai API and update Gemini provider (#238)
* refactor: migrate to new genai API and update Gemini provider implementation - Replace deprecated google/generative-ai-go/genai package with google.golang.org/genai and update imports accordingly - Refactor Gemini provider: update API usage to reflect new genai client structure and methods - Update completion and function call logic to use explicit content/message structs rather than genai.Text - Add internal error handling for invalid function call responses - Improve usage metrics extraction with additional nil checks - Remove unused imports and update client initialization logic to match new genai API - Add new indirect dependencies in go.mod (google.golang.org/genai, github.com/google/go-cmp, github.com/gorilla/websocket) - Remove obsolete and unused dependencies in go.mod Signed-off-by: appleboy <[email protected]> * feat: add debug dumping of Gemini API responses using godump - Add github.com/yassinebenaid/godump v0.11.1 as a new dependency - Import godump in the gemini provider code - Dump Gemini API response candidates to console when debug mode is enabled Signed-off-by: appleboy <[email protected]> * refactor: refactor error handling and simplify prefix assignment - Remove the ErrInvalidFunctionCall error declaration - Refactor error handling in GetSummaryPrefix to provide more granular error messages for missing candidates, content, function calls, or prefix values - Store the extracted prefix in a variable to simplify assignment to the response object Signed-off-by: appleboy <[email protected]> --------- Signed-off-by: appleboy <[email protected]>
1 parent 76cda3c commit 3e76cf1

File tree

4 files changed

+109
-91
lines changed

4 files changed

+109
-91
lines changed

‎go.mod

Lines changed: 4 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -9,24 +9,21 @@ require (
99
github.com/charmbracelet/bubbletea v1.3.4
1010
github.com/erikgeiser/promptkit v0.9.0
1111
github.com/fatih/color v1.18.0
12-
github.com/google/generative-ai-go v0.19.0
1312
github.com/joho/godotenv v1.5.1
1413
github.com/liushuangls/go-anthropic/v2 v2.14.1
1514
github.com/rodaine/table v1.3.0
1615
github.com/sashabaranov/go-openai v1.38.0
1716
github.com/spf13/cobra v1.9.1
1817
github.com/spf13/viper v1.19.0
18+
github.com/yassinebenaid/godump v0.11.1
1919
golang.org/x/net v0.39.0
20-
google.golang.org/api v0.223.0
20+
google.golang.org/genai v1.3.0
2121
)
2222

2323
require (
2424
cloud.google.com/go v0.118.3 // indirect
25-
cloud.google.com/go/ai v0.10.0 // indirect
2625
cloud.google.com/go/auth v0.15.0 // indirect
27-
cloud.google.com/go/auth/oauth2adapt v0.2.7 // indirect
2826
cloud.google.com/go/compute/metadata v0.6.0 // indirect
29-
cloud.google.com/go/longrunning v0.6.4 // indirect
3027
github.com/atotto/clipboard v0.1.4 // indirect
3128
github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect
3229
github.com/charmbracelet/lipgloss v1.0.0 // indirect
@@ -37,10 +34,11 @@ require (
3734
github.com/fsnotify/fsnotify v1.8.0 // indirect
3835
github.com/go-logr/logr v1.4.2 // indirect
3936
github.com/go-logr/stdr v1.2.2 // indirect
37+
github.com/google/go-cmp v0.7.0 // indirect
4038
github.com/google/s2a-go v0.1.9 // indirect
41-
github.com/google/uuid v1.6.0 // indirect
4239
github.com/googleapis/enterprise-certificate-proxy v0.3.5 // indirect
4340
github.com/googleapis/gax-go/v2 v2.14.1 // indirect
41+
github.com/gorilla/websocket v1.5.3 // indirect
4442
github.com/hashicorp/hcl v1.0.0 // indirect
4543
github.com/inconshreveable/mousetrap v1.1.0 // indirect
4644
github.com/lucasb-eyer/go-colorful v1.2.0 // indirect
@@ -64,20 +62,16 @@ require (
6462
github.com/spf13/pflag v1.0.6 // indirect
6563
github.com/subosito/gotenv v1.6.0 // indirect
6664
go.opentelemetry.io/auto/sdk v1.1.0 // indirect
67-
go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.60.0 // indirect
6865
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.60.0 // indirect
6966
go.opentelemetry.io/otel v1.35.0 // indirect
7067
go.opentelemetry.io/otel/metric v1.35.0 // indirect
7168
go.opentelemetry.io/otel/trace v1.35.0 // indirect
7269
go.uber.org/multierr v1.11.0 // indirect
7370
golang.org/x/crypto v0.37.0 // indirect
7471
golang.org/x/exp v0.0.0-20250305212735-054e65f0b394 // indirect
75-
golang.org/x/oauth2 v0.28.0 // indirect
7672
golang.org/x/sync v0.13.0 // indirect
7773
golang.org/x/sys v0.32.0 // indirect
7874
golang.org/x/text v0.24.0 // indirect
79-
golang.org/x/time v0.11.0 // indirect
80-
google.golang.org/genproto/googleapis/api v0.0.0-20250303144028-a0af3efb3deb // indirect
8175
google.golang.org/genproto/googleapis/rpc v0.0.0-20250303144028-a0af3efb3deb // indirect
8276
google.golang.org/grpc v1.71.0 // indirect
8377
google.golang.org/protobuf v1.36.5 // indirect

‎go.sum

Lines changed: 6 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,9 @@
11
cloud.google.com/go v0.118.3 h1:jsypSnrE/w4mJysioGdMBg4MiW/hHx/sArFpaBWHdME=
22
cloud.google.com/go v0.118.3/go.mod h1:Lhs3YLnBlwJ4KA6nuObNMZ/fCbOQBPuWKPoE0Wa/9Vc=
3-
cloud.google.com/go/ai v0.10.0 h1:hwj6CI6sMKubXodoJJGTy/c2T1RbbLGM6TL3QoAvzU8=
4-
cloud.google.com/go/ai v0.10.0/go.mod h1:kvnt2KeHqX8+41PVeMRBETDyQAp/RFvBWGdx/aGjNMo=
53
cloud.google.com/go/auth v0.15.0 h1:Ly0u4aA5vG/fsSsxu98qCQBemXtAtJf+95z9HK+cxps=
64
cloud.google.com/go/auth v0.15.0/go.mod h1:WJDGqZ1o9E9wKIL+IwStfyn/+s59zl4Bi+1KQNVXLZ8=
7-
cloud.google.com/go/auth/oauth2adapt v0.2.7 h1:/Lc7xODdqcEw8IrZ9SvwnlLX6j9FHQM74z6cBk9Rw6M=
8-
cloud.google.com/go/auth/oauth2adapt v0.2.7/go.mod h1:NTbTTzfvPl1Y3V1nPpOgl2w6d/FjO7NNUQaWSox6ZMc=
95
cloud.google.com/go/compute/metadata v0.6.0 h1:A6hENjEsCDtC1k8byVsgwvVcioamEHvZ4j01OwKxG9I=
106
cloud.google.com/go/compute/metadata v0.6.0/go.mod h1:FjyFAW1MW0C203CEOMDTu3Dk1FlqW3Rga40jzHL4hfg=
11-
cloud.google.com/go/longrunning v0.6.4 h1:3tyw9rO3E2XVXzSApn1gyEEnH2K9SynNQjMlBi3uHLg=
12-
cloud.google.com/go/longrunning v0.6.4/go.mod h1:ttZpLCe6e7EXvn9OxpBRx7kZEB0efv8yBO6YnVMfhJs=
137
github.com/MakeNowJust/heredoc v1.0.0 h1:cXCdzVdstXyiTqTvfqk9SDHpKNjxuom+DOlyEeQ4pzQ=
148
github.com/MakeNowJust/heredoc v1.0.0/go.mod h1:mG5amYoWBHf8vpLOuehzbGGw0EHxpZZ6lCpQ4fNJ8LE=
159
github.com/appleboy/com v0.3.0 h1:omze/tJPyi2YVH+m23GSrCGt90A+4vQNpEYBW+GuSr4=
@@ -54,8 +48,6 @@ github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
5448
github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
5549
github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek=
5650
github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps=
57-
github.com/google/generative-ai-go v0.19.0 h1:R71szggh8wHMCUlEMsW2A/3T+5LdEIkiaHSYgSpUgdg=
58-
github.com/google/generative-ai-go v0.19.0/go.mod h1:JYolL13VG7j79kM5BtHz4qwONHkeJQzOCkKXnpqtS/E=
5951
github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
6052
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
6153
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
@@ -67,6 +59,8 @@ github.com/googleapis/enterprise-certificate-proxy v0.3.5 h1:VgzTY2jogw3xt39CusE
6759
github.com/googleapis/enterprise-certificate-proxy v0.3.5/go.mod h1:MkHOF77EYAE7qfSuSS9PU6g4Nt4e11cnsDUowfwewLA=
6860
github.com/googleapis/gax-go/v2 v2.14.1 h1:hb0FFeiPaQskmvakKu5EbCbpntQn48jyHuvrkurSS/Q=
6961
github.com/googleapis/gax-go/v2 v2.14.1/go.mod h1:Hb/NubMaVM88SrNkvl8X/o8XWwDJEPqouaLeN2IUxoA=
62+
github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
63+
github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
7064
github.com/hashicorp/hcl v1.0.0 h1:0Anlzjpi4vEasTeNFn2mLJgTSwt0+6sfsiTG8qcWGx4=
7165
github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T20WEQ=
7266
github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
@@ -146,10 +140,10 @@ github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOf
146140
github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
147141
github.com/subosito/gotenv v1.6.0 h1:9NlTDc1FTs4qu0DDq7AEtTPNw6SVm7uBMsUCUjABIf8=
148142
github.com/subosito/gotenv v1.6.0/go.mod h1:Dk4QP5c2W3ibzajGcXpNraDfq2IrhjMIvMSWPKKo0FU=
143+
github.com/yassinebenaid/godump v0.11.1 h1:SPujx/XaYqGDfmNh7JI3dOyCUVrG0bG2duhO3Eh2EhI=
144+
github.com/yassinebenaid/godump v0.11.1/go.mod h1:dc/0w8wmg6kVIvNGAzbKH1Oa54dXQx8SNKh4dPRyW44=
149145
go.opentelemetry.io/auto/sdk v1.1.0 h1:cH53jehLUN6UFLY71z+NDOiNJqDdPRaXzTel0sJySYA=
150146
go.opentelemetry.io/auto/sdk v1.1.0/go.mod h1:3wSPjt5PWp2RhlCcmmOial7AvC4DQqZb7a7wCow3W8A=
151-
go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.60.0 h1:x7wzEgXfnzJcHDwStJT+mxOz4etr2EcexjqhBvmoakw=
152-
go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.60.0/go.mod h1:rg+RlpR5dKwaS95IyyZqj5Wd4E13lk/msnTS0Xl9lJM=
153147
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.60.0 h1:sbiXRNDSWJOTobXh5HyQKjq6wUC5tNybqjIqDpAY4CU=
154148
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.60.0/go.mod h1:69uWxva0WgAA/4bu2Yy70SLDBwZXuQ6PbBpbsa5iZrQ=
155149
go.opentelemetry.io/otel v1.35.0 h1:xKWKPxrxB6OtMCbmMY021CqC45J+3Onta9MqjhnusiQ=
@@ -170,8 +164,6 @@ golang.org/x/exp v0.0.0-20250305212735-054e65f0b394 h1:nDVHiLt8aIbd/VzvPWN6kSOPE
170164
golang.org/x/exp v0.0.0-20250305212735-054e65f0b394/go.mod h1:sIifuuw/Yco/y6yb6+bDNfyeQ/MdPUy/hKEMYQV17cM=
171165
golang.org/x/net v0.39.0 h1:ZCu7HMWDxpXpaiKdhzIfaltL9Lp31x/3fCP11bc6/fY=
172166
golang.org/x/net v0.39.0/go.mod h1:X7NRbYVEA+ewNkCNyJ513WmMdQ3BineSwVtN2zD/d+E=
173-
golang.org/x/oauth2 v0.28.0 h1:CrgCKl8PPAVtLnU3c+EDw6x11699EWlsDeWNWKdIOkc=
174-
golang.org/x/oauth2 v0.28.0/go.mod h1:onh5ek6nERTohokkhCD/y2cV4Do3fxFHFuAejCkRWT8=
175167
golang.org/x/sync v0.13.0 h1:AauUjRAJ9OSnvULf/ARrrVywoJDy0YS2AwQ98I37610=
176168
golang.org/x/sync v0.13.0/go.mod h1:1dzgHSNfp02xaA81J2MS99Qcpr2w7fw1gpm99rleRqA=
177169
golang.org/x/sys v0.0.0-20210809222454-d867a43fc93e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
@@ -180,12 +172,8 @@ golang.org/x/sys v0.32.0 h1:s77OFDvIQeibCmezSnk/q6iAfkdiQaJi4VzroCFrN20=
180172
golang.org/x/sys v0.32.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
181173
golang.org/x/text v0.24.0 h1:dd5Bzh4yt5KYA8f9CJHCP4FB4D51c2c6JvN37xJJkJ0=
182174
golang.org/x/text v0.24.0/go.mod h1:L8rBsPeo2pSS+xqN0d5u2ikmjtmoJbDBT1b7nHvFCdU=
183-
golang.org/x/time v0.11.0 h1:/bpjEDfN9tkoN/ryeYHnv5hcMlc8ncjMcM4XBk5NWV0=
184-
golang.org/x/time v0.11.0/go.mod h1:CDIdPxbZBQxdj6cxyCIdrNogrJKMJ7pr37NYpMcMDSg=
185-
google.golang.org/api v0.223.0 h1:JUTaWEriXmEy5AhvdMgksGGPEFsYfUKaPEYXd4c3Wvc=
186-
google.golang.org/api v0.223.0/go.mod h1:C+RS7Z+dDwds2b+zoAk5hN/eSfsiCn0UDrYof/M4d2M=
187-
google.golang.org/genproto/googleapis/api v0.0.0-20250303144028-a0af3efb3deb h1:p31xT4yrYrSM/G4Sn2+TNUkVhFCbG9y8itM2S6Th950=
188-
google.golang.org/genproto/googleapis/api v0.0.0-20250303144028-a0af3efb3deb/go.mod h1:jbe3Bkdp+Dh2IrslsFCklNhweNTBgSYanP1UXhJDhKg=
175+
google.golang.org/genai v1.3.0 h1:tXhPJF30skOjnnDY7ZnjK3q7IKy4PuAlEA0fk7uEaEI=
176+
google.golang.org/genai v1.3.0/go.mod h1:TyfOKRz/QyCaj6f/ZDt505x+YreXnY40l2I6k8TvgqY=
189177
google.golang.org/genproto/googleapis/rpc v0.0.0-20250303144028-a0af3efb3deb h1:TLPQVbx1GJ8VKZxz52VAxl1EBgKXXbTiU9Fc5fZeLn4=
190178
google.golang.org/genproto/googleapis/rpc v0.0.0-20250303144028-a0af3efb3deb/go.mod h1:LuRYeWDFV6WOn90g357N17oMCaxpgCnbi/44qJvDn2I=
191179
google.golang.org/grpc v1.71.0 h1:kF77BGdPTQ4/JZWMlb9VpJ5pa25aqvVqogsxNHHdeBg=

‎provider/gemini/func.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
package gemini
22

3-
import "github.com/google/generative-ai-go/genai"
3+
import "google.golang.org/genai"
44

55
var summaryPrefixFunc = &genai.Tool{
66
FunctionDeclarations: []*genai.FunctionDeclaration{{

‎provider/gemini/gemini.go

Lines changed: 98 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -2,22 +2,21 @@ package gemini
22

33
import (
44
"context"
5-
"fmt"
5+
"errors"
66
"net/http"
7-
"strings"
87

98
"github.com/appleboy/CodeGPT/core"
109
"github.com/appleboy/CodeGPT/core/transport"
1110
"github.com/appleboy/CodeGPT/version"
12-
"github.com/appleboy/com/convert"
1311

14-
"github.com/google/generative-ai-go/genai"
12+
"github.com/appleboy/com/convert"
1513
"github.com/sashabaranov/go-openai"
16-
"google.golang.org/api/option"
14+
"github.com/yassinebenaid/godump"
15+
"google.golang.org/genai"
1716
)
1817

1918
type Client struct {
20-
client *genai.GenerativeModel
19+
client *genai.Client
2120
model string
2221
maxTokens int32
2322
temperature float32
@@ -27,78 +26,115 @@ type Client struct {
2726

2827
// Completion is a method on the Client struct that takes a context.Context and a string argument
2928
func (c *Client) Completion(ctx context.Context, content string) (*core.Response, error) {
30-
resp, err := c.client.GenerateContent(ctx, genai.Text(content))
31-
if err != nil {
32-
return nil, err
29+
cfg := &genai.GenerateContentConfig{
30+
TopP: convert.ToPtr(c.topP),
31+
Temperature: convert.ToPtr(c.temperature),
32+
MaxOutputTokens: c.maxTokens,
3333
}
34-
35-
var ret string
36-
37-
for _, cand := range resp.Candidates {
38-
for _, part := range cand.Content.Parts {
39-
ret += fmt.Sprintf("%v", part)
40-
}
34+
data := []*genai.Content{
35+
{
36+
Role: "user",
37+
Parts: []*genai.Part{
38+
{
39+
Text: content,
40+
},
41+
},
42+
},
4143
}
4244

43-
usage := core.Usage{
44-
PromptTokens: int(resp.UsageMetadata.PromptTokenCount),
45-
CompletionTokens: int(resp.UsageMetadata.CandidatesTokenCount),
46-
TotalTokens: int(resp.UsageMetadata.TotalTokenCount),
45+
resp, err := c.client.Models.GenerateContent(ctx, c.model, data, cfg)
46+
if err != nil {
47+
return nil, err
4748
}
4849

49-
if resp.UsageMetadata.CachedContentTokenCount > 0 {
50-
usage.PromptTokensDetails = &openai.PromptTokensDetails{
51-
CachedTokens: int(resp.UsageMetadata.CachedContentTokenCount),
50+
usage := core.Usage{}
51+
if resp.UsageMetadata != nil {
52+
usage.PromptTokens = int(resp.UsageMetadata.PromptTokenCount)
53+
usage.CompletionTokens = int(resp.UsageMetadata.CandidatesTokenCount)
54+
usage.TotalTokens = int(resp.UsageMetadata.TotalTokenCount)
55+
if resp.UsageMetadata.CachedContentTokenCount > 0 {
56+
usage.PromptTokensDetails = &openai.PromptTokensDetails{
57+
CachedTokens: int(resp.UsageMetadata.CachedContentTokenCount),
58+
}
5259
}
5360
}
5461

5562
return &core.Response{
56-
Content: ret,
63+
Content: resp.Text(),
5764
Usage: usage,
5865
}, nil
5966
}
6067

6168
// GetSummaryPrefix is an API call to get a summary prefix using function call.
6269
func (c *Client) GetSummaryPrefix(ctx context.Context, content string) (*core.Response, error) {
63-
c.client.Tools = []*genai.Tool{summaryPrefixFunc}
64-
65-
// Start new chat session.
66-
session := c.client.StartChat()
70+
cfg := &genai.GenerateContentConfig{
71+
MaxOutputTokens: c.maxTokens,
72+
TopP: convert.ToPtr(c.topP),
73+
Temperature: convert.ToPtr(c.temperature),
74+
Tools: []*genai.Tool{summaryPrefixFunc},
75+
ToolConfig: &genai.ToolConfig{
76+
FunctionCallingConfig: &genai.FunctionCallingConfig{
77+
Mode: genai.FunctionCallingConfigModeAny,
78+
AllowedFunctionNames: []string{
79+
"get_summary_prefix",
80+
},
81+
},
82+
},
83+
}
84+
data := []*genai.Content{
85+
{
86+
Role: "user",
87+
Parts: []*genai.Part{
88+
{
89+
Text: content,
90+
},
91+
},
92+
},
93+
}
6794

68-
// Send the message to the generative model.
69-
resp, err := session.SendMessage(ctx, genai.Text(content))
95+
resp, err := c.client.Models.GenerateContent(ctx, c.model, data, cfg)
7096
if err != nil {
7197
return nil, err
7298
}
7399

74-
part := resp.Candidates[0].Content.Parts[0]
100+
usage := core.Usage{}
101+
if resp.UsageMetadata != nil {
102+
usage.PromptTokens = int(resp.UsageMetadata.PromptTokenCount)
103+
usage.CompletionTokens = int(resp.UsageMetadata.CandidatesTokenCount)
104+
usage.TotalTokens = int(resp.UsageMetadata.TotalTokenCount)
105+
if resp.UsageMetadata.CachedContentTokenCount > 0 {
106+
usage.PromptTokensDetails = &openai.PromptTokensDetails{
107+
CachedTokens: int(resp.UsageMetadata.CachedContentTokenCount),
108+
}
109+
}
110+
}
75111

76-
usage := core.Usage{
77-
PromptTokens: int(resp.UsageMetadata.PromptTokenCount),
78-
CompletionTokens: int(resp.UsageMetadata.CandidatesTokenCount),
79-
TotalTokens: int(resp.UsageMetadata.TotalTokenCount),
112+
if len(resp.Candidates) == 0 {
113+
return nil, errors.New("no candidates found")
80114
}
81115

82-
if resp.UsageMetadata.CachedContentTokenCount > 0 {
83-
usage.PromptTokensDetails = &openai.PromptTokensDetails{
84-
CachedTokens: int(resp.UsageMetadata.CachedContentTokenCount),
85-
}
116+
cand := resp.Candidates[0]
117+
if len(cand.Content.Parts) == 0 {
118+
return nil, errors.New("no content found")
86119
}
87120

88-
r := &core.Response{
89-
Content: strings.TrimSpace(strings.TrimSuffix(fmt.Sprintf("%v", part), "\n")),
90-
Usage: usage,
121+
part := cand.Content.Parts[0]
122+
if part.FunctionCall == nil || part.FunctionCall.Name != "get_summary_prefix" {
123+
return nil, errors.New("no function call found")
124+
}
125+
126+
prefix, ok := part.FunctionCall.Args["prefix"].(string)
127+
if !ok || prefix == "" {
128+
return nil, errors.New("no prefix found")
91129
}
92130

93131
if c.debug {
94-
// Check that you got the expected function call back.
95-
funcall, ok := part.(genai.FunctionCall)
96-
if !ok {
97-
return nil, fmt.Errorf("expected type FunctionCall, got %T", part)
98-
}
99-
if g, e := funcall.Name, summaryPrefixFunc.FunctionDeclarations[0].Name; g != e {
100-
return nil, fmt.Errorf("expected FunctionCall.Name %q, got %q", e, g)
101-
}
132+
_ = godump.Dump(resp.Candidates)
133+
}
134+
135+
r := &core.Response{
136+
Content: prefix,
137+
Usage: usage,
102138
}
103139

104140
return r, nil
@@ -113,13 +149,6 @@ func New(ctx context.Context, opts ...Option) (c *Client, err error) {
113149
return nil, err
114150
}
115151

116-
// Create a new client instance with the necessary fields.
117-
engine := &Client{
118-
model: cfg.model,
119-
maxTokens: cfg.maxTokens,
120-
temperature: cfg.temperature,
121-
}
122-
123152
// Inject x-app-name and x-app-version headers using core/transport.DefaultHeaderTransport
124153
httpClient := &http.Client{
125154
Transport: &transport.DefaultHeaderTransport{
@@ -130,15 +159,22 @@ func New(ctx context.Context, opts ...Option) (c *Client, err error) {
130159
},
131160
}
132161

133-
client, err := genai.NewClient(ctx, option.WithAPIKey(cfg.token), option.WithHTTPClient(httpClient))
162+
client, err := genai.NewClient(ctx, &genai.ClientConfig{
163+
APIKey: cfg.token,
164+
HTTPClient: httpClient,
165+
Backend: genai.BackendGeminiAPI,
166+
})
134167
if err != nil {
135168
return nil, err
136169
}
137170

138-
engine.client = client.GenerativeModel(engine.model)
139-
engine.client.MaxOutputTokens = convert.ToPtr(engine.maxTokens)
140-
engine.client.Temperature = convert.ToPtr(engine.temperature)
141-
engine.client.TopP = convert.ToPtr(engine.topP)
171+
engine := &Client{
172+
client: client,
173+
model: cfg.model,
174+
maxTokens: cfg.maxTokens,
175+
temperature: cfg.temperature,
176+
topP: cfg.topP,
177+
}
142178

143179
return engine, nil
144180
}

0 commit comments

Comments
 (0)