1package provider
2
3import (
4 "context"
5 "encoding/json"
6 "errors"
7 "fmt"
8 "io"
9 "strings"
10 "time"
11
12 "github.com/charmbracelet/crush/internal/config"
13 "github.com/charmbracelet/crush/internal/llm/tools"
14 "github.com/charmbracelet/crush/internal/logging"
15 "github.com/charmbracelet/crush/internal/message"
16 "github.com/google/uuid"
17 "google.golang.org/genai"
18)
19
20type geminiClient struct {
21 providerOptions providerClientOptions
22 client *genai.Client
23}
24
25type GeminiClient ProviderClient
26
27func newGeminiClient(opts providerClientOptions) GeminiClient {
28 client, err := genai.NewClient(context.Background(), &genai.ClientConfig{APIKey: opts.apiKey, Backend: genai.BackendGeminiAPI})
29 if err != nil {
30 logging.Error("Failed to create Gemini client", "error", err)
31 return nil
32 }
33
34 return &geminiClient{
35 providerOptions: opts,
36 client: client,
37 }
38}
39
40func (g *geminiClient) convertMessages(messages []message.Message) []*genai.Content {
41 var history []*genai.Content
42 for _, msg := range messages {
43 switch msg.Role {
44 case message.User:
45 var parts []*genai.Part
46 parts = append(parts, &genai.Part{Text: msg.Content().String()})
47 for _, binaryContent := range msg.BinaryContent() {
48 imageFormat := strings.Split(binaryContent.MIMEType, "/")
49 parts = append(parts, &genai.Part{InlineData: &genai.Blob{
50 MIMEType: imageFormat[1],
51 Data: binaryContent.Data,
52 }})
53 }
54 history = append(history, &genai.Content{
55 Parts: parts,
56 Role: "user",
57 })
58 case message.Assistant:
59 var assistantParts []*genai.Part
60
61 if msg.Content().String() != "" {
62 assistantParts = append(assistantParts, &genai.Part{Text: msg.Content().String()})
63 }
64
65 if len(msg.ToolCalls()) > 0 {
66 for _, call := range msg.ToolCalls() {
67 args, _ := parseJsonToMap(call.Input)
68 assistantParts = append(assistantParts, &genai.Part{
69 FunctionCall: &genai.FunctionCall{
70 Name: call.Name,
71 Args: args,
72 },
73 })
74 }
75 }
76
77 if len(assistantParts) > 0 {
78 history = append(history, &genai.Content{
79 Role: "model",
80 Parts: assistantParts,
81 })
82 }
83
84 case message.Tool:
85 for _, result := range msg.ToolResults() {
86 response := map[string]any{"result": result.Content}
87 parsed, err := parseJsonToMap(result.Content)
88 if err == nil {
89 response = parsed
90 }
91
92 var toolCall message.ToolCall
93 for _, m := range messages {
94 if m.Role == message.Assistant {
95 for _, call := range m.ToolCalls() {
96 if call.ID == result.ToolCallID {
97 toolCall = call
98 break
99 }
100 }
101 }
102 }
103
104 history = append(history, &genai.Content{
105 Parts: []*genai.Part{
106 {
107 FunctionResponse: &genai.FunctionResponse{
108 Name: toolCall.Name,
109 Response: response,
110 },
111 },
112 },
113 Role: "function",
114 })
115 }
116 }
117 }
118
119 return history
120}
121
122func (g *geminiClient) convertTools(tools []tools.BaseTool) []*genai.Tool {
123 geminiTool := &genai.Tool{}
124 geminiTool.FunctionDeclarations = make([]*genai.FunctionDeclaration, 0, len(tools))
125
126 for _, tool := range tools {
127 info := tool.Info()
128 declaration := &genai.FunctionDeclaration{
129 Name: info.Name,
130 Description: info.Description,
131 Parameters: &genai.Schema{
132 Type: genai.TypeObject,
133 Properties: convertSchemaProperties(info.Parameters),
134 Required: info.Required,
135 },
136 }
137
138 geminiTool.FunctionDeclarations = append(geminiTool.FunctionDeclarations, declaration)
139 }
140
141 return []*genai.Tool{geminiTool}
142}
143
144func (g *geminiClient) finishReason(reason genai.FinishReason) message.FinishReason {
145 switch reason {
146 case genai.FinishReasonStop:
147 return message.FinishReasonEndTurn
148 case genai.FinishReasonMaxTokens:
149 return message.FinishReasonMaxTokens
150 default:
151 return message.FinishReasonUnknown
152 }
153}
154
155func (g *geminiClient) send(ctx context.Context, messages []message.Message, tools []tools.BaseTool) (*ProviderResponse, error) {
156 // Convert messages
157 geminiMessages := g.convertMessages(messages)
158
159 cfg := config.Get()
160 if cfg.Options.Debug {
161 jsonData, _ := json.Marshal(geminiMessages)
162 logging.Debug("Prepared messages", "messages", string(jsonData))
163 }
164
165 history := geminiMessages[:len(geminiMessages)-1] // All but last message
166 lastMsg := geminiMessages[len(geminiMessages)-1]
167 config := &genai.GenerateContentConfig{
168 MaxOutputTokens: int32(g.providerOptions.maxTokens),
169 SystemInstruction: &genai.Content{
170 Parts: []*genai.Part{{Text: g.providerOptions.systemMessage}},
171 },
172 }
173 if len(tools) > 0 {
174 config.Tools = g.convertTools(tools)
175 }
176 model := g.providerOptions.model(g.providerOptions.modelType)
177 chat, _ := g.client.Chats.Create(ctx, model.ID, config, history)
178
179 attempts := 0
180 for {
181 attempts++
182 var toolCalls []message.ToolCall
183
184 var lastMsgParts []genai.Part
185 for _, part := range lastMsg.Parts {
186 lastMsgParts = append(lastMsgParts, *part)
187 }
188 resp, err := chat.SendMessage(ctx, lastMsgParts...)
189 // If there is an error we are going to see if we can retry the call
190 if err != nil {
191 retry, after, retryErr := g.shouldRetry(attempts, err)
192 if retryErr != nil {
193 return nil, retryErr
194 }
195 if retry {
196 logging.WarnPersist(fmt.Sprintf("Retrying due to rate limit... attempt %d of %d", attempts, maxRetries), logging.PersistTimeArg, time.Millisecond*time.Duration(after+100))
197 select {
198 case <-ctx.Done():
199 return nil, ctx.Err()
200 case <-time.After(time.Duration(after) * time.Millisecond):
201 continue
202 }
203 }
204 return nil, retryErr
205 }
206
207 content := ""
208
209 if len(resp.Candidates) > 0 && resp.Candidates[0].Content != nil {
210 for _, part := range resp.Candidates[0].Content.Parts {
211 switch {
212 case part.Text != "":
213 content = string(part.Text)
214 case part.FunctionCall != nil:
215 id := "call_" + uuid.New().String()
216 args, _ := json.Marshal(part.FunctionCall.Args)
217 toolCalls = append(toolCalls, message.ToolCall{
218 ID: id,
219 Name: part.FunctionCall.Name,
220 Input: string(args),
221 Type: "function",
222 Finished: true,
223 })
224 }
225 }
226 }
227 finishReason := message.FinishReasonEndTurn
228 if len(resp.Candidates) > 0 {
229 finishReason = g.finishReason(resp.Candidates[0].FinishReason)
230 }
231 if len(toolCalls) > 0 {
232 finishReason = message.FinishReasonToolUse
233 }
234
235 return &ProviderResponse{
236 Content: content,
237 ToolCalls: toolCalls,
238 Usage: g.usage(resp),
239 FinishReason: finishReason,
240 }, nil
241 }
242}
243
244func (g *geminiClient) stream(ctx context.Context, messages []message.Message, tools []tools.BaseTool) <-chan ProviderEvent {
245 // Convert messages
246 geminiMessages := g.convertMessages(messages)
247
248 cfg := config.Get()
249 if cfg.Options.Debug {
250 jsonData, _ := json.Marshal(geminiMessages)
251 logging.Debug("Prepared messages", "messages", string(jsonData))
252 }
253
254 history := geminiMessages[:len(geminiMessages)-1] // All but last message
255 lastMsg := geminiMessages[len(geminiMessages)-1]
256 config := &genai.GenerateContentConfig{
257 MaxOutputTokens: int32(g.providerOptions.maxTokens),
258 SystemInstruction: &genai.Content{
259 Parts: []*genai.Part{{Text: g.providerOptions.systemMessage}},
260 },
261 }
262 if len(tools) > 0 {
263 config.Tools = g.convertTools(tools)
264 }
265 model := g.providerOptions.model(g.providerOptions.modelType)
266 chat, _ := g.client.Chats.Create(ctx, model.ID, config, history)
267
268 attempts := 0
269 eventChan := make(chan ProviderEvent)
270
271 go func() {
272 defer close(eventChan)
273
274 for {
275 attempts++
276
277 currentContent := ""
278 toolCalls := []message.ToolCall{}
279 var finalResp *genai.GenerateContentResponse
280
281 eventChan <- ProviderEvent{Type: EventContentStart}
282
283 var lastMsgParts []genai.Part
284
285 for _, part := range lastMsg.Parts {
286 lastMsgParts = append(lastMsgParts, *part)
287 }
288 for resp, err := range chat.SendMessageStream(ctx, lastMsgParts...) {
289 if err != nil {
290 retry, after, retryErr := g.shouldRetry(attempts, err)
291 if retryErr != nil {
292 eventChan <- ProviderEvent{Type: EventError, Error: retryErr}
293 return
294 }
295 if retry {
296 logging.WarnPersist(fmt.Sprintf("Retrying due to rate limit... attempt %d of %d", attempts, maxRetries), logging.PersistTimeArg, time.Millisecond*time.Duration(after+100))
297 select {
298 case <-ctx.Done():
299 if ctx.Err() != nil {
300 eventChan <- ProviderEvent{Type: EventError, Error: ctx.Err()}
301 }
302
303 return
304 case <-time.After(time.Duration(after) * time.Millisecond):
305 break
306 }
307 } else {
308 eventChan <- ProviderEvent{Type: EventError, Error: err}
309 return
310 }
311 }
312
313 finalResp = resp
314
315 if len(resp.Candidates) > 0 && resp.Candidates[0].Content != nil {
316 for _, part := range resp.Candidates[0].Content.Parts {
317 switch {
318 case part.Text != "":
319 delta := string(part.Text)
320 if delta != "" {
321 eventChan <- ProviderEvent{
322 Type: EventContentDelta,
323 Content: delta,
324 }
325 currentContent += delta
326 }
327 case part.FunctionCall != nil:
328 id := "call_" + uuid.New().String()
329 args, _ := json.Marshal(part.FunctionCall.Args)
330 newCall := message.ToolCall{
331 ID: id,
332 Name: part.FunctionCall.Name,
333 Input: string(args),
334 Type: "function",
335 Finished: true,
336 }
337
338 isNew := true
339 for _, existing := range toolCalls {
340 if existing.Name == newCall.Name && existing.Input == newCall.Input {
341 isNew = false
342 break
343 }
344 }
345
346 if isNew {
347 toolCalls = append(toolCalls, newCall)
348 }
349 }
350 }
351 }
352 }
353
354 eventChan <- ProviderEvent{Type: EventContentStop}
355
356 if finalResp != nil {
357 finishReason := message.FinishReasonEndTurn
358 if len(finalResp.Candidates) > 0 {
359 finishReason = g.finishReason(finalResp.Candidates[0].FinishReason)
360 }
361 if len(toolCalls) > 0 {
362 finishReason = message.FinishReasonToolUse
363 }
364 eventChan <- ProviderEvent{
365 Type: EventComplete,
366 Response: &ProviderResponse{
367 Content: currentContent,
368 ToolCalls: toolCalls,
369 Usage: g.usage(finalResp),
370 FinishReason: finishReason,
371 },
372 }
373 return
374 }
375 }
376 }()
377
378 return eventChan
379}
380
381func (g *geminiClient) shouldRetry(attempts int, err error) (bool, int64, error) {
382 // Check if error is a rate limit error
383 if attempts > maxRetries {
384 return false, 0, fmt.Errorf("maximum retry attempts reached for rate limit: %d retries", maxRetries)
385 }
386
387 // Gemini doesn't have a standard error type we can check against
388 // So we'll check the error message for rate limit indicators
389 if errors.Is(err, io.EOF) {
390 return false, 0, err
391 }
392
393 errMsg := err.Error()
394 isRateLimit := contains(errMsg, "rate limit", "quota exceeded", "too many requests")
395
396 // Check for common rate limit error messages
397
398 if !isRateLimit {
399 return false, 0, err
400 }
401
402 // Calculate backoff with jitter
403 backoffMs := 2000 * (1 << (attempts - 1))
404 jitterMs := int(float64(backoffMs) * 0.2)
405 retryMs := backoffMs + jitterMs
406
407 return true, int64(retryMs), nil
408}
409
410func (g *geminiClient) usage(resp *genai.GenerateContentResponse) TokenUsage {
411 if resp == nil || resp.UsageMetadata == nil {
412 return TokenUsage{}
413 }
414
415 return TokenUsage{
416 InputTokens: int64(resp.UsageMetadata.PromptTokenCount),
417 OutputTokens: int64(resp.UsageMetadata.CandidatesTokenCount),
418 CacheCreationTokens: 0, // Not directly provided by Gemini
419 CacheReadTokens: int64(resp.UsageMetadata.CachedContentTokenCount),
420 }
421}
422
423func (g *geminiClient) Model() config.Model {
424 return g.providerOptions.model(g.providerOptions.modelType)
425}
426
427// Helper functions
428func parseJsonToMap(jsonStr string) (map[string]any, error) {
429 var result map[string]any
430 err := json.Unmarshal([]byte(jsonStr), &result)
431 return result, err
432}
433
434func convertSchemaProperties(parameters map[string]any) map[string]*genai.Schema {
435 properties := make(map[string]*genai.Schema)
436
437 for name, param := range parameters {
438 properties[name] = convertToSchema(param)
439 }
440
441 return properties
442}
443
444func convertToSchema(param any) *genai.Schema {
445 schema := &genai.Schema{Type: genai.TypeString}
446
447 paramMap, ok := param.(map[string]any)
448 if !ok {
449 return schema
450 }
451
452 if desc, ok := paramMap["description"].(string); ok {
453 schema.Description = desc
454 }
455
456 typeVal, hasType := paramMap["type"]
457 if !hasType {
458 return schema
459 }
460
461 typeStr, ok := typeVal.(string)
462 if !ok {
463 return schema
464 }
465
466 schema.Type = mapJSONTypeToGenAI(typeStr)
467
468 switch typeStr {
469 case "array":
470 schema.Items = processArrayItems(paramMap)
471 case "object":
472 if props, ok := paramMap["properties"].(map[string]any); ok {
473 schema.Properties = convertSchemaProperties(props)
474 }
475 }
476
477 return schema
478}
479
480func processArrayItems(paramMap map[string]any) *genai.Schema {
481 items, ok := paramMap["items"].(map[string]any)
482 if !ok {
483 return nil
484 }
485
486 return convertToSchema(items)
487}
488
489func mapJSONTypeToGenAI(jsonType string) genai.Type {
490 switch jsonType {
491 case "string":
492 return genai.TypeString
493 case "number":
494 return genai.TypeNumber
495 case "integer":
496 return genai.TypeInteger
497 case "boolean":
498 return genai.TypeBoolean
499 case "array":
500 return genai.TypeArray
501 case "object":
502 return genai.TypeObject
503 default:
504 return genai.TypeString // Default to string for unknown types
505 }
506}
507
508func contains(s string, substrs ...string) bool {
509 for _, substr := range substrs {
510 if strings.Contains(strings.ToLower(s), strings.ToLower(substr)) {
511 return true
512 }
513 }
514 return false
515}