[goose-llm] kotlin fn for getting structured outputs (#2547)

2025-05-20 07:08:54 -07:00
parent 81332ab914
commit f153204dde
13 changed files with 340 additions and 202 deletions
@@ -135,6 +135,66 @@ fun main() = runBlocking {
    )

    val response = completion(req)
-    println("\nCompletion Response:")
-    println(response.message)
+    println("\nCompletion Response:\n${response.message}")
+    println()
+
+    // ---- UI Extraction (custom schema) ----
+    runUiExtraction(providerName, providerConfig)
+}
+
+
+suspend fun runUiExtraction(providerName: String, providerConfig: String) {
+    val systemPrompt = "You are a UI generator AI. Convert the user input into a JSON-driven UI."
+    val messages = listOf(
+        Message(
+            role = Role.USER,
+            created = System.currentTimeMillis() / 1000,
+            content = listOf(
+                MessageContent.Text(
+                    TextContent("Make a User Profile Form")
+                )
+            )
+        )
+    )
+    val schema = """{
+        "type": "object",
+        "properties": {
+            "type": {
+                "type": "string",
+                "enum": ["div","button","header","section","field","form"]
+            },
+            "label":   { "type": "string" },
+            "children": {
+                "type": "array",
+                "items": { "${'$'}ref": "#" }
+            },
+            "attributes": {
+                "type": "array",
+                "items": {
+                    "type": "object",
+                    "properties": {
+                        "name":  { "type": "string" },
+                        "value": { "type": "string" }
+                    },
+                    "required": ["name","value"],
+                    "additionalProperties": false
+                }
+            }
+        },
+        "required": ["type","label","children","attributes"],
+        "additionalProperties": false
+    }""".trimIndent();
+
+    try {
+        val response = generateStructuredOutputs(
+            providerName = providerName,
+            providerConfig = providerConfig,
+            systemPrompt = systemPrompt,
+            messages = messages,
+            schema = schema
+        )
+        println("\nUI Extraction Output:\n${response}")
+    } catch (e: ProviderException) {
+        println("\nUI Extraction failed:\n${e.message}")
+    }
 }
@@ -775,6 +775,8 @@ internal interface IntegrityCheckingUniffiLib : Library {

    fun uniffi_goose_llm_checksum_func_generate_session_name(): Short

+    fun uniffi_goose_llm_checksum_func_generate_structured_outputs(): Short
+
    fun uniffi_goose_llm_checksum_func_generate_tooltip(): Short

    fun uniffi_goose_llm_checksum_func_print_messages(): Short
@@ -847,6 +849,14 @@ internal interface UniffiLib : Library {
        `messages`: RustBuffer.ByValue,
    ): Long

+    fun uniffi_goose_llm_fn_func_generate_structured_outputs(
+        `providerName`: RustBuffer.ByValue,
+        `providerConfig`: RustBuffer.ByValue,
+        `systemPrompt`: RustBuffer.ByValue,
+        `messages`: RustBuffer.ByValue,
+        `schema`: RustBuffer.ByValue,
+    ): Long
+
    fun uniffi_goose_llm_fn_func_generate_tooltip(
        `providerName`: RustBuffer.ByValue,
        `providerConfig`: RustBuffer.ByValue,
@@ -1090,16 +1100,19 @@ private fun uniffiCheckApiChecksums(lib: IntegrityCheckingUniffiLib) {
    if (lib.uniffi_goose_llm_checksum_func_completion() != 47457.toShort()) {
        throw RuntimeException("UniFFI API checksum mismatch: try cleaning and rebuilding your project")
    }
-    if (lib.uniffi_goose_llm_checksum_func_create_completion_request() != 51008.toShort()) {
+    if (lib.uniffi_goose_llm_checksum_func_create_completion_request() != 39068.toShort()) {
        throw RuntimeException("UniFFI API checksum mismatch: try cleaning and rebuilding your project")
    }
-    if (lib.uniffi_goose_llm_checksum_func_create_tool_config() != 22809.toShort()) {
+    if (lib.uniffi_goose_llm_checksum_func_create_tool_config() != 49910.toShort()) {
        throw RuntimeException("UniFFI API checksum mismatch: try cleaning and rebuilding your project")
    }
-    if (lib.uniffi_goose_llm_checksum_func_generate_session_name() != 9810.toShort()) {
+    if (lib.uniffi_goose_llm_checksum_func_generate_session_name() != 64087.toShort()) {
        throw RuntimeException("UniFFI API checksum mismatch: try cleaning and rebuilding your project")
    }
-    if (lib.uniffi_goose_llm_checksum_func_generate_tooltip() != 15466.toShort()) {
+    if (lib.uniffi_goose_llm_checksum_func_generate_structured_outputs() != 43426.toShort()) {
+        throw RuntimeException("UniFFI API checksum mismatch: try cleaning and rebuilding your project")
+    }
+    if (lib.uniffi_goose_llm_checksum_func_generate_tooltip() != 41121.toShort()) {
        throw RuntimeException("UniFFI API checksum mismatch: try cleaning and rebuilding your project")
    }
    if (lib.uniffi_goose_llm_checksum_func_print_messages() != 30278.toShort()) {
@@ -1626,6 +1639,54 @@ public object FfiConverterTypeProviderCompleteResponse : FfiConverterRustBuffer<
    }
 }

+/**
+ * Response from a structured‐extraction call
+ */
+data class ProviderExtractResponse(
+    /**
+     * The extracted JSON object
+     */
+    var `data`: Value,
+    /**
+     * Which model produced it
+     */
+    var `model`: kotlin.String,
+    /**
+     * Token usage stats
+     */
+    var `usage`: Usage,
+) {
+    companion object
+}
+
+/**
+ * @suppress
+ */
+public object FfiConverterTypeProviderExtractResponse : FfiConverterRustBuffer<ProviderExtractResponse> {
+    override fun read(buf: ByteBuffer): ProviderExtractResponse =
+        ProviderExtractResponse(
+            FfiConverterTypeValue.read(buf),
+            FfiConverterString.read(buf),
+            FfiConverterTypeUsage.read(buf),
+        )
+
+    override fun allocationSize(value: ProviderExtractResponse) =
+        (
+            FfiConverterTypeValue.allocationSize(value.`data`) +
+                FfiConverterString.allocationSize(value.`model`) +
+                FfiConverterTypeUsage.allocationSize(value.`usage`)
+        )
+
+    override fun write(
+        value: ProviderExtractResponse,
+        buf: ByteBuffer,
+    ) {
+        FfiConverterTypeValue.write(value.`data`, buf)
+        FfiConverterString.write(value.`model`, buf)
+        FfiConverterTypeUsage.write(value.`usage`, buf)
+    }
+}
+
 data class RedactedThinkingContent(
    var `data`: kotlin.String,
 ) {
@@ -1750,6 +1811,46 @@ public object FfiConverterTypeThinkingContent : FfiConverterRustBuffer<ThinkingC
    }
 }

+data class ToolConfig(
+    var `name`: kotlin.String,
+    var `description`: kotlin.String,
+    var `inputSchema`: Value,
+    var `approvalMode`: ToolApprovalMode,
+) {
+    companion object
+}
+
+/**
+ * @suppress
+ */
+public object FfiConverterTypeToolConfig : FfiConverterRustBuffer<ToolConfig> {
+    override fun read(buf: ByteBuffer): ToolConfig =
+        ToolConfig(
+            FfiConverterString.read(buf),
+            FfiConverterString.read(buf),
+            FfiConverterTypeValue.read(buf),
+            FfiConverterTypeToolApprovalMode.read(buf),
+        )
+
+    override fun allocationSize(value: ToolConfig) =
+        (
+            FfiConverterString.allocationSize(value.`name`) +
+                FfiConverterString.allocationSize(value.`description`) +
+                FfiConverterTypeValue.allocationSize(value.`inputSchema`) +
+                FfiConverterTypeToolApprovalMode.allocationSize(value.`approvalMode`)
+        )
+
+    override fun write(
+        value: ToolConfig,
+        buf: ByteBuffer,
+    ) {
+        FfiConverterString.write(value.`name`, buf)
+        FfiConverterString.write(value.`description`, buf)
+        FfiConverterTypeValue.write(value.`inputSchema`, buf)
+        FfiConverterTypeToolApprovalMode.write(value.`approvalMode`, buf)
+    }
+}
+
 data class ToolRequest(
    var `id`: kotlin.String,
    var `toolCall`: ToolRequestToolCall,
@@ -2737,34 +2838,6 @@ public object FfiConverterSequenceTypeMessage : FfiConverterRustBuffer<List<Mess
    }
 }

-/**
- * @suppress
- */
-public object FfiConverterSequenceTypeMessageContent : FfiConverterRustBuffer<List<MessageContent>> {
-    override fun read(buf: ByteBuffer): List<MessageContent> {
-        val len = buf.getInt()
-        return List<MessageContent>(len) {
-            FfiConverterTypeMessageContent.read(buf)
-        }
-    }
-
-    override fun allocationSize(value: List<MessageContent>): ULong {
-        val sizeForLength = 4UL
-        val sizeForItems = value.map { FfiConverterTypeMessageContent.allocationSize(it) }.sum()
-        return sizeForLength + sizeForItems
-    }
-
-    override fun write(
-        value: List<MessageContent>,
-        buf: ByteBuffer,
-    ) {
-        buf.putInt(value.size)
-        value.iterator().forEach {
-            FfiConverterTypeMessageContent.write(it, buf)
-        }
-    }
-}
-
 /**
 * @suppress
 */
@@ -2793,6 +2866,34 @@ public object FfiConverterSequenceTypeToolConfig : FfiConverterRustBuffer<List<T
    }
 }

+/**
+ * @suppress
+ */
+public object FfiConverterSequenceTypeMessageContent : FfiConverterRustBuffer<List<MessageContent>> {
+    override fun read(buf: ByteBuffer): List<MessageContent> {
+        val len = buf.getInt()
+        return List<MessageContent>(len) {
+            FfiConverterTypeMessageContent.read(buf)
+        }
+    }
+
+    override fun allocationSize(value: List<MessageContent>): ULong {
+        val sizeForLength = 4UL
+        val sizeForItems = value.map { FfiConverterTypeMessageContent.allocationSize(it) }.sum()
+        return sizeForLength + sizeForItems
+    }
+
+    override fun write(
+        value: List<MessageContent>,
+        buf: ByteBuffer,
+    ) {
+        buf.putInt(value.size)
+        value.iterator().forEach {
+            FfiConverterTypeMessageContent.write(it, buf)
+        }
+    }
+}
+
 /**
 * Typealias from the type name used in the UDL file to the builtin type.  This
 * is needed because the UDL type name is used in function/method signatures.
@@ -2809,22 +2910,6 @@ public typealias FfiConverterTypeCompletionRequest = FfiConverterString
 public typealias Contents = List<MessageContent>
 public typealias FfiConverterTypeContents = FfiConverterSequenceTypeMessageContent

-/**
- * Typealias from the type name used in the UDL file to the builtin type.  This
- * is needed because the UDL type name is used in function/method signatures.
- * It's also what we have an external type that references a custom type.
- */
-public typealias JsonValueFfi = kotlin.String
-public typealias FfiConverterTypeJsonValueFfi = FfiConverterString
-
-/**
- * Typealias from the type name used in the UDL file to the builtin type.  This
- * is needed because the UDL type name is used in function/method signatures.
- * It's also what we have an external type that references a custom type.
- */
-public typealias ToolConfig = kotlin.String
-public typealias FfiConverterTypeToolConfig = FfiConverterString
-
 /**
 * Typealias from the type name used in the UDL file to the builtin type.  This
 * is needed because the UDL type name is used in function/method signatures.
@@ -2841,6 +2926,14 @@ public typealias FfiConverterTypeToolRequestToolCall = FfiConverterString
 public typealias ToolResponseToolResult = kotlin.String
 public typealias FfiConverterTypeToolResponseToolResult = FfiConverterString

+/**
+ * Typealias from the type name used in the UDL file to the builtin type.  This
+ * is needed because the UDL type name is used in function/method signatures.
+ * It's also what we have an external type that references a custom type.
+ */
+public typealias Value = kotlin.String
+public typealias FfiConverterTypeValue = FfiConverterString
+
 /**
 * Public API for the Goose LLM completion function
 */
@@ -2860,7 +2953,7 @@ suspend fun `completion`(`req`: CompletionRequest): CompletionResponse =

 fun `createCompletionRequest`(
    `providerName`: kotlin.String,
-    `providerConfig`: JsonValueFfi,
+    `providerConfig`: Value,
    `modelConfig`: ModelConfig,
    `systemPreamble`: kotlin.String,
    `messages`: List<Message>,
@@ -2870,7 +2963,7 @@ fun `createCompletionRequest`(
        uniffiRustCall { _status ->
            UniffiLib.INSTANCE.uniffi_goose_llm_fn_func_create_completion_request(
                FfiConverterString.lower(`providerName`),
-                FfiConverterTypeJsonValueFfi.lower(`providerConfig`),
+                FfiConverterTypeValue.lower(`providerConfig`),
                FfiConverterTypeModelConfig.lower(`modelConfig`),
                FfiConverterString.lower(`systemPreamble`),
                FfiConverterSequenceTypeMessage.lower(`messages`),
@@ -2883,7 +2976,7 @@ fun `createCompletionRequest`(
 fun `createToolConfig`(
    `name`: kotlin.String,
    `description`: kotlin.String,
-    `inputSchema`: JsonValueFfi,
+    `inputSchema`: Value,
    `approvalMode`: ToolApprovalMode,
 ): ToolConfig =
    FfiConverterTypeToolConfig.lift(
@@ -2891,7 +2984,7 @@ fun `createToolConfig`(
            UniffiLib.INSTANCE.uniffi_goose_llm_fn_func_create_tool_config(
                FfiConverterString.lower(`name`),
                FfiConverterString.lower(`description`),
-                FfiConverterTypeJsonValueFfi.lower(`inputSchema`),
+                FfiConverterTypeValue.lower(`inputSchema`),
                FfiConverterTypeToolApprovalMode.lower(`approvalMode`),
                _status,
            )
@@ -2905,13 +2998,13 @@ fun `createToolConfig`(
@Suppress("ASSIGNED_BUT_NEVER_ACCESSED_VARIABLE")
 suspend fun `generateSessionName`(
    `providerName`: kotlin.String,
-    `providerConfig`: JsonValueFfi,
+    `providerConfig`: Value,
    `messages`: List<Message>,
 ): kotlin.String =
    uniffiRustCallAsync(
        UniffiLib.INSTANCE.uniffi_goose_llm_fn_func_generate_session_name(
            FfiConverterString.lower(`providerName`),
-            FfiConverterTypeJsonValueFfi.lower(`providerConfig`),
+            FfiConverterTypeValue.lower(`providerConfig`),
            FfiConverterSequenceTypeMessage.lower(`messages`),
        ),
        { future, callback, continuation -> UniffiLib.INSTANCE.ffi_goose_llm_rust_future_poll_rust_buffer(future, callback, continuation) },
@@ -2923,6 +3016,36 @@ suspend fun `generateSessionName`(
        ProviderException.ErrorHandler,
    )

+/**
+ * Generates a structured output based on the provided schema,
+ * system prompt and user messages.
+ */
+@Throws(ProviderException::class)
+@Suppress("ASSIGNED_BUT_NEVER_ACCESSED_VARIABLE")
+suspend fun `generateStructuredOutputs`(
+    `providerName`: kotlin.String,
+    `providerConfig`: Value,
+    `systemPrompt`: kotlin.String,
+    `messages`: List<Message>,
+    `schema`: Value,
+): ProviderExtractResponse =
+    uniffiRustCallAsync(
+        UniffiLib.INSTANCE.uniffi_goose_llm_fn_func_generate_structured_outputs(
+            FfiConverterString.lower(`providerName`),
+            FfiConverterTypeValue.lower(`providerConfig`),
+            FfiConverterString.lower(`systemPrompt`),
+            FfiConverterSequenceTypeMessage.lower(`messages`),
+            FfiConverterTypeValue.lower(`schema`),
+        ),
+        { future, callback, continuation -> UniffiLib.INSTANCE.ffi_goose_llm_rust_future_poll_rust_buffer(future, callback, continuation) },
+        { future, continuation -> UniffiLib.INSTANCE.ffi_goose_llm_rust_future_complete_rust_buffer(future, continuation) },
+        { future -> UniffiLib.INSTANCE.ffi_goose_llm_rust_future_free_rust_buffer(future) },
+        // lift function
+        { FfiConverterTypeProviderExtractResponse.lift(it) },
+        // Error FFI converter
+        ProviderException.ErrorHandler,
+    )
+
 /**
 * Generates a tooltip summarizing the last two messages in the session,
 * including any tool calls or results.
@@ -2931,13 +3054,13 @@ suspend fun `generateSessionName`(
@Suppress("ASSIGNED_BUT_NEVER_ACCESSED_VARIABLE")
 suspend fun `generateTooltip`(
    `providerName`: kotlin.String,
-    `providerConfig`: JsonValueFfi,
+    `providerConfig`: Value,
    `messages`: List<Message>,
 ): kotlin.String =
    uniffiRustCallAsync(
        UniffiLib.INSTANCE.uniffi_goose_llm_fn_func_generate_tooltip(
            FfiConverterString.lower(`providerName`),
-            FfiConverterTypeJsonValueFfi.lower(`providerConfig`),
+            FfiConverterTypeValue.lower(`providerConfig`),
            FfiConverterSequenceTypeMessage.lower(`messages`),
        ),
        { future, callback, continuation -> UniffiLib.INSTANCE.ffi_goose_llm_rust_future_poll_rust_buffer(future, callback, continuation) },