2024年6月15日

如何使用Gemini的Golang API

作者 柳 永强
内容目录

如何使用Gemini的Golang API

在本文中,我们将讨论如何使用Gemini的Golang API来生成内容。我们会提供一个示例代码并详细解释每个部分的功能。

使用Gemini API前提条件

在开始之前,请确保您已经设置了Gemini API密钥,并将其作为环境变量GEMINI_API_KEY

Gemini API示例代码

以下是一个使用Gemini API的完整Golang示例:

package gemini

import (
    "context"
    "fmt"
    "os"

    "nuwa-engineer/pkg/llms"

    "github.com/google/generative-ai-go/genai"
    "google.golang.org/api/option"
)

// Gemini是Gemini API的包装器。
type Gemini struct {
    Client *genai.Client
    Model  *genai.GenerativeModel
}

// NewGemini返回一个新的Gemini客户端。
func NewGemini(ctx context.Context, modelName string) (llms.Model, error) {
    // 从环境变量中获取API密钥
    client, err := genai.NewClient(ctx, option.WithAPIKey(os.Getenv("GEMINI_API_KEY")))
    if err != nil {
        return nil, fmt.Errorf("创建Gemini客户端失败: %w", err)
    }

    model := client.GenerativeModel(modelName)

    return &Gemini{
        Client: client,
        Model:  model,
    }, nil
}

// ContentToString将内容转换为字符串
func (g *Gemini) ContentToString(content *genai.Content) string {
    var str string
    for _, part := range content.Pparts {
        // 检查part是否是Text类型
        if _, ok := part.(genai.Text); ok {
            str += string(part.(genai.Text))
        }
    }
    return str
}

// GenerateContent根据提示生成内容。
func (g *Gemini) GenerateContent(ctx context.Context, prompt string) (string, error) {
    resp, err := g.Model.GenerateContent(ctx, genai.Text(prompt))
    if err != nil {
        return "", fmt.Errorf("生成内容失败: %w", err)
    }

    // 将响应转换为字符串
    return g.ContentToString(resp.Candidates[0].Content), nil
}

// 关闭客户端。
func (g *Gemini) CloseBackend() error {
    return g.Client.Close()
}

详细解释

包和依赖导入

首先,我们导入了需要的包和依赖:

import (
    "context"
    "fmt"
    "os"

    "nuwa-engineer/pkg/llms"

    "github.com/google/generative-ai-go/genai"
    "google.golang.org/api/option"
)

定义Gemini结构体

定义一个Gemini结构体来封装Gemini API客户端和模型:

type Gemini struct {
    Client *genai.Client
    Model  *genai.GenerativeModel
}

创建新的Gemini客户端

NewGemini函数用于创建一个新的Gemini客户端:

func NewGemini(ctx context.Context, modelName string) (llms.Model, error) {
    client, err := genai.NewClient(ctx, option.WithAPIKey(os.Getenv("GEMINI_API_KEY")))
    if err != nil {
        return nil, fmt.Errorf("创建Gemini客户端失败: %w", err)
    }

    model := client.GenerativeModel(modelName)

    return &Gemini{
        Client: client,
        Model:  model,
    }, nil
}

将内容转换为字符串

ContentToString方法将genai.Content转换为字符串:

func (g *Gemini) ContentToString(content *genai.Content) string {
    var str string
    for _, part := range content.Parts {
        if _, ok := part.(genai.Text); ok {
            str += string(part.(genai.Text))
        }
    }
    return str
}

根据提示生成内容

GenerateContent方法根据给定的提示生成内容:

func (g *Gemini) GenerateContent(ctx context.Context, prompt string) (string, error) {
    resp, err := g.Model.GenerateContent(ctx, genai.Text(prompt))
    if err != nil {
        return "", fmt.Errorf("生成内容失败: %w", err)
    }

    return g.ContentToString(resp.Candidates[0].Content), nil
}

关闭客户端

CloseBackend方法用于关闭Gemini客户端:

func (g *Gemini) CloseBackend() error {
    return g.Client.Close()
}

Main 函数

package main

import (
    "context"
    "fmt"
    "log"
    "os"
    "your/package/path/gemini" // 替换为您的实际包路径
)

func main() {
    // 设置上下文
    ctx := context.Background()

    // 定义模型名称(替换为您实际使用的模型名称)
    modelName := "your-model-name"

    // 创建Gemini客户端
    client, err := gemini.NewGemini(ctx, modelName)
    if err != nil {
        log.Fatalf("创建Gemini客户端失败: %v", err)
    }
    defer func() {
        if err := client.CloseBackend(); err != nil {
            log.Fatalf("关闭Gemini客户端失败: %v", err)
        }
    }()

    // 定义提示
    prompt := "请生成一段关于人工智能的文章。"

    // 生成内容
    content, err := client.GenerateContent(ctx, prompt)
    if err != nil {
        log.Fatalf("生成内容失败: %v", err)
    }

    // 打印生成的内容
    fmt.Println("生成的内容:")
    fmt.Println(content)
}