Gemini APIをGoで呼び出してJsonで返してもらう

概要

Vertex AI API for Gemini(以下Gemini APIと呼びます)を使って指示した内容をJsonで返してもらえるようにしました。

事前準備

以下のクイックスタートガイドの通りに、

1. GCPのセットアップ
2. Gemini APIの有効化
3. gcloud cliのインストール
   以上を行います。

https://cloud.google.com/vertex-ai/generative-ai/docs/start/quickstarts/quickstart-multimodal?hl=ja

コード

最終的なコードを示します。

genai.go

package main

import (
  "context"
  "encoding/json"
  "fmt"
  "os"
  "strings"

  "cloud.google.com/go/vertexai/genai"
)

func main() {
  ctx := context.Background()

  GenerateContentFromText(ctx, "Q. パンはパンでも食べられないパンってなーんだ。\nA. 硬すぎるパン")
}

const (
  location  = "asia-northeast1"
  modelName = "gemini-1.5-pro-002"
)

type GenerateContentResponse struct {
  Message string `json:"message"`
  Score   int    `json:"score"`
}

func GenerateContentFromText(ctx context.Context, message string) (*GenerateContentResponse, error) {
  gc, err := genai.NewClient(ctx, os.Getenv("GCP_PROJECT_ID"), location)
  if err != nil {
    return nil, err
  }

  systemInstruction := `クイズをします。私が解答するので採点してください`
  schema := &genai.Schema{
    Type: genai.TypeObject,
    Properties: map[string]*genai.Schema{
      "message": {
        Type:        genai.TypeString,
        Description: "返信内容",
      },
      "score": {
        Type:        genai.TypeInteger,
        Description: "解答の点数。0~100の範囲で採点してください。",
      },
    },
    Required: []string{"message", "score"},
  }

  gemini := client.GenerativeModel(modelName)
  gemini.SystemInstruction = &genai.Content{
    Parts: []genai.Part{genai.Text(systemInstruction)},
  }
  gemini.GenerationConfig.ResponseMIMEType = "application/json"
  gemini.GenerationConfig.ResponseSchema = schema

  prompt := genai.Text(message)
  resp, err := gemini.GenerateContent(ctx, prompt)
  if err != nil {
    return nil, fmt.Errorf("error generating content: %w", err)
  }
  if resp.PromptFeedback != nil {
    return nil, fmt.Errorf("generating content is blocked: %s", resp.PromptFeedback.BlockReasonMessage)
  }
  builder := strings.Builder{}
  for _, candidate := range resp.Candidates {
    for _, part := range candidate.Content.Parts {
      builder.WriteString(fmt.Sprintf("%s", part))
    }
  }
  response := GenerateContentResponse{}
  if err := json.Unmarshal([]byte(builder.String()), &response); err != nil {
    return nil, err
  }
  return &response, nil
}

解説

内容を解説します。

返事をもらうだけの処理

まずGemini APIを使ってただ返事をしてもらうだけの場合は以下のようなコードになります。

func GenerateContentFromText(ctx context.Context, message string) error {
  gc, err := genai.NewClient(ctx, os.Getenv("GCP_PROJECT_ID"), location)
  if err != nil {
    return err
  }

  gemini := client.GenerativeModel(modelName)
  gemini.GenerationConfig.ResponseMIMEType = "application/json"

  prompt := genai.Text(message)
  resp, err := gemini.GenerateContent(ctx, prompt)
  if err != nil {
    return fmt.Errorf("error generating content: %w", err)
  }
  if resp.PromptFeedback != nil {
    return fmt.Errorf("generating content is blocked: %s", resp.PromptFeedback.BlockReasonMessage)
  }
  log.Println(resp)
  return nil
}

このとき返ってくるrespはこのような型をしており、必要な部分だけ抜き出す必要があります。

type GenerateContentResponse struct {
	// Output only. Generated candidates.
	Candidates []*Candidate
	// Output only. Content filter results for a prompt sent in the request.
	// Note: Sent only in the first stream chunk.
	// Only happens when no candidates were generated due to content violations.
	PromptFeedback *PromptFeedback
	// Usage metadata about the response(s).
	UsageMetadata *UsageMetadata
}

https://pkg.go.dev/cloud.google.com/go/vertexai/genai#GenerateContentResponse

Part の中身の型は返事の内容次第で変わるようですが、今回はテキストで返事をしてもらっているので文字列として取り出します。

builder := strings.Builder{}
for _, candidate := range resp.Candidates {
  for _, part := range candidate.Content.Parts {
    builder.WriteString(fmt.Sprintf("%s", part))
  }
}
log.Println(builder.String())

schemeの指示

返事をjsonで返してもらうために、schemeの指定をします。

schema := &genai.Schema{
  Type: genai.TypeObject,
  Properties: map[string]*genai.Schema{
    "message": {
      Type:        genai.TypeString,
      Description: "返信内容",
    },
    "score": {
      Type:        genai.TypeInteger,
      Description: "解答の点数。0~100の範囲で採点してください。",
    },
  },
  Required: []string{"message", "score"},
}
gemini.GenerationConfig.ResponseSchema = schema

すると Description を読んでよしなに値を入れて返してくれるのでパースすれば完了です。

type GenerateContentResponse struct {
  Message string `json:"message"`
  Score   int    `json:"score"`
}

...

response := GenerateContentResponse{}
if err := json.Unmarshal([]byte(builder.String()), &response); err != nil {
  return nil, err
}
log.Println(response)