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