gemini.go

  1package gemini
  2
  3import (
  4	"bytes"
  5	"context"
  6	"encoding/json"
  7	"fmt"
  8	"io"
  9	"net/http"
 10)
 11
 12// https://ai.google.dev/api/generate-content#request-body
 13type Request struct {
 14	// Field order matters for JSON serialization - stable fields should come first
 15	// to maximize prefix deduplication when storing LLM requests.
 16	CachedContent     string            `json:"cachedContent,omitempty"` // format: "cachedContents/{name}"
 17	GenerationConfig  *GenerationConfig `json:"generationConfig,omitempty"`
 18	SystemInstruction *Content          `json:"systemInstruction,omitempty"`
 19	Tools             []Tool            `json:"tools,omitempty"`
 20	// ToolConfig has been left out because it does not appear to be useful.
 21	// Contents comes last since it grows with each request in a conversation
 22	Contents []Content `json:"contents"`
 23}
 24
 25// https://ai.google.dev/api/generate-content#response-body
 26type Response struct {
 27	Candidates []Candidate `json:"candidates"`
 28	headers    http.Header // captured HTTP response headers
 29}
 30
 31// Header returns the HTTP response headers.
 32func (r *Response) Header() http.Header {
 33	return r.headers
 34}
 35
 36type Candidate struct {
 37	Content Content `json:"content"`
 38}
 39
 40type Content struct {
 41	Parts []Part `json:"parts"`
 42	Role  string `json:"role,omitempty"`
 43}
 44
 45// Part is a part of the content.
 46// This is a union data structure, only one-of the fields can be set.
 47type Part struct {
 48	Text                string               `json:"text,omitempty"`
 49	FunctionCall        *FunctionCall        `json:"functionCall,omitempty"`
 50	FunctionResponse    *FunctionResponse    `json:"functionResponse,omitempty"`
 51	ExecutableCode      *ExecutableCode      `json:"executableCode,omitempty"`
 52	CodeExecutionResult *CodeExecutionResult `json:"codeExecutionResult,omitempty"`
 53	// ThoughtSignature is required for Gemini 3 models when using function calling.
 54	// It must be passed back exactly as received when sending the conversation history.
 55	ThoughtSignature string `json:"thoughtSignature,omitempty"`
 56	// TODO inlineData
 57	// TODO fileData
 58}
 59
 60type FunctionCall struct {
 61	Name string         `json:"name"`
 62	Args map[string]any `json:"args"`
 63}
 64
 65type FunctionResponse struct {
 66	Name     string         `json:"name"`
 67	Response map[string]any `json:"response"`
 68}
 69
 70type ExecutableCode struct {
 71	Language Language `json:"language"`
 72	Code     string   `json:"code"`
 73}
 74
 75type Language int
 76
 77const (
 78	LanguageUnspecified Language = 0
 79	LanguagePython      Language = 1 // python >= 3.10 with numpy and simpy
 80)
 81
 82type CodeExecutionResult struct {
 83	Outcome Outcome `json:"outcome"`
 84	Output  string  `json:"output"`
 85}
 86
 87type Outcome int
 88
 89const (
 90	OutcomeUnspecified      Outcome = 0
 91	OutcomeOK               Outcome = 1
 92	OutcomeFailed           Outcome = 2
 93	OutcomeDeadlineExceeded Outcome = 3
 94)
 95
 96// https://ai.google.dev/api/generate-content#v1beta.GenerationConfig
 97type GenerationConfig struct {
 98	ResponseMimeType string  `json:"responseMimeType,omitempty"` // text/plain, application/json, or text/x.enum
 99	ResponseSchema   *Schema `json:"responseSchema,omitempty"`   // for JSON
100}
101
102// https://ai.google.dev/api/caching#Tool
103type Tool struct {
104	FunctionDeclarations []FunctionDeclaration `json:"functionDeclarations"`
105	CodeExecution        *struct{}             `json:"codeExecution,omitempty"` // if present, enables the model to execute code
106	// TODO googleSearchRetrieval https://ai.google.dev/api/caching#GoogleSearchRetrieval
107}
108
109// https://ai.google.dev/api/caching#FunctionDeclaration
110type FunctionDeclaration struct {
111	Name        string `json:"name"`
112	Description string `json:"description"`
113	Parameters  Schema `json:"parameters"`
114}
115
116// https://ai.google.dev/api/caching#Schema
117type Schema struct {
118	Type        DataType          `json:"type"`
119	Format      string            `json:"string,omitempty"` // for NUMBER type: float, double for INTEGER type: int32, int64 for STRING type: enum
120	Description string            `json:"description,omitempty"`
121	Nullable    *bool             `json:"nullable,omitempty"`
122	Enum        []string          `json:"enum,omitempty"`
123	MaxItems    string            `json:"maxItems,omitempty"`   // for ARRAY
124	MinItems    string            `json:"minItems,omitempty"`   // for ARRAY
125	Properties  map[string]Schema `json:"properties,omitempty"` // for OBJECT
126	Required    []string          `json:"required,omitempty"`   // for OBJECT
127	Items       *Schema           `json:"items,omitempty"`      // for ARRAY
128}
129
130type DataType int
131
132const (
133	DataTypeUNSPECIFIED = DataType(0) // Not specified, should not be used.
134	DataTypeSTRING      = DataType(1)
135	DataTypeNUMBER      = DataType(2)
136	DataTypeINTEGER     = DataType(3)
137	DataTypeBOOLEAN     = DataType(4)
138	DataTypeARRAY       = DataType(5)
139	DataTypeOBJECT      = DataType(6)
140)
141
142const defaultEndpoint = "https://generativelanguage.googleapis.com/v1beta"
143
144type Model struct {
145	Model    string // e.g. "models/gemini-1.5-flash"
146	APIKey   string
147	HTTPC    *http.Client // if nil, http.DefaultClient is used
148	Endpoint string       // if empty, DefaultEndpoint is used
149}
150
151func (m Model) GenerateContent(ctx context.Context, req *Request) (*Response, error) {
152	reqBytes, err := json.Marshal(req)
153	if err != nil {
154		return nil, fmt.Errorf("marshaling request: %w", err)
155	}
156	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, fmt.Sprintf("%s/%s:generateContent?key=%s", m.endpoint(), m.Model, m.APIKey), bytes.NewReader(reqBytes))
157	if err != nil {
158		return nil, fmt.Errorf("creating HTTP request: %w", err)
159	}
160	httpReq.Header.Add("Content-Type", "application/json")
161	httpResp, err := m.httpc().Do(httpReq)
162	if err != nil {
163		return nil, fmt.Errorf("GenerateContent: do: %w", err)
164	}
165	defer httpResp.Body.Close()
166	body, err := io.ReadAll(httpResp.Body)
167	if err != nil {
168		return nil, fmt.Errorf("GenerateContent: reading response body: %w", err)
169	}
170	if httpResp.StatusCode != http.StatusOK {
171		return nil, fmt.Errorf("GenerateContent: HTTP status: %d, %s", httpResp.StatusCode, string(body))
172	}
173	var res Response
174	if err := json.Unmarshal(body, &res); err != nil {
175		return nil, fmt.Errorf("GenerateContent: unmarshaling response: %w, %s", err, string(body))
176	}
177	res.headers = httpResp.Header
178	return &res, nil
179}
180
181func (m Model) endpoint() string {
182	if m.Endpoint != "" {
183		return m.Endpoint
184	}
185	return defaultEndpoint
186}
187
188func (m Model) httpc() *http.Client {
189	if m.HTTPC != nil {
190		return m.HTTPC
191	}
192	return http.DefaultClient
193}