From 37a3eb5b1972ce25f815070b6d6de78a926a9894 Mon Sep 17 00:00:00 2001
From: Rodrigo Lazo Paz <rlazo@google.com>
Date: Mon, 15 Jul 2024 19:38:23 -0400
Subject: [PATCH] Controlled generation in Kotlin

---
 .../samples/controlled_generation.kt          | 59 +++++++++++++++++++
 1 file changed, 59 insertions(+)

diff --git a/samples/src/main/java/com/google/ai/client/generative/samples/controlled_generation.kt b/samples/src/main/java/com/google/ai/client/generative/samples/controlled_generation.kt
index c5fceb55..011c1370 100644
--- a/samples/src/main/java/com/google/ai/client/generative/samples/controlled_generation.kt
+++ b/samples/src/main/java/com/google/ai/client/generative/samples/controlled_generation.kt
@@ -16,6 +16,11 @@
 
 package com.google.ai.client.generative.samples
 
+import com.google.ai.client.generativeai.GenerativeModel
+import com.google.ai.client.generativeai.type.FunctionType
+import com.google.ai.client.generativeai.type.Schema
+import com.google.ai.client.generativeai.type.generationConfig
+
 // Set up your API Key
 // ====================
 //
@@ -25,10 +30,64 @@ package com.google.ai.client.generative.samples
 
 suspend fun json_controlled_generation() {
     // [START json_controlled_generation]
+    val jsonSchema = Schema(
+        name = "recipes",
+        description = "List of recipes",
+        type = FunctionType.ARRAY,
+        items = Schema(
+            name = "recipe",
+            description = "A recipe",
+            type = FunctionType.OBJECT,
+            properties = mapOf(
+                "recipeName" to Schema(
+                    name = "recipeName",
+                    description = "Name of the recipe",
+                    type = FunctionType.STRING,
+                    nullable = false
+                ),
+            ),
+            required = listOf("recipeName")
+        ),
+    )
+
+    val generativeModel =
+        GenerativeModel(
+            // Specify a Gemini model appropriate for your use case
+            modelName = "gemini-1.5-pro",
+            // Access your API key as a Build Configuration variable (see "Set up your API key" above)
+            apiKey = BuildConfig.apiKey,
+            generationConfig = generationConfig {
+                responseMimeType = "application/json"
+                responseSchema = jsonSchema
+            })
+
+    val prompt = "List a few popular cookie recipes."
+    val response = generativeModel.generateContent(prompt)
+    print(response.text)
+
     // [END json_controlled_generation]
 }
 
 suspend fun json_no_schema() {
     // [START json_no_schema]
+
+    val generativeModel =
+        GenerativeModel(
+            // Specify a Gemini model appropriate for your use case
+            modelName = "gemini-1.5-flash",
+            // Access your API key as a Build Configuration variable (see "Set up your API key" above)
+            apiKey = BuildConfig.apiKey,
+            generationConfig = generationConfig {
+                responseMimeType = "application/json"
+            })
+
+    val prompt = """
+           List a few popular cookie recipes using this JSON schema:
+           Recipe = {'recipeName': string}
+           Return: Array<Recipe>
+    """.trimIndent()
+    val response = generativeModel.generateContent(prompt)
+    print(response.text)
+
     // [END json_no_schema]
 }