A
arktype•4w ago
jacksteamdev

Using ArkType with @google/genai

I had some trouble getting ArkType working with the latest Google AI SDK. Gemini complains about certain JSON Schema fields, so I had to remove them. This seems to work for my use case: 
import { GoogleGenAI } from '@google/genai'
import { config } from '$/config'
import { toGeminiSchema } from './utils'
import { type, type Type } from 'arktype'

const genAI = new GoogleGenAI({ apiKey: config.apiKeys.google })

/**
 * Generates a reponse as JSON data constrained to an Arktype schema.
 */
export async function generateJsonContent<t extends type>({
  text,
  prompt,
  schema,
}: {
  text: string
  prompt: string
  schema: t
}): Promise<t['infer'] | undefined> {
  const response = await genAI.models.generateContent({
    model: config.models.generate,
    contents: `User input: ${text}
    Assistant:`,
    config: {
      systemInstruction: prompt,
      responseMimeType: 'application/json',
      responseJsonSchema: toGeminiSchema(schema),
    },
  })

  const data = parseJsonString.to(schema as type.cast<t['t']>)(response.text)
  if (data instanceof type.errors) {
    console.error(new Error('GoogleGenAI returned an invalid JSON response', { cause: data }))
    return
  }

  return data
}

const data = await generateJsonContent({
  text: 'hello',
  prompt: 'You are a helpful assistant',
  schema: type({ a: 'string', b: 'number' }),
}) // { a: string; b: number }
import { GoogleGenAI } from '@google/genai'
import { config } from '$/config'
import { toGeminiSchema } from './utils'
import { type, type Type } from 'arktype'

const genAI = new GoogleGenAI({ apiKey: config.apiKeys.google })

/**
 * Generates a reponse as JSON data constrained to an Arktype schema.
 */
export async function generateJsonContent<t extends type>({
  text,
  prompt,
  schema,
}: {
  text: string
  prompt: string
  schema: t
}): Promise<t['infer'] | undefined> {
  const response = await genAI.models.generateContent({
    model: config.models.generate,
    contents: `User input: ${text}
    Assistant:`,
    config: {
      systemInstruction: prompt,
      responseMimeType: 'application/json',
      responseJsonSchema: toGeminiSchema(schema),
    },
  })

  const data = parseJsonString.to(schema as type.cast<t['t']>)(response.text)
  if (data instanceof type.errors) {
    console.error(new Error('GoogleGenAI returned an invalid JSON response', { cause: data }))
    return
  }

  return data
}

const data = await generateJsonContent({
  text: 'hello',
  prompt: 'You are a helpful assistant',
  schema: type({ a: 'string', b: 'number' }),
}) // { a: string; b: number }
Hope this saves someone some time!
2 Replies
jacksteamdev
jacksteamdevOP•4w ago
// ./utils.ts

import type { Schema } from '@google/genai'
import { type JsonSchema } from 'arktype'

/** Gemini complains about certain JSON Schema fields */
export function toGeminiSchema<T extends { toJsonSchema(): JsonSchema }>(schema: T): Schema {
  const jsonSchema = schema.toJsonSchema()

  // Recursively remove unsupported fields for Gemini API
  function cleanSchema(obj: any): any {
    if (Array.isArray(obj)) {
      return obj.map(cleanSchema)
    }
    if (typeof obj === 'object' && obj !== null) {
      const newObj: { [key: string]: any } = {}
      for (const key in obj) {
        if (key === '$schema' || key === 'additionalProperties') {
          continue // Skip unsupported fields
        }
        newObj[key] = cleanSchema(obj[key])
      }
      return newObj
    }
    return obj
  }

  return cleanSchema(jsonSchema) as Schema
}
// ./utils.ts

import type { Schema } from '@google/genai'
import { type JsonSchema } from 'arktype'

/** Gemini complains about certain JSON Schema fields */
export function toGeminiSchema<T extends { toJsonSchema(): JsonSchema }>(schema: T): Schema {
  const jsonSchema = schema.toJsonSchema()

  // Recursively remove unsupported fields for Gemini API
  function cleanSchema(obj: any): any {
    if (Array.isArray(obj)) {
      return obj.map(cleanSchema)
    }
    if (typeof obj === 'object' && obj !== null) {
      const newObj: { [key: string]: any } = {}
      for (const key in obj) {
        if (key === '$schema' || key === 'additionalProperties') {
          continue // Skip unsupported fields
        }
        newObj[key] = cleanSchema(obj[key])
      }
      return newObj
    }
    return obj
  }

  return cleanSchema(jsonSchema) as Schema
}
There was also had an issue with the way type.enumerated(...) is converted to a JSON Schema, but I'm not encountering it now, I'll update later if I hit it again.
ssalbdivad
ssalbdivad•3w ago
It is is impossible when a JSON schema consumer has arbitrary restrictions that don't align with the spec to have output that would satisfy those constraints for all possible tools. That said, we have made some changes in recent versions to use the most ubiquitous representations wherever possible wherever multiple are valid. If there are cases you find where we output something that could be simplified in a way that would be better understood by most tools, please open an issue so we can do that by default! We also have lots of upcoming features in 2.2 for JSON schema generation that will help customize some of these behavior directly. Thanks for posting this solution- hopefully if another Gemini users runs into the same issues, they will find this and/or I will point them to it 💓

Did you find this page helpful?