90 lines
2 KiB
Go
90 lines
2 KiB
Go
package gemini
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"net/http"
|
|
"net/url"
|
|
|
|
"github.com/google/generative-ai-go/genai"
|
|
"google.golang.org/api/option"
|
|
)
|
|
|
|
type Client struct {
|
|
client *genai.Client
|
|
model string
|
|
}
|
|
|
|
func NewClient(ctx context.Context, apiKey, baseURL, modelName string) (*Client, error) {
|
|
var opts []option.ClientOption
|
|
if apiKey != "" {
|
|
opts = append(opts, option.WithAPIKey(apiKey))
|
|
}
|
|
|
|
if baseURL != "" {
|
|
parsedURL, err := url.Parse(baseURL)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("invalid base url: %w", err)
|
|
}
|
|
// Custom HTTP client endpoint override for Gemini API proxy/baseURL
|
|
customClient := &http.Client{
|
|
Transport: &baseURITransport{
|
|
BaseURL: parsedURL,
|
|
Transport: http.DefaultTransport,
|
|
},
|
|
}
|
|
opts = append(opts, option.WithHTTPClient(customClient))
|
|
}
|
|
|
|
c, err := genai.NewClient(ctx, opts...)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return &Client{
|
|
client: c,
|
|
model: modelName,
|
|
}, nil
|
|
}
|
|
|
|
func (c *Client) 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")
|
|
}
|
|
|
|
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 *Client) Close() {
|
|
c.client.Close()
|
|
}
|
|
|
|
type baseURITransport struct {
|
|
BaseURL *url.URL
|
|
Transport http.RoundTripper
|
|
}
|
|
|
|
func (t *baseURITransport) RoundTrip(req *http.Request) (*http.Response, error) {
|
|
// Rewrite request URL to use custom BaseURL while preserving path & query
|
|
req.URL.Scheme = t.BaseURL.Scheme
|
|
req.URL.Host = t.BaseURL.Host
|
|
if t.BaseURL.Path != "" {
|
|
req.URL.Path = t.BaseURL.Path + req.URL.Path
|
|
}
|
|
tr := t.Transport
|
|
if tr == nil {
|
|
tr = http.DefaultTransport
|
|
}
|
|
return tr.RoundTrip(req)
|
|
}
|