fix: schema handling in transformer

PiperOrigin-RevId: 745790224

Author: yyyu-google

examples/src/main/java/com/google/genai/examples/GenerateContentWithFunctionCall.java

@@ -39,6 +39,7 @@
  */
 package com.google.genai.examples;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.genai.Client;
 import com.google.genai.types.FunctionDeclaration;
@@ -64,7 +65,13 @@ public static void main(String[] args) {
                             Schema.builder()
                                 .type("string")
                                 .description("The location to get the weather for.")
+                                .build(),
+                            "unit",
+                            Schema.builder()
+                                .type("string")
+                                .description("The unit to return the weather in, e.g. 'celsius'.")
                                 .build()))
+                    .required(ImmutableList.of("location", "unit"))
                     .build())
             .build();
 

examples/src/main/java/com/google/genai/examples/GenerateContentWithResponseSchema.java

@@ -40,6 +40,7 @@
 package com.google.genai.examples;
 
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableList;
 import com.google.genai.Client;
 import com.google.genai.types.GenerateContentConfig;
 import com.google.genai.types.GenerateContentResponse;
@@ -53,12 +54,23 @@ public class GenerateContentWithResponseSchema {
   public static void main(String[] args) {
     Client client = new Client();
 
-    Schema schema =
-        Schema.builder()
-            .type("object")
-            .properties(
-                ImmutableMap.of(
-                    "name", Schema.builder().type("string").description("Your Name").build()))
+    Schema schema = Schema.builder()
+            .type("ARRAY")
+            .items(
+                Schema.builder()
+                    .type("OBJECT")
+                    .properties(
+                        ImmutableMap.of(
+                            "recipe_name",
+                            Schema.builder().type("STRING").build(),
+                            "ingredients",
+                            Schema.builder()
+                                .type("ARRAY")
+                                .items(Schema.builder().type("STRING").build())
+                                .build()
+                            ))
+                    .required(ImmutableList.of("recipe_name", "ingredients"))
+                    .build())
             .build();
     GenerateContentConfig config =
         GenerateContentConfig.builder()
@@ -68,7 +80,7 @@ public static void main(String[] args) {
             .build();
 
     GenerateContentResponse response =
-        client.models.generateContent("gemini-2.0-flash-001", "Tell me your name", config);
+        client.models.generateContent("gemini-2.0-flash-001", "List a few popular cookie recipes.", config);
 
     System.out.println("Response: " + response.text());
   }

src/main/java/com/google/genai/Transformers.java

@@ -28,6 +28,7 @@
 import com.google.genai.types.Tool;
 import com.google.genai.types.VoiceConfig;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import org.jspecify.annotations.Nullable;
@@ -263,20 +264,18 @@ private static Schema processSchema(ApiClient apiClient, Schema schema) {
     }
 
     if (schema.items().isPresent()) {
-      schema = schema.toBuilder().items(processSchema(apiClient, schema.items().get())).build();
+      Schema items = processSchema(apiClient, schema.items().get());
+      schema = schema.toBuilder().items(items).build();
     }
 
     if (schema.properties().isPresent()) {
+      Map<String, Schema> properties = new HashMap<>();
       for (Map.Entry<String, Schema> entry : schema.properties().get().entrySet()) {
         String propertyName = entry.getKey();
         Schema propertySchema = entry.getValue();
-        schema =
-            schema.toBuilder()
-                .properties(
-                    ImmutableMap.of(
-                        propertyName, processSchema(apiClient, propertySchema).toBuilder().build()))
-                .build();
+        properties.put(propertyName, processSchema(apiClient, propertySchema));
       }
+      schema = schema.toBuilder().properties(ImmutableMap.copyOf(properties)).build();
     }
 
     return schema;

src/test/java/com/google/genai/TransformerTest.java

@@ -88,4 +88,28 @@ public void testTSchema_Items_success() {
     assertEquals("STRING", transformedSchema.items().get().type().get());
     assertEquals("ARRAY", transformedSchema.type().get());
   }
+
+  @Test
+  public void testTSchema_Required_sucess() {
+    Schema schema = Schema.builder()
+            .type("ARRAY")
+            .items(
+                Schema.builder()
+                    .type("OBJECT")
+                    .properties(
+                        ImmutableMap.of(
+                            "recipe_name",
+                            Schema.builder().type("STRING").build(),
+                            "ingredients",
+                            Schema.builder()
+                                .type("ARRAY")
+                                .items(Schema.builder().type("STRING").build())
+                                .build()
+                            ))
+                    .required(ImmutableList.of("recipe_name", "ingredients"))
+                    .build())
+            .build();
+    Schema transformedSchema = Transformers.tSchema(GEMINI_API_CLIENT, schema);
+    assertEquals(schema, transformedSchema);
+  }
 }