1package provider
2
3import (
4 "bytes"
5 "context"
6 "encoding/base64"
7 "encoding/json"
8 "fmt"
9 "io"
10 "net/http"
11 "os"
12 "path/filepath"
13 "strings"
14 "time"
15
16 "github.com/google/uuid"
17)
18
19// Google Gemini provider, talking directly to the Generative Language
20// REST API at generativelanguage.googleapis.com. We hand-roll the wire
21// format instead of pulling in @google/genai (the SDK is large and we
22// only need streamGenerateContent + the models list).
23//
24// Auth model: API key only. Google does NOT issue OAuth tokens for
25// consumer Gemini Advanced / Google One AI subscriptions; programmatic
26// access requires either an AI Studio API key (this client) or Vertex
27// AI / GCP service-account credentials (separate provider, not yet
28// implemented in zot).
29
30const geminiDefaultBaseURL = "https://generativelanguage.googleapis.com"
31
32// geminiClient implements Client against the Gemini Generative Language API.
33type geminiClient struct {
34 apiKey string
35 baseURL string
36 http *http.Client
37}
38
39// NewGemini creates a Gemini client using an AI Studio API key.
40// baseURL may be empty; defaults to https://generativelanguage.googleapis.com.
41func NewGemini(apiKey, baseURL string) Client {
42 if baseURL == "" {
43 baseURL = geminiDefaultBaseURL
44 }
45 return &geminiClient{
46 apiKey: apiKey,
47 baseURL: strings.TrimRight(baseURL, "/"),
48 http: &http.Client{Timeout: 0},
49 }
50}
51
52func (c *geminiClient) Name() string { return "google" }
53
54// ---- wire types ----
55//
56// Subset of Gemini's Content / Part / GenerateContentRequest schema.
57// Only the fields zot actually emits or consumes are declared here.
58
59type gemInlineData struct {
60 MimeType string `json:"mimeType"`
61 Data string `json:"data"` // base64
62}
63
64type gemFunctionCall struct {
65 Name string `json:"name"`
66 Args json.RawMessage `json:"args,omitempty"`
67}
68
69type gemFunctionResponse struct {
70 Name string `json:"name"`
71 Response json.RawMessage `json:"response"`
72}
73
74// gemPart is one element of a Content. Exactly one of the optional
75// fields is populated per part. We use pointers so empty values are
76// omitted from the wire, which matters for tool responses (Gemini
77// rejects parts with both `text` and `functionResponse` set).
78type gemPart struct {
79 Text string `json:"text,omitempty"`
80 InlineData *gemInlineData `json:"inlineData,omitempty"`
81 FunctionCall *gemFunctionCall `json:"functionCall,omitempty"`
82 FunctionResponse *gemFunctionResponse `json:"functionResponse,omitempty"`
83 // Thought: true marks a thought-summary part. Outgoing parts
84 // from zot never set this; incoming chunks might.
85 Thought bool `json:"thought,omitempty"`
86}
87
88type gemContent struct {
89 Role string `json:"role"` // "user" | "model"
90 Parts []gemPart `json:"parts"`
91}
92
93type gemFunctionDecl struct {
94 Name string `json:"name"`
95 Description string `json:"description,omitempty"`
96 Parameters json.RawMessage `json:"parameters,omitempty"`
97}
98
99type gemTool struct {
100 FunctionDeclarations []gemFunctionDecl `json:"functionDeclarations,omitempty"`
101}
102
103type gemThinkingConfig struct {
104 IncludeThoughts bool `json:"includeThoughts,omitempty"`
105 ThinkingBudget *int `json:"thinkingBudget,omitempty"`
106 ThinkingLevel string `json:"thinkingLevel,omitempty"`
107}
108
109type gemGenerationConfig struct {
110 Temperature *float32 `json:"temperature,omitempty"`
111 MaxOutputTokens *int `json:"maxOutputTokens,omitempty"`
112 ThinkingConfig *gemThinkingConfig `json:"thinkingConfig,omitempty"`
113}
114
115type gemSystemInstruction struct {
116 Parts []gemPart `json:"parts"`
117}
118
119type gemRequest struct {
120 Contents []gemContent `json:"contents"`
121 SystemInstruction *gemSystemInstruction `json:"systemInstruction,omitempty"`
122 Tools []gemTool `json:"tools,omitempty"`
123 GenerationConfig *gemGenerationConfig `json:"generationConfig,omitempty"`
124}
125
126// ---- request building ----
127
128func (c *geminiClient) buildRequest(req Request) (*gemRequest, string, error) {
129 m, err := FindModel("google", req.Model)
130 if err != nil {
131 // Not in the catalog — still allow custom ids by falling back
132 // to defaults so users can point at unreleased models or
133 // alternate base URLs.
134 m = Model{
135 Provider: "google",
136 ID: req.Model,
137 ContextWindow: 1_000_000,
138 MaxOutput: 8192,
139 Reasoning: strings.Contains(req.Model, "2.5") || strings.Contains(req.Model, "3"),
140 }
141 }
142
143 out := &gemRequest{}
144
145 // System prompt → systemInstruction.parts[0].text.
146 if strings.TrimSpace(req.System) != "" {
147 out.SystemInstruction = &gemSystemInstruction{
148 Parts: []gemPart{{Text: req.System}},
149 }
150 }
151
152 functionsEnabled := geminiSupportsFunctionCalling(m.ID)
153
154 // Convert tool defs. Gemini image-generation models reject function
155 // declarations with "Function calling is not enabled for this model";
156 // for those models, send a direct multimodal prompt instead.
157 if functionsEnabled && len(req.Tools) > 0 {
158 decls := make([]gemFunctionDecl, 0, len(req.Tools))
159 for _, t := range req.Tools {
160 schema := sanitizeGeminiToolSchema(t.Schema)
161 decls = append(decls, gemFunctionDecl{
162 Name: t.Name,
163 Description: t.Description,
164 Parameters: schema,
165 })
166 }
167 out.Tools = []gemTool{{FunctionDeclarations: decls}}
168 }
169
170 // Generation config.
171 maxTok := req.MaxTokens
172 if maxTok <= 0 {
173 maxTok = m.MaxOutput
174 }
175 gc := &gemGenerationConfig{Temperature: req.Temperature}
176 if maxTok > 0 {
177 gc.MaxOutputTokens = &maxTok
178 }
179 if req.Reasoning != "" && m.Reasoning {
180 tc := geminiThinkingConfig(m.ID, req.Reasoning)
181 if tc != nil {
182 gc.ThinkingConfig = tc
183 }
184 }
185 out.GenerationConfig = gc
186
187 // Convert messages. When function calling is disabled for the target
188 // model, also remove historical functionCall/functionResponse parts;
189 // image models should receive only text/image content.
190 msgs := req.Messages
191 if functionsEnabled {
192 msgs = RepairOrphanedToolResults(req.Messages)
193 }
194 for _, msg := range msgs {
195 switch msg.Role {
196 case RoleUser:
197 parts := convertGemUserParts(msg.Content)
198 if len(parts) == 0 {
199 continue
200 }
201 out.Contents = append(out.Contents, gemContent{Role: "user", Parts: parts})
202 case RoleAssistant:
203 parts := convertGemAssistantParts(msg.Content, functionsEnabled)
204 if len(parts) == 0 {
205 continue
206 }
207 out.Contents = append(out.Contents, gemContent{Role: "model", Parts: parts})
208 case RoleTool:
209 if !functionsEnabled {
210 continue
211 }
212 // Each tool_result becomes a user-role message with a
213 // functionResponse part. Gemini's protocol uses
214 // "user" role for tool replies.
215 parts := convertGemToolResultParts(msg.Content)
216 if len(parts) == 0 {
217 continue
218 }
219 out.Contents = append(out.Contents, gemContent{Role: "user", Parts: parts})
220 }
221 }
222
223 return out, m.ID, nil
224}
225
226func sanitizeGeminiToolSchema(schema json.RawMessage) json.RawMessage {
227 if len(schema) == 0 || !json.Valid(schema) {
228 return json.RawMessage(`{"type":"object","properties":{}}`)
229 }
230 var v any
231 if err := json.Unmarshal(schema, &v); err != nil {
232 return json.RawMessage(`{"type":"object","properties":{}}`)
233 }
234 v = stripGeminiUnsupportedSchemaFields(v)
235 out, err := json.Marshal(v)
236 if err != nil || len(out) == 0 || !json.Valid(out) {
237 return json.RawMessage(`{"type":"object","properties":{}}`)
238 }
239 return out
240}
241
242func stripGeminiUnsupportedSchemaFields(v any) any {
243 switch x := v.(type) {
244 case map[string]any:
245 out := make(map[string]any, len(x))
246 for k, val := range x {
247 switch k {
248 case "additionalProperties", "$schema":
249 continue
250 default:
251 out[k] = stripGeminiUnsupportedSchemaFields(val)
252 }
253 }
254 return out
255 case []any:
256 out := make([]any, len(x))
257 for i, val := range x {
258 out[i] = stripGeminiUnsupportedSchemaFields(val)
259 }
260 return out
261 default:
262 return v
263 }
264}
265
266func convertGemUserParts(blocks []Content) []gemPart {
267 var parts []gemPart
268 for _, b := range blocks {
269 switch v := b.(type) {
270 case TextBlock:
271 if v.Text == "" {
272 continue
273 }
274 parts = append(parts, gemPart{Text: v.Text})
275 case ImageBlock:
276 parts = append(parts, gemPart{
277 InlineData: &gemInlineData{
278 MimeType: v.MimeType,
279 Data: base64.StdEncoding.EncodeToString(v.Data),
280 },
281 })
282 }
283 }
284 return parts
285}
286
287func convertGemAssistantParts(blocks []Content, functionsEnabled bool) []gemPart {
288 var parts []gemPart
289 for _, b := range blocks {
290 switch v := b.(type) {
291 case TextBlock:
292 if strings.TrimSpace(v.Text) == "" {
293 continue
294 }
295 parts = append(parts, gemPart{Text: v.Text})
296 case ToolCallBlock:
297 if !functionsEnabled {
298 continue
299 }
300 args := v.Arguments
301 if len(args) == 0 || !json.Valid(args) {
302 args = json.RawMessage("{}")
303 }
304 parts = append(parts, gemPart{
305 FunctionCall: &gemFunctionCall{
306 Name: v.Name,
307 Args: args,
308 },
309 })
310 }
311 }
312 return parts
313}
314
315func geminiSupportsFunctionCalling(modelID string) bool {
316 id := strings.ToLower(modelID)
317 // Gemini image generation/editing models accept direct multimodal
318 // prompts but reject tools/function declarations.
319 if strings.Contains(id, "flash-image") || strings.Contains(id, "image-generation") {
320 return false
321 }
322 return true
323}
324
325func saveGeminiImageToWorkingDir(mimeType string, data []byte) (string, error) {
326 ext := ".png"
327 switch strings.ToLower(mimeType) {
328 case "image/jpeg", "image/jpg":
329 ext = ".jpg"
330 case "image/webp":
331 ext = ".webp"
332 case "image/gif":
333 ext = ".gif"
334 }
335 name := "zot-gemini-image-" + uuid.NewString() + ext
336 path := filepath.Join(".", name)
337 if err := os.WriteFile(path, data, 0o644); err != nil {
338 return "", err
339 }
340 return path, nil
341}
342
343func convertGemToolResultParts(blocks []Content) []gemPart {
344 var parts []gemPart
345 for _, b := range blocks {
346 tr, ok := b.(ToolResultBlock)
347 if !ok {
348 continue
349 }
350 // Flatten text content. Image content in tool results is dropped
351 // for Gemini < 3 (multimodal function responses are 3+); for the
352 // common path (text output) we wrap as {"output": "..."} or
353 // {"error": "..."} per the SDK's convention.
354 var sb strings.Builder
355 for _, c := range tr.Content {
356 if tb, ok := c.(TextBlock); ok {
357 if sb.Len() > 0 {
358 sb.WriteString("\n")
359 }
360 sb.WriteString(tb.Text)
361 }
362 }
363 key := "output"
364 if tr.IsError {
365 key = "error"
366 }
367 // Gemini wants response to be an object, not a string.
368 respObj := map[string]string{key: sb.String()}
369 respBytes, err := json.Marshal(respObj)
370 if err != nil {
371 respBytes = []byte(`{"output":""}`)
372 }
373 // Tool name is required on functionResponse; the original
374 // ToolCallBlock has it but ToolResultBlock only carries the
375 // call id. We thread the name back via the call id by using
376 // the id as the name fallback — Gemini ignores the name on
377 // the response side as long as it's non-empty.
378 name := tr.CallID
379 parts = append(parts, gemPart{
380 FunctionResponse: &gemFunctionResponse{
381 Name: name,
382 Response: respBytes,
383 },
384 })
385 }
386 return parts
387}
388
389// geminiThinkingConfig maps zot's reasoning level ("low"/"medium"/"high")
390// to Gemini's thinkingConfig. The right knob depends on the model
391// generation: 2.5 family uses thinkingBudget (tokens), 3.x uses
392// thinkingLevel (enum). Returns nil when the level is unrecognised.
393func geminiThinkingConfig(modelID, level string) *gemThinkingConfig {
394 level = NormalizeReasoning(level)
395 id := strings.ToLower(modelID)
396
397 // Gemini 3.x: enum-based thinkingLevel. Pro can't go below LOW.
398 if strings.Contains(id, "gemini-3") {
399 isPro := strings.Contains(id, "-pro")
400 var lvl string
401 switch level {
402 case "minimum":
403 if isPro {
404 lvl = "LOW"
405 } else {
406 lvl = "MINIMAL"
407 }
408 case "low":
409 lvl = "LOW"
410 case "medium":
411 if isPro {
412 lvl = "HIGH"
413 } else {
414 lvl = "MEDIUM"
415 }
416 case "high", "maximum":
417 lvl = "HIGH"
418 default:
419 return nil
420 }
421 return &gemThinkingConfig{IncludeThoughts: true, ThinkingLevel: lvl}
422 }
423
424 // Gemini 2.5 family: token-budget per-model.
425 budget := ReasoningBudget(level)
426 switch {
427 case strings.Contains(id, "2.5-pro"):
428 if budget > 32768 {
429 budget = 32768
430 }
431 case strings.Contains(id, "2.5-flash-lite"):
432 if budget > 24576 {
433 budget = 24576
434 }
435 case strings.Contains(id, "2.5-flash"):
436 if budget > 24576 {
437 budget = 24576
438 }
439 default:
440 return nil
441 }
442 if budget <= 0 {
443 return nil
444 }
445 return &gemThinkingConfig{IncludeThoughts: true, ThinkingBudget: &budget}
446}
447
448// ---- streaming ----
449
450func (c *geminiClient) Stream(ctx context.Context, req Request) (<-chan Event, error) {
451 wire, modelID, err := c.buildRequest(req)
452 if err != nil {
453 return nil, err
454 }
455 body, err := json.Marshal(wire)
456 if err != nil {
457 return nil, err
458 }
459
460 // Gemini's streaming endpoint is a POST with ?alt=sse to get
461 // an EventSource-compatible response. Without alt=sse the
462 // server returns a JSON array (one element per chunk), which
463 // we'd need a different parser for.
464 apiPath := fmt.Sprintf("/v1beta/models/%s:streamGenerateContent", modelID)
465 url := c.baseURL + apiPath + "?alt=sse"
466
467 newReq := func() (*http.Request, error) {
468 httpReq, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(body))
469 if err != nil {
470 return nil, err
471 }
472 httpReq.Header.Set("content-type", "application/json")
473 httpReq.Header.Set("accept", "text/event-stream")
474 // Two ways to authenticate against the Generative Language
475 // API: x-goog-api-key header or ?key= query param. We use
476 // the header so the key never lands in proxy access logs.
477 httpReq.Header.Set("x-goog-api-key", c.apiKey)
478 return httpReq, nil
479 }
480
481 resp, err := doStreamWithRetry(ctx, c.http, newReq)
482 if err != nil {
483 return nil, fmt.Errorf("google: %w", err)
484 }
485 if resp.StatusCode != http.StatusOK {
486 b, _ := io.ReadAll(resp.Body)
487 resp.Body.Close()
488 return nil, fmt.Errorf("google: http %d: %s", resp.StatusCode, strings.TrimSpace(string(b)))
489 }
490
491 out := make(chan Event, 16)
492 go c.runStream(ctx, resp, req, out)
493 return out, nil
494}
495
496func (c *geminiClient) runStream(ctx context.Context, resp *http.Response, req Request, out chan<- Event) {
497 defer close(out)
498 defer resp.Body.Close()
499
500 model, _ := FindModel("google", req.Model)
501 out <- EventStart{Model: req.Model, Provider: "google"}
502
503 raw := make(chan sseEvent, 16)
504 go readSSE(resp.Body, raw)
505
506 // Gemini's SSE stream is a sequence of complete JSON
507 // GenerateContentResponse objects, one per data: line. Each
508 // candidate carries a list of parts (text or functionCall),
509 // possibly accumulating across chunks.
510 type blockEntry struct {
511 kind string // "text" | "image" | "tool_use"
512 textBuf strings.Builder
513 image *ImageBlock
514 imagePath string
515 toolID string
516 toolName string
517 toolArgs strings.Builder
518 }
519 var (
520 blocks []*blockEntry
521 currentText *blockEntry
522 usage Usage
523 stop StopReason = StopEnd
524 finalErr error
525 toolCounter int
526 )
527
528 appendText := func(delta string) {
529 if currentText == nil {
530 currentText = &blockEntry{kind: "text"}
531 blocks = append(blocks, currentText)
532 }
533 currentText.textBuf.WriteString(delta)
534 }
535
536 appendImage := func(mimeType, dataB64 string) {
537 data, err := base64.StdEncoding.DecodeString(dataB64)
538 if err != nil || len(data) == 0 {
539 return
540 }
541 img := ImageBlock{MimeType: mimeType, Data: data}
542 path, _ := saveGeminiImageToWorkingDir(mimeType, data)
543 blocks = append(blocks, &blockEntry{kind: "image", image: &img, imagePath: path})
544 // Image blocks break the current text run.
545 currentText = nil
546 }
547
548 startTool := func(name string, providedID string, args json.RawMessage) *blockEntry {
549 toolCounter++
550 id := providedID
551 if id == "" {
552 id = fmt.Sprintf("%s_%d_%d", name, time.Now().UnixNano(), toolCounter)
553 }
554 t := &blockEntry{
555 kind: "tool_use",
556 toolID: id,
557 toolName: name,
558 }
559 if len(args) > 0 && json.Valid(args) {
560 t.toolArgs.Write(args)
561 }
562 blocks = append(blocks, t)
563 // New tool block breaks the current text run.
564 currentText = nil
565 return t
566 }
567
568 assembleMsg := func() Message {
569 content := []Content{}
570 for _, b := range blocks {
571 switch b.kind {
572 case "text":
573 if b.textBuf.Len() > 0 {
574 content = append(content, TextBlock{Text: b.textBuf.String()})
575 }
576 case "image":
577 if b.image != nil && len(b.image.Data) > 0 {
578 content = append(content, *b.image)
579 if b.imagePath != "" {
580 content = append(content, TextBlock{Text: fmt.Sprintf("Saved image: `%s`", b.imagePath)})
581 }
582 }
583 case "tool_use":
584 args := b.toolArgs.String()
585 if args == "" || !json.Valid([]byte(args)) {
586 args = "{}"
587 }
588 content = append(content, ToolCallBlock{
589 ID: b.toolID, Name: b.toolName, Arguments: json.RawMessage(args),
590 })
591 }
592 }
593 return Message{Role: RoleAssistant, Content: content, Time: time.Now()}
594 }
595
596 sendDone := func() {
597 usage.CostUSD = ComputeCost(model, usage)
598 out <- EventUsage{Usage: usage}
599 out <- EventDone{Stop: stop, Err: finalErr, Message: assembleMsg()}
600 }
601
602 for {
603 select {
604 case <-ctx.Done():
605 stop = StopAborted
606 finalErr = ctx.Err()
607 sendDone()
608 return
609 case ev, ok := <-raw:
610 if !ok {
611 sendDone()
612 return
613 }
614 if strings.TrimSpace(ev.Data) == "" {
615 continue
616 }
617 var chunk struct {
618 Candidates []struct {
619 Content struct {
620 Role string `json:"role"`
621 Parts []struct {
622 Text string `json:"text"`
623 InlineData *gemInlineData `json:"inlineData"`
624 Thought bool `json:"thought"`
625 FunctionCall *gemFunctionCall `json:"functionCall"`
626 FunctionCallID string `json:"id"`
627 FunctionCallName string `json:"name"`
628 } `json:"parts"`
629 } `json:"content"`
630 FinishReason string `json:"finishReason"`
631 } `json:"candidates"`
632 UsageMetadata *struct {
633 PromptTokenCount int `json:"promptTokenCount"`
634 CandidatesTokenCount int `json:"candidatesTokenCount"`
635 ThoughtsTokenCount int `json:"thoughtsTokenCount"`
636 CachedContentTokenCount int `json:"cachedContentTokenCount"`
637 TotalTokenCount int `json:"totalTokenCount"`
638 } `json:"usageMetadata"`
639 Error *struct {
640 Code int `json:"code"`
641 Message string `json:"message"`
642 Status string `json:"status"`
643 } `json:"error"`
644 }
645 if err := json.Unmarshal([]byte(ev.Data), &chunk); err != nil {
646 continue
647 }
648 if chunk.Error != nil {
649 stop = StopError
650 finalErr = fmt.Errorf("google: %s", chunk.Error.Message)
651 sendDone()
652 return
653 }
654 for _, cand := range chunk.Candidates {
655 for _, part := range cand.Content.Parts {
656 if part.InlineData != nil {
657 appendImage(part.InlineData.MimeType, part.InlineData.Data)
658 continue
659 }
660 if part.FunctionCall != nil {
661 var args json.RawMessage
662 if len(part.FunctionCall.Args) > 0 {
663 args = part.FunctionCall.Args
664 } else {
665 args = json.RawMessage("{}")
666 }
667 t := startTool(part.FunctionCall.Name, part.FunctionCallID, args)
668 out <- EventToolStart{ID: t.toolID, Name: t.toolName}
669 out <- EventToolArgs{ID: t.toolID, Delta: t.toolArgs.String()}
670 out <- EventToolEnd{ID: t.toolID}
671 continue
672 }
673 if part.Text == "" {
674 continue
675 }
676 if part.Thought {
677 // Thinking summaries arrive as text parts
678 // with thought=true. Not surfaced to the
679 // user in v1; could be exposed via a
680 // future ReasoningBlock if useful.
681 continue
682 }
683 appendText(part.Text)
684 out <- EventTextDelta{Delta: part.Text}
685 }
686 switch cand.FinishReason {
687 case "STOP", "":
688 // "" arrives on intermediate chunks; only
689 // promote the explicit terminal value.
690 if cand.FinishReason == "STOP" {
691 stop = StopEnd
692 }
693 case "MAX_TOKENS":
694 stop = StopLength
695 case "SAFETY", "RECITATION", "BLOCKLIST", "PROHIBITED_CONTENT", "SPII", "IMAGE_SAFETY":
696 stop = StopError
697 finalErr = fmt.Errorf("google: response blocked (%s)", cand.FinishReason)
698 }
699 }
700 if chunk.UsageMetadata != nil {
701 // Gemini reports cumulative totals on every chunk
702 // (or close to it), so assign rather than sum.
703 um := chunk.UsageMetadata
704 input := um.PromptTokenCount - um.CachedContentTokenCount
705 if input < 0 {
706 input = um.PromptTokenCount
707 }
708 usage.InputTokens = input
709 usage.OutputTokens = um.CandidatesTokenCount + um.ThoughtsTokenCount
710 usage.CacheReadTokens = um.CachedContentTokenCount
711 }
712 // Promote ToolUse stop when tool calls are present and
713 // the candidate finished cleanly.
714 if stop == StopEnd {
715 for _, b := range blocks {
716 if b.kind == "tool_use" {
717 stop = StopToolUse
718 break
719 }
720 }
721 }
722 }
723 }
724}