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 := t.Schema
161 if len(schema) == 0 || !json.Valid(schema) {
162 schema = json.RawMessage(`{"type":"object","properties":{}}`)
163 }
164 decls = append(decls, gemFunctionDecl{
165 Name: t.Name,
166 Description: t.Description,
167 Parameters: schema,
168 })
169 }
170 out.Tools = []gemTool{{FunctionDeclarations: decls}}
171 }
172
173 // Generation config.
174 maxTok := req.MaxTokens
175 if maxTok <= 0 {
176 maxTok = m.MaxOutput
177 }
178 gc := &gemGenerationConfig{Temperature: req.Temperature}
179 if maxTok > 0 {
180 gc.MaxOutputTokens = &maxTok
181 }
182 if req.Reasoning != "" && m.Reasoning {
183 tc := geminiThinkingConfig(m.ID, req.Reasoning)
184 if tc != nil {
185 gc.ThinkingConfig = tc
186 }
187 }
188 out.GenerationConfig = gc
189
190 // Convert messages. When function calling is disabled for the target
191 // model, also remove historical functionCall/functionResponse parts;
192 // image models should receive only text/image content.
193 msgs := req.Messages
194 if functionsEnabled {
195 msgs = RepairOrphanedToolResults(req.Messages)
196 }
197 for _, msg := range msgs {
198 switch msg.Role {
199 case RoleUser:
200 parts := convertGemUserParts(msg.Content)
201 if len(parts) == 0 {
202 continue
203 }
204 out.Contents = append(out.Contents, gemContent{Role: "user", Parts: parts})
205 case RoleAssistant:
206 parts := convertGemAssistantParts(msg.Content, functionsEnabled)
207 if len(parts) == 0 {
208 continue
209 }
210 out.Contents = append(out.Contents, gemContent{Role: "model", Parts: parts})
211 case RoleTool:
212 if !functionsEnabled {
213 continue
214 }
215 // Each tool_result becomes a user-role message with a
216 // functionResponse part. Gemini's protocol uses
217 // "user" role for tool replies.
218 parts := convertGemToolResultParts(msg.Content)
219 if len(parts) == 0 {
220 continue
221 }
222 out.Contents = append(out.Contents, gemContent{Role: "user", Parts: parts})
223 }
224 }
225
226 return out, m.ID, nil
227}
228
229func convertGemUserParts(blocks []Content) []gemPart {
230 var parts []gemPart
231 for _, b := range blocks {
232 switch v := b.(type) {
233 case TextBlock:
234 if v.Text == "" {
235 continue
236 }
237 parts = append(parts, gemPart{Text: v.Text})
238 case ImageBlock:
239 parts = append(parts, gemPart{
240 InlineData: &gemInlineData{
241 MimeType: v.MimeType,
242 Data: base64.StdEncoding.EncodeToString(v.Data),
243 },
244 })
245 }
246 }
247 return parts
248}
249
250func convertGemAssistantParts(blocks []Content, functionsEnabled bool) []gemPart {
251 var parts []gemPart
252 for _, b := range blocks {
253 switch v := b.(type) {
254 case TextBlock:
255 if strings.TrimSpace(v.Text) == "" {
256 continue
257 }
258 parts = append(parts, gemPart{Text: v.Text})
259 case ToolCallBlock:
260 if !functionsEnabled {
261 continue
262 }
263 args := v.Arguments
264 if len(args) == 0 || !json.Valid(args) {
265 args = json.RawMessage("{}")
266 }
267 parts = append(parts, gemPart{
268 FunctionCall: &gemFunctionCall{
269 Name: v.Name,
270 Args: args,
271 },
272 })
273 }
274 }
275 return parts
276}
277
278func geminiSupportsFunctionCalling(modelID string) bool {
279 id := strings.ToLower(modelID)
280 // Gemini image generation/editing models accept direct multimodal
281 // prompts but reject tools/function declarations.
282 if strings.Contains(id, "flash-image") || strings.Contains(id, "image-generation") {
283 return false
284 }
285 return true
286}
287
288func saveGeminiImageToWorkingDir(mimeType string, data []byte) (string, error) {
289 ext := ".png"
290 switch strings.ToLower(mimeType) {
291 case "image/jpeg", "image/jpg":
292 ext = ".jpg"
293 case "image/webp":
294 ext = ".webp"
295 case "image/gif":
296 ext = ".gif"
297 }
298 name := "zot-gemini-image-" + uuid.NewString() + ext
299 path := filepath.Join(".", name)
300 if err := os.WriteFile(path, data, 0o644); err != nil {
301 return "", err
302 }
303 return path, nil
304}
305
306func convertGemToolResultParts(blocks []Content) []gemPart {
307 var parts []gemPart
308 for _, b := range blocks {
309 tr, ok := b.(ToolResultBlock)
310 if !ok {
311 continue
312 }
313 // Flatten text content. Image content in tool results is dropped
314 // for Gemini < 3 (multimodal function responses are 3+); for the
315 // common path (text output) we wrap as {"output": "..."} or
316 // {"error": "..."} per the SDK's convention.
317 var sb strings.Builder
318 for _, c := range tr.Content {
319 if tb, ok := c.(TextBlock); ok {
320 if sb.Len() > 0 {
321 sb.WriteString("\n")
322 }
323 sb.WriteString(tb.Text)
324 }
325 }
326 key := "output"
327 if tr.IsError {
328 key = "error"
329 }
330 // Gemini wants response to be an object, not a string.
331 respObj := map[string]string{key: sb.String()}
332 respBytes, err := json.Marshal(respObj)
333 if err != nil {
334 respBytes = []byte(`{"output":""}`)
335 }
336 // Tool name is required on functionResponse; the original
337 // ToolCallBlock has it but ToolResultBlock only carries the
338 // call id. We thread the name back via the call id by using
339 // the id as the name fallback — Gemini ignores the name on
340 // the response side as long as it's non-empty.
341 name := tr.CallID
342 parts = append(parts, gemPart{
343 FunctionResponse: &gemFunctionResponse{
344 Name: name,
345 Response: respBytes,
346 },
347 })
348 }
349 return parts
350}
351
352// geminiThinkingConfig maps zot's reasoning level ("low"/"medium"/"high")
353// to Gemini's thinkingConfig. The right knob depends on the model
354// generation: 2.5 family uses thinkingBudget (tokens), 3.x uses
355// thinkingLevel (enum). Returns nil when the level is unrecognised.
356func geminiThinkingConfig(modelID, level string) *gemThinkingConfig {
357 level = strings.ToLower(level)
358 id := strings.ToLower(modelID)
359
360 // Gemini 3.x: enum-based thinkingLevel. Pro can't go below LOW.
361 if strings.Contains(id, "gemini-3") {
362 isPro := strings.Contains(id, "-pro")
363 var lvl string
364 switch level {
365 case "low":
366 lvl = "LOW"
367 case "medium":
368 if isPro {
369 lvl = "HIGH"
370 } else {
371 lvl = "MEDIUM"
372 }
373 case "high":
374 lvl = "HIGH"
375 default:
376 return nil
377 }
378 return &gemThinkingConfig{IncludeThoughts: true, ThinkingLevel: lvl}
379 }
380
381 // Gemini 2.5 family: token-budget per-model.
382 var budget int
383 switch {
384 case strings.Contains(id, "2.5-pro"):
385 switch level {
386 case "low":
387 budget = 2048
388 case "medium":
389 budget = 8192
390 case "high":
391 budget = 32768
392 }
393 case strings.Contains(id, "2.5-flash-lite"):
394 switch level {
395 case "low":
396 budget = 2048
397 case "medium":
398 budget = 8192
399 case "high":
400 budget = 24576
401 }
402 case strings.Contains(id, "2.5-flash"):
403 switch level {
404 case "low":
405 budget = 2048
406 case "medium":
407 budget = 8192
408 case "high":
409 budget = 24576
410 }
411 default:
412 return nil
413 }
414 if budget <= 0 {
415 return nil
416 }
417 return &gemThinkingConfig{IncludeThoughts: true, ThinkingBudget: &budget}
418}
419
420// ---- streaming ----
421
422func (c *geminiClient) Stream(ctx context.Context, req Request) (<-chan Event, error) {
423 wire, modelID, err := c.buildRequest(req)
424 if err != nil {
425 return nil, err
426 }
427 body, err := json.Marshal(wire)
428 if err != nil {
429 return nil, err
430 }
431
432 // Gemini's streaming endpoint is a POST with ?alt=sse to get
433 // an EventSource-compatible response. Without alt=sse the
434 // server returns a JSON array (one element per chunk), which
435 // we'd need a different parser for.
436 apiPath := fmt.Sprintf("/v1beta/models/%s:streamGenerateContent", modelID)
437 url := c.baseURL + apiPath + "?alt=sse"
438
439 newReq := func() (*http.Request, error) {
440 httpReq, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(body))
441 if err != nil {
442 return nil, err
443 }
444 httpReq.Header.Set("content-type", "application/json")
445 httpReq.Header.Set("accept", "text/event-stream")
446 // Two ways to authenticate against the Generative Language
447 // API: x-goog-api-key header or ?key= query param. We use
448 // the header so the key never lands in proxy access logs.
449 httpReq.Header.Set("x-goog-api-key", c.apiKey)
450 return httpReq, nil
451 }
452
453 resp, err := doStreamWithRetry(ctx, c.http, newReq)
454 if err != nil {
455 return nil, fmt.Errorf("google: %w", err)
456 }
457 if resp.StatusCode != http.StatusOK {
458 b, _ := io.ReadAll(resp.Body)
459 resp.Body.Close()
460 return nil, fmt.Errorf("google: http %d: %s", resp.StatusCode, strings.TrimSpace(string(b)))
461 }
462
463 out := make(chan Event, 16)
464 go c.runStream(ctx, resp, req, out)
465 return out, nil
466}
467
468func (c *geminiClient) runStream(ctx context.Context, resp *http.Response, req Request, out chan<- Event) {
469 defer close(out)
470 defer resp.Body.Close()
471
472 model, _ := FindModel("google", req.Model)
473 out <- EventStart{Model: req.Model, Provider: "google"}
474
475 raw := make(chan sseEvent, 16)
476 go readSSE(resp.Body, raw)
477
478 // Gemini's SSE stream is a sequence of complete JSON
479 // GenerateContentResponse objects, one per data: line. Each
480 // candidate carries a list of parts (text or functionCall),
481 // possibly accumulating across chunks.
482 type blockEntry struct {
483 kind string // "text" | "image" | "tool_use"
484 textBuf strings.Builder
485 image *ImageBlock
486 imagePath string
487 toolID string
488 toolName string
489 toolArgs strings.Builder
490 }
491 var (
492 blocks []*blockEntry
493 currentText *blockEntry
494 usage Usage
495 stop StopReason = StopEnd
496 finalErr error
497 toolCounter int
498 )
499
500 appendText := func(delta string) {
501 if currentText == nil {
502 currentText = &blockEntry{kind: "text"}
503 blocks = append(blocks, currentText)
504 }
505 currentText.textBuf.WriteString(delta)
506 }
507
508 appendImage := func(mimeType, dataB64 string) {
509 data, err := base64.StdEncoding.DecodeString(dataB64)
510 if err != nil || len(data) == 0 {
511 return
512 }
513 img := ImageBlock{MimeType: mimeType, Data: data}
514 path, _ := saveGeminiImageToWorkingDir(mimeType, data)
515 blocks = append(blocks, &blockEntry{kind: "image", image: &img, imagePath: path})
516 // Image blocks break the current text run.
517 currentText = nil
518 }
519
520 startTool := func(name string, providedID string, args json.RawMessage) *blockEntry {
521 toolCounter++
522 id := providedID
523 if id == "" {
524 id = fmt.Sprintf("%s_%d_%d", name, time.Now().UnixNano(), toolCounter)
525 }
526 t := &blockEntry{
527 kind: "tool_use",
528 toolID: id,
529 toolName: name,
530 }
531 if len(args) > 0 && json.Valid(args) {
532 t.toolArgs.Write(args)
533 }
534 blocks = append(blocks, t)
535 // New tool block breaks the current text run.
536 currentText = nil
537 return t
538 }
539
540 assembleMsg := func() Message {
541 content := []Content{}
542 for _, b := range blocks {
543 switch b.kind {
544 case "text":
545 if b.textBuf.Len() > 0 {
546 content = append(content, TextBlock{Text: b.textBuf.String()})
547 }
548 case "image":
549 if b.image != nil && len(b.image.Data) > 0 {
550 content = append(content, *b.image)
551 if b.imagePath != "" {
552 content = append(content, TextBlock{Text: fmt.Sprintf("Saved image: `%s`", b.imagePath)})
553 }
554 }
555 case "tool_use":
556 args := b.toolArgs.String()
557 if args == "" || !json.Valid([]byte(args)) {
558 args = "{}"
559 }
560 content = append(content, ToolCallBlock{
561 ID: b.toolID, Name: b.toolName, Arguments: json.RawMessage(args),
562 })
563 }
564 }
565 return Message{Role: RoleAssistant, Content: content, Time: time.Now()}
566 }
567
568 sendDone := func() {
569 usage.CostUSD = ComputeCost(model, usage)
570 out <- EventUsage{Usage: usage}
571 out <- EventDone{Stop: stop, Err: finalErr, Message: assembleMsg()}
572 }
573
574 for {
575 select {
576 case <-ctx.Done():
577 stop = StopAborted
578 finalErr = ctx.Err()
579 sendDone()
580 return
581 case ev, ok := <-raw:
582 if !ok {
583 sendDone()
584 return
585 }
586 if strings.TrimSpace(ev.Data) == "" {
587 continue
588 }
589 var chunk struct {
590 Candidates []struct {
591 Content struct {
592 Role string `json:"role"`
593 Parts []struct {
594 Text string `json:"text"`
595 InlineData *gemInlineData `json:"inlineData"`
596 Thought bool `json:"thought"`
597 FunctionCall *gemFunctionCall `json:"functionCall"`
598 FunctionCallID string `json:"id"`
599 FunctionCallName string `json:"name"`
600 } `json:"parts"`
601 } `json:"content"`
602 FinishReason string `json:"finishReason"`
603 } `json:"candidates"`
604 UsageMetadata *struct {
605 PromptTokenCount int `json:"promptTokenCount"`
606 CandidatesTokenCount int `json:"candidatesTokenCount"`
607 ThoughtsTokenCount int `json:"thoughtsTokenCount"`
608 CachedContentTokenCount int `json:"cachedContentTokenCount"`
609 TotalTokenCount int `json:"totalTokenCount"`
610 } `json:"usageMetadata"`
611 Error *struct {
612 Code int `json:"code"`
613 Message string `json:"message"`
614 Status string `json:"status"`
615 } `json:"error"`
616 }
617 if err := json.Unmarshal([]byte(ev.Data), &chunk); err != nil {
618 continue
619 }
620 if chunk.Error != nil {
621 stop = StopError
622 finalErr = fmt.Errorf("google: %s", chunk.Error.Message)
623 sendDone()
624 return
625 }
626 for _, cand := range chunk.Candidates {
627 for _, part := range cand.Content.Parts {
628 if part.InlineData != nil {
629 appendImage(part.InlineData.MimeType, part.InlineData.Data)
630 continue
631 }
632 if part.FunctionCall != nil {
633 var args json.RawMessage
634 if len(part.FunctionCall.Args) > 0 {
635 args = part.FunctionCall.Args
636 } else {
637 args = json.RawMessage("{}")
638 }
639 t := startTool(part.FunctionCall.Name, part.FunctionCallID, args)
640 out <- EventToolStart{ID: t.toolID, Name: t.toolName}
641 out <- EventToolArgs{ID: t.toolID, Delta: t.toolArgs.String()}
642 out <- EventToolEnd{ID: t.toolID}
643 continue
644 }
645 if part.Text == "" {
646 continue
647 }
648 if part.Thought {
649 // Thinking summaries arrive as text parts
650 // with thought=true. Not surfaced to the
651 // user in v1; could be exposed via a
652 // future ReasoningBlock if useful.
653 continue
654 }
655 appendText(part.Text)
656 out <- EventTextDelta{Delta: part.Text}
657 }
658 switch cand.FinishReason {
659 case "STOP", "":
660 // "" arrives on intermediate chunks; only
661 // promote the explicit terminal value.
662 if cand.FinishReason == "STOP" {
663 stop = StopEnd
664 }
665 case "MAX_TOKENS":
666 stop = StopLength
667 case "SAFETY", "RECITATION", "BLOCKLIST", "PROHIBITED_CONTENT", "SPII", "IMAGE_SAFETY":
668 stop = StopError
669 finalErr = fmt.Errorf("google: response blocked (%s)", cand.FinishReason)
670 }
671 }
672 if chunk.UsageMetadata != nil {
673 // Gemini reports cumulative totals on every chunk
674 // (or close to it), so assign rather than sum.
675 um := chunk.UsageMetadata
676 input := um.PromptTokenCount - um.CachedContentTokenCount
677 if input < 0 {
678 input = um.PromptTokenCount
679 }
680 usage.InputTokens = input
681 usage.OutputTokens = um.CandidatesTokenCount + um.ThoughtsTokenCount
682 usage.CacheReadTokens = um.CachedContentTokenCount
683 }
684 // Promote ToolUse stop when tool calls are present and
685 // the candidate finished cleanly.
686 if stop == StopEnd {
687 for _, b := range blocks {
688 if b.kind == "tool_use" {
689 stop = StopToolUse
690 break
691 }
692 }
693 }
694 }
695 }
696}