162 lines
3.7 KiB
Go
162 lines
3.7 KiB
Go
package gemini
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
|
|
"github.com/google/generative-ai-go/genai"
|
|
"google.golang.org/api/option"
|
|
)
|
|
|
|
// Client defines the common interface for Gemini interactions
|
|
type Client interface {
|
|
Generate(ctx context.Context, prompt string) (string, error)
|
|
Close() error
|
|
}
|
|
|
|
// --- REST Client (for custom BaseURL / proxy like lar) ---
|
|
|
|
type RestClient struct {
|
|
httpClient *http.Client
|
|
baseURL string
|
|
model string
|
|
}
|
|
|
|
func NewRestClient(apiKey, baseURL, model string) (*RestClient, error) {
|
|
baseURL = strings.TrimSuffix(baseURL, "/")
|
|
return &RestClient{
|
|
httpClient: &http.Client{},
|
|
baseURL: baseURL,
|
|
model: model,
|
|
}, nil
|
|
}
|
|
|
|
type restRequest struct {
|
|
Contents []restContent `json:"contents"`
|
|
}
|
|
|
|
type restContent struct {
|
|
Parts []restPart `json:"parts"`
|
|
}
|
|
|
|
type restPart struct {
|
|
Text string `json:"text"`
|
|
}
|
|
|
|
type restResponse struct {
|
|
Candidates []restCandidate `json:"candidates"`
|
|
}
|
|
|
|
type restCandidate struct {
|
|
Content restContent `json:"content"`
|
|
}
|
|
|
|
func (c *RestClient) Generate(ctx context.Context, prompt string) (string, error) {
|
|
url := fmt.Sprintf("%s/v1beta/models/%s:generateContent", c.baseURL, c.model)
|
|
|
|
reqBody := restRequest{
|
|
Contents: []restContent{
|
|
{
|
|
Parts: []restPart{
|
|
{Text: prompt},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
jsonBytes, err := json.Marshal(reqBody)
|
|
if err != nil {
|
|
return "", fmt.Errorf("failed to marshal rest request: %w", err)
|
|
}
|
|
|
|
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonBytes))
|
|
if err != nil {
|
|
return "", fmt.Errorf("failed to create http request: %w", err)
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
|
|
// Note: Authentication (x-goog-api-key) is injected transparently by the lar proxy transport.
|
|
|
|
resp, err := c.httpClient.Do(req)
|
|
if err != nil {
|
|
return "", fmt.Errorf("http request failed: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
bodyBytes, _ := io.ReadAll(resp.Body)
|
|
if resp.StatusCode != http.StatusOK {
|
|
return "", fmt.Errorf("gemini rest api error (status %d): %s", resp.StatusCode, string(bodyBytes))
|
|
}
|
|
|
|
var res restResponse
|
|
if err := json.Unmarshal(bodyBytes, &res); err != nil {
|
|
return "", fmt.Errorf("failed to decode rest response: %w (body: %s)", err, string(bodyBytes))
|
|
}
|
|
|
|
if len(res.Candidates) == 0 || len(res.Candidates[0].Content.Parts) == 0 {
|
|
return "", fmt.Errorf("empty candidates in rest response")
|
|
}
|
|
|
|
return res.Candidates[0].Content.Parts[0].Text, nil
|
|
}
|
|
|
|
func (c *RestClient) Close() error {
|
|
return nil
|
|
}
|
|
|
|
|
|
// --- gRPC Client (official SDK for direct Google API connection) ---
|
|
|
|
type GrpcClient struct {
|
|
client *genai.Client
|
|
model string
|
|
}
|
|
|
|
func NewGrpcClient(ctx context.Context, apiKey, model string) (*GrpcClient, error) {
|
|
c, err := genai.NewClient(ctx, option.WithAPIKey(apiKey))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &GrpcClient{
|
|
client: c,
|
|
model: model,
|
|
}, nil
|
|
}
|
|
|
|
func (c *GrpcClient) Generate(ctx context.Context, prompt string) (string, error) {
|
|
model := c.client.GenerativeModel(c.model)
|
|
resp, err := model.GenerateContent(ctx, genai.Text(prompt))
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
if len(resp.Candidates) == 0 || len(resp.Candidates[0].Content.Parts) == 0 {
|
|
return "", fmt.Errorf("empty response from Gemini gRPC")
|
|
}
|
|
|
|
part := resp.Candidates[0].Content.Parts[0]
|
|
if textPart, ok := part.(genai.Text); ok {
|
|
return string(textPart), nil
|
|
}
|
|
|
|
return fmt.Sprintf("%v", part), nil
|
|
}
|
|
|
|
func (c *GrpcClient) Close() error {
|
|
return c.client.Close()
|
|
}
|
|
|
|
|
|
// --- Factory helper ---
|
|
|
|
func NewClient(ctx context.Context, apiKey, baseURL, model string) (Client, error) {
|
|
if baseURL != "" {
|
|
return NewRestClient(apiKey, baseURL, model)
|
|
}
|
|
return NewGrpcClient(ctx, apiKey, model)
|
|
}
|