git.iko.soy/crush
on main help sign in
fix gemini provider
a3a04d8Kujtim Hoxha
internal/llm/provider/gemini.go+38−71
compare with main
unifiedsplit
@@ -54,19 +54,6 @@ func newGeminiClient(opts providerClientOptions) GeminiClient {
54
54
55func (g *geminiClient) convertMessages(messages []message.Message) []*genai.Content {
55func (g *geminiClient) convertMessages(messages []message.Message) []*genai.Content {
56 var history []*genai.Content
56 var history []*genai.Content
57−
58− // Add system message first
59− history = append(history, &genai.Content{
60− Parts: []genai.Part{genai.Text(g.providerOptions.systemMessage)},
61− Role: "user",
62− })
63−
64− // Add a system response to acknowledge the system message
65− history = append(history, &genai.Content{
66− Parts: []genai.Part{genai.Text("I'll help you with that.")},
67− Role: "model",
68− })
69−
70 for _, msg := range messages {
57 for _, msg := range messages {
71 switch msg.Role {
58 switch msg.Role {
72 case message.User:
59 case message.User:
@@ -154,14 +141,11 @@ func (g *geminiClient) convertTools(tools []tools.BaseTool) []*genai.Tool {
154}
141}
155
142
156func (g *geminiClient) finishReason(reason genai.FinishReason) message.FinishReason {
143func (g *geminiClient) finishReason(reason genai.FinishReason) message.FinishReason {
157− reasonStr := reason.String()
158 switch {
144 switch {
159− case reasonStr == "STOP":
145+ case reason == genai.FinishReasonStop:
160 return message.FinishReasonEndTurn
146 return message.FinishReasonEndTurn
161− case reasonStr == "MAX_TOKENS":
147+ case reason == genai.FinishReasonMaxTokens:
162 return message.FinishReasonMaxTokens
148 return message.FinishReasonMaxTokens
163− case strings.Contains(reasonStr, "FUNCTION") || strings.Contains(reasonStr, "TOOL"):
164− return message.FinishReasonToolUse
165 default:
149 default:
166 return message.FinishReasonUnknown
150 return message.FinishReasonUnknown
167 }
151 }
@@ -170,7 +154,11 @@ func (g *geminiClient) finishReason(reason genai.FinishReason) message.FinishRea
170func (g *geminiClient) send(ctx context.Context, messages []message.Message, tools []tools.BaseTool) (*ProviderResponse, error) {
154func (g *geminiClient) send(ctx context.Context, messages []message.Message, tools []tools.BaseTool) (*ProviderResponse, error) {
171 model := g.client.GenerativeModel(g.providerOptions.model.APIModel)
155 model := g.client.GenerativeModel(g.providerOptions.model.APIModel)
172 model.SetMaxOutputTokens(int32(g.providerOptions.maxTokens))
156 model.SetMaxOutputTokens(int32(g.providerOptions.maxTokens))
173−
157+ model.SystemInstruction = &genai.Content{
158+ Parts: []genai.Part{
159+ genai.Text(g.providerOptions.systemMessage),
160+ },
161+ }
174 // Convert tools
162 // Convert tools
175 if len(tools) > 0 {
163 if len(tools) > 0 {
176 model.Tools = g.convertTools(tools)
164 model.Tools = g.convertTools(tools)
@@ -188,19 +176,13 @@ func (g *geminiClient) send(ctx context.Context, messages []message.Message, too
188 attempts := 0
176 attempts := 0
189 for {
177 for {
190 attempts++
178 attempts++
179+ var toolCalls []message.ToolCall
191 chat := model.StartChat()
180 chat := model.StartChat()
192 chat.History = geminiMessages[:len(geminiMessages)-1] // All but last message
181 chat.History = geminiMessages[:len(geminiMessages)-1] // All but last message
193
182
194 lastMsg := geminiMessages[len(geminiMessages)-1]
183 lastMsg := geminiMessages[len(geminiMessages)-1]
195− var lastText string
196− for _, part := range lastMsg.Parts {
197− if text, ok := part.(genai.Text); ok {
198− lastText = string(text)
199− break
200− }
201− }
202
184
203− resp, err := chat.SendMessage(ctx, genai.Text(lastText))
185+ resp, err := chat.SendMessage(ctx, lastMsg.Parts...)
204 // If there is an error we are going to see if we can retry the call
186 // If there is an error we are going to see if we can retry the call
205 if err != nil {
187 if err != nil {
206 retry, after, retryErr := g.shouldRetry(attempts, err)
188 retry, after, retryErr := g.shouldRetry(attempts, err)
@@ -220,7 +202,6 @@ func (g *geminiClient) send(ctx context.Context, messages []message.Message, too
220 }
202 }
221
203
222 content := ""
204 content := ""
223− var toolCalls []message.ToolCall
224
205
225 if len(resp.Candidates) > 0 && resp.Candidates[0].Content != nil {
206 if len(resp.Candidates) > 0 && resp.Candidates[0].Content != nil {
226 for _, part := range resp.Candidates[0].Content.Parts {
207 for _, part := range resp.Candidates[0].Content.Parts {
@@ -231,20 +212,25 @@ func (g *geminiClient) send(ctx context.Context, messages []message.Message, too
231 id := "call_" + uuid.New().String()
212 id := "call_" + uuid.New().String()
232 args, _ := json.Marshal(p.Args)
213 args, _ := json.Marshal(p.Args)
233 toolCalls = append(toolCalls, message.ToolCall{
214 toolCalls = append(toolCalls, message.ToolCall{
234− ID: id,
215+ ID: id,
235− Name: p.Name,
216+ Name: p.Name,
236− Input: string(args),
217+ Input: string(args),
237− Type: "function",
218+ Type: "function",
219+ Finished: true,
238 })
220 })
239 }
221 }
240 }
222 }
241 }
223 }
224+ finishReason := g.finishReason(resp.Candidates[0].FinishReason)
225+ if len(toolCalls) > 0 {
226+ finishReason = message.FinishReasonToolUse
227+ }
242
228
243 return &ProviderResponse{
229 return &ProviderResponse{
244 Content: content,
230 Content: content,
245 ToolCalls: toolCalls,
231 ToolCalls: toolCalls,
246 Usage: g.usage(resp),
232 Usage: g.usage(resp),
247− FinishReason: g.finishReason(resp.Candidates[0].FinishReason),
233+ FinishReason: finishReason,
248 }, nil
234 }, nil
249 }
235 }
250}
236}
@@ -252,7 +238,11 @@ func (g *geminiClient) send(ctx context.Context, messages []message.Message, too
252func (g *geminiClient) stream(ctx context.Context, messages []message.Message, tools []tools.BaseTool) <-chan ProviderEvent {
238func (g *geminiClient) stream(ctx context.Context, messages []message.Message, tools []tools.BaseTool) <-chan ProviderEvent {
253 model := g.client.GenerativeModel(g.providerOptions.model.APIModel)
239 model := g.client.GenerativeModel(g.providerOptions.model.APIModel)
254 model.SetMaxOutputTokens(int32(g.providerOptions.maxTokens))
240 model.SetMaxOutputTokens(int32(g.providerOptions.maxTokens))
255−
241+ model.SystemInstruction = &genai.Content{
242+ Parts: []genai.Part{
243+ genai.Text(g.providerOptions.systemMessage),
244+ },
245+ }
256 // Convert tools
246 // Convert tools
257 if len(tools) > 0 {
247 if len(tools) > 0 {
258 model.Tools = g.convertTools(tools)
248 model.Tools = g.convertTools(tools)
@@ -276,18 +266,10 @@ func (g *geminiClient) stream(ctx context.Context, messages []message.Message, t
276 for {
266 for {
277 attempts++
267 attempts++
278 chat := model.StartChat()
268 chat := model.StartChat()
279− chat.History = geminiMessages[:len(geminiMessages)-1] // All but last message
269+ chat.History = geminiMessages[:len(geminiMessages)-1]
280−
281 lastMsg := geminiMessages[len(geminiMessages)-1]
270 lastMsg := geminiMessages[len(geminiMessages)-1]
282− var lastText string
283− for _, part := range lastMsg.Parts {
284− if text, ok := part.(genai.Text); ok {
285− lastText = string(text)
286− break
287− }
288− }
289
271
290− iter := chat.SendMessageStream(ctx, genai.Text(lastText))
272+ iter := chat.SendMessageStream(ctx, lastMsg.Parts...)
291
273
292 currentContent := ""
274 currentContent := ""
293 toolCalls := []message.ToolCall{}
275 toolCalls := []message.ToolCall{}
@@ -330,23 +312,23 @@ func (g *geminiClient) stream(ctx context.Context, messages []message.Message, t
330 for _, part := range resp.Candidates[0].Content.Parts {
312 for _, part := range resp.Candidates[0].Content.Parts {
331 switch p := part.(type) {
313 switch p := part.(type) {
332 case genai.Text:
314 case genai.Text:
333− newText := string(p)
315+ delta := string(p)
334− delta := newText[len(currentContent):]
335 if delta != "" {
316 if delta != "" {
336 eventChan <- ProviderEvent{
317 eventChan <- ProviderEvent{
337 Type: EventContentDelta,
318 Type: EventContentDelta,
338 Content: delta,
319 Content: delta,
339 }
320 }
340− currentContent = newText
321+ currentContent += delta
341 }
322 }
342 case genai.FunctionCall:
323 case genai.FunctionCall:
343 id := "call_" + uuid.New().String()
324 id := "call_" + uuid.New().String()
344 args, _ := json.Marshal(p.Args)
325 args, _ := json.Marshal(p.Args)
345 newCall := message.ToolCall{
326 newCall := message.ToolCall{
346− ID: id,
327+ ID: id,
347− Name: p.Name,
328+ Name: p.Name,
348− Input: string(args),
329+ Input: string(args),
349− Type: "function",
330+ Type: "function",
331+ Finished: true,
350 }
332 }
351
333
352 isNew := true
334 isNew := true
@@ -368,37 +350,22 @@ func (g *geminiClient) stream(ctx context.Context, messages []message.Message, t
368 eventChan <- ProviderEvent{Type: EventContentStop}
350 eventChan <- ProviderEvent{Type: EventContentStop}
369
351
370 if finalResp != nil {
352 if finalResp != nil {
353+ finishReason := g.finishReason(finalResp.Candidates[0].FinishReason)
354+ if len(toolCalls) > 0 {
355+ finishReason = message.FinishReasonToolUse
356+ }
371 eventChan <- ProviderEvent{
357 eventChan <- ProviderEvent{
372 Type: EventComplete,
358 Type: EventComplete,
373 Response: &ProviderResponse{
359 Response: &ProviderResponse{
374 Content: currentContent,
360 Content: currentContent,
375 ToolCalls: toolCalls,
361 ToolCalls: toolCalls,
376 Usage: g.usage(finalResp),
362 Usage: g.usage(finalResp),
377− FinishReason: g.finishReason(finalResp.Candidates[0].FinishReason),
363+ FinishReason: finishReason,
378 },
364 },
379 }
365 }
380 return
366 return
381 }
367 }
382
368
383− // If we get here, we need to retry
384− if attempts > maxRetries {
385− eventChan <- ProviderEvent{
386− Type: EventError,
387− Error: fmt.Errorf("maximum retry attempts reached: %d retries", maxRetries),
388− }
389− return
390− }
391−
392− // Wait before retrying
393− select {
394− case <-ctx.Done():
395− if ctx.Err() != nil {
396− eventChan <- ProviderEvent{Type: EventError, Error: ctx.Err()}
397− }
398− return
399− case <-time.After(time.Duration(2000*(1<<(attempts-1))) * time.Millisecond):
400− continue
401− }
402 }
369 }
403 }()
370 }()
404
371