🐕

Gemini で型安全にレスポンスを扱うユーティリティを Valibot を使って実装してみる

に公開

早い・安い・うまい(?)な gemini-2.0-flashですが、gemini は response format に JSON が適用されていないため、自前でパースしてあげる必要があります。(違ったらすみません)
今回は、valibot で型安全にパースするユーティリティを実装してみます。

方針

以下の方針で実装していきます。

  1. レスポンスのフォーマットを渡す
    a. 今回は json or text で実装
  2. プロンプトを gemini に投げる
  3. レスポンスから不要なマークダウンなどを取り除く
  4. 取り除いたら valibot でパース
  5. レスポンスを返す

型定義

responseType ごとにレスポンスの型を定義 > union で結合という感じで書いていきます。

import type { GenerativeModel } from '@google/generative-ai'
import type { BaseSchema, InferOutput } from 'valibot'

/**
 * プロンプトのレスポンスの型
 */
type PromptingResponseType = 'json' | 'text'

/**
 * プロンプトのレスポンスの型
 * responseTypeがjsonの場合は、valibotの型推論を利用して型を推論する
 * responseTypeがtextの場合は、stringを返す
 */
type PromptingResponse<
  T extends PromptingResponseType,
  Schema extends BaseSchema<any, any, any>
> = T extends 'json' ? InferOutput<Schema> : string

/**
 * responseTypeがtextの場合のオプション
 */
type PromptingResponseTextOptions<T extends 'text'> = {
  responseType: T
  options?: never
}

/**
 * responseTypeがjsonの場合のオプション
 */
type PromptingResponseJsonOptions<
  T extends 'json',
  Schema extends BaseSchema<any, any, any>
> = {
  responseType: T
  options: {
    schema: Schema
  }
}

/**
 * プロンプトのオプションの型
 */
type PromptingOptions<
  T extends PromptingResponseType,
  Schema extends BaseSchema<any, any, any>
> = T extends 'json'
  ? PromptingResponseJsonOptions<T, Schema>
  : T extends 'text'
    ? PromptingResponseTextOptions<T>
    : never

/**
 * プロンプトのプロパティの型
 */
type PromptingProps<
  T extends PromptingResponseType,
  Schema extends BaseSchema<any, any, any>
> = {
  prompt: string
  gemini: GenerativeModel
} & PromptingOptions<T, Schema>

関数

const prompting = async <
  T extends PromptingResponseType,
  Schema extends BaseSchema<any, any, any>
>({
  prompt,
  gemini,
  responseType,
  options
}: PromptingProps<T, Schema>): Promise<PromptingResponse<T, Schema>> => {
  try {
    const { response } = await gemini.generateContent(prompt)
    const text = response.text()

    switch (responseType) {
      case 'json':
        return parse(
          options.schema,
      // geminiから返されるマークダウン記法を削除してからパース
          JSON.parse(text.replace(/^```json\n/, '').replace(/\n```$/, ''))
        )
      case 'text':
        return text
      default:
        throw new Error(`Invalid response type: ${responseType}`)
    }
  } catch (error) {
    console.error(error)
    throw new Error(
      `Failed to get response: ${error}, prompt: ${prompt}`
    )
  }
}

使ってみる

const genAI = new GoogleGenerativeAI(apiKey)
const gemini = genAI.getGenerativeModel({ model: 'gemini-2.0-flash' })

const jsonResult = prompting({
  prompt: "URLからメディアタイプを判定〜(レスポンスのフォーマットなども書く)",
  gemini,
  responseType: 'json',
  options: {
    schema: object({
      category: union([
        literal('video'),
        literal('image'),
        literal('audio')
      ]),
      url: optional(string())
    })
  }
})

const textResult = await prompting({
  prompt: 'Hello, gemini!',
  gemini: this.gemini,
  responseType: 'text'
})

下のような感じで型安全に使えます。
そもそもJSON形式でレスポンスが帰ってきていない場合や、こちらから指示したフォーマットに沿っていない場合は、これである程度は検知できるようになったと思います。

おわり

LLM を触っていると、ハルシネーション周りのデバッグ難しいなあと思います。
お、いい感じ!と思っても、ちょっとインプット変わるとあれ?となることも多く、気づいたら時間が溶けていることがしばしば🐕
本番稼働しているサービスは本当にすごいです

Discussion