feat(ai): introduce Schema.fromEnum<E>()

Author: thatfiredev

firebase-ai/api.txt

@@ -774,6 +774,9 @@ package com.google.firebase.ai.type {
     method public static com.google.firebase.ai.type.Schema enumeration(java.util.List<java.lang.String> values);
     method public static com.google.firebase.ai.type.Schema enumeration(java.util.List<java.lang.String> values, String? description = null);
     method public static com.google.firebase.ai.type.Schema enumeration(java.util.List<java.lang.String> values, String? description = null, boolean nullable = false);
+    method public static <E extends java.lang.Enum<E>> com.google.firebase.ai.type.Schema fromEnum(Class<E> enumClass);
+    method public static <E extends java.lang.Enum<E>> com.google.firebase.ai.type.Schema fromEnum(Class<E> enumClass, String? description = null);
+    method public static <E extends java.lang.Enum<E>> com.google.firebase.ai.type.Schema fromEnum(Class<E> enumClass, String? description = null, boolean nullable = false);
     method public String? getDescription();
     method public java.util.List<java.lang.String>? getEnum();
     method public String? getFormat();
@@ -823,6 +826,10 @@ package com.google.firebase.ai.type {
     method public com.google.firebase.ai.type.Schema enumeration(java.util.List<java.lang.String> values);
     method public com.google.firebase.ai.type.Schema enumeration(java.util.List<java.lang.String> values, String? description = null);
     method public com.google.firebase.ai.type.Schema enumeration(java.util.List<java.lang.String> values, String? description = null, boolean nullable = false);
+    method public <E extends java.lang.Enum<E>> com.google.firebase.ai.type.Schema fromEnum(Class<E> enumClass);
+    method public <E extends java.lang.Enum<E>> com.google.firebase.ai.type.Schema fromEnum(Class<E> enumClass, String? description = null);
+    method public <E extends java.lang.Enum<E>> com.google.firebase.ai.type.Schema fromEnum(Class<E> enumClass, String? description = null, boolean nullable = false);
+    method public inline <reified E extends java.lang.Enum<E>> com.google.firebase.ai.type.Schema fromEnum(String? description = null, boolean nullable = false);
     method public com.google.firebase.ai.type.Schema numDouble();
     method public com.google.firebase.ai.type.Schema numDouble(String? description = null);
     method public com.google.firebase.ai.type.Schema numDouble(String? description = null, boolean nullable = false);

firebase-ai/src/main/kotlin/com/google/firebase/ai/type/Schema.kt

@@ -16,6 +16,7 @@
 
 package com.google.firebase.ai.type
 
+import java.util.EnumSet
 import kotlinx.serialization.Serializable
 
 public abstract class StringFormat private constructor(internal val value: String) {
@@ -179,7 +180,11 @@ internal constructor(
     ): Schema {
       if (!properties.keys.containsAll(optionalProperties)) {
         throw IllegalArgumentException(
-          "All optional properties must be present in properties. Missing: ${optionalProperties.minus(properties.keys)}"
+          "All optional properties must be present in properties. Missing: ${
+            optionalProperties.minus(
+              properties.keys
+            )
+          }"
         )
       }
       return Schema(
@@ -239,6 +244,49 @@ internal constructor(
         enum = values,
         type = "STRING",
       )
+
+    /**
+     * Returns a [Schema] for the given Kotlin Enum.
+     *
+     * For example, the cardinal directions can be represented as:
+     *
+     * ```
+     * enum class CardinalDirection { NORTH, EAST, SOUTH, WEST }
+     *
+     * Schema.fromEnum<CardinalDirection>()
+     * ```
+     *
+     * @param description The description of what the parameter should contain or represent
+     * @param nullable Indicates whether the value can be `null`. Defaults to `false`.
+     */
+    @JvmOverloads
+    public inline fun <reified E : Enum<E>> fromEnum(
+      description: String? = null,
+      nullable: Boolean = false
+    ): Schema = enumeration(enumValues<E>().map { it.toString() }, description, nullable)
+
+    /**
+     * Returns a [Schema] for the given Java Enum.
+     *
+     * For example, the cardinal directions can be represented as:
+     *
+     * ```
+     * enum CardinalDirection { NORTH, EAST, SOUTH, WEST }
+     *
+     * Schema.fromEnum(CardinalDirection.class);
+     * ```
+     *
+     * @param enumClass The Enum's Java class.
+     * @param description The description of what the parameter should contain or represent
+     * @param nullable Indicates whether the value can be `null`. Defaults to `false`.
+     */
+    @JvmStatic
+    @JvmOverloads
+    public fun <E : Enum<E>> fromEnum(
+      enumClass: Class<E>,
+      description: String? = null,
+      nullable: Boolean = false
+    ): Schema = enumeration(EnumSet.allOf(enumClass).map { it.name }, description, nullable)
   }
 
   internal fun toInternal(): Internal =
@@ -252,6 +300,7 @@ internal constructor(
       required,
       items?.toInternal(),
     )
+
   @Serializable
   internal data class Internal(
     val type: String,

firebase-ai/src/test/java/com/google/firebase/ai/SchemaTests.kt

@@ -218,4 +218,42 @@ internal class SchemaTests {
 
     Json.encodeToString(schemaDeclaration.toInternal()).shouldEqualJson(expectedJson)
   }
+
+  enum class TestEnum {
+    BASIC,
+    INTERMEDIATE,
+    ADVANCED
+  }
+
+  @Test
+  fun `basic Kotlin enum class`() {
+    val schema = Schema.fromEnum<TestEnum>()
+    val expectedJson =
+      """
+      {
+        "type": "STRING",
+        "format": "enum",
+        "enum": ["BASIC", "INTERMEDIATE", "ADVANCED"]
+      }
+    """
+        .trimIndent()
+
+    Json.encodeToString(schema.toInternal()).shouldEqualJson(expectedJson)
+  }
+
+  @Test
+  fun `basic Java enum`() {
+    val schema = Schema.fromEnum(TestEnum::class.java)
+    val expectedJson =
+      """
+      {
+        "type": "STRING",
+        "format": "enum",
+        "enum": ["BASIC", "INTERMEDIATE", "ADVANCED"]
+      }
+    """
+        .trimIndent()
+
+    Json.encodeToString(schema.toInternal()).shouldEqualJson(expectedJson)
+  }
 }