Add support for configuring model response channels in LiteRT-LM Kotlin API ConversationConfig.

LiteRT-LM-PiperOrigin-RevId: 888850632

Author: matthewchan-g

kotlin/java/com/google/ai/edge/litertlm/Benchmark.kt

@@ -78,6 +78,7 @@ fun benchmark(
         null, // SamplerConfig
         "[]", // messagesJsonString
         "[]", // toolsDescriptionJsonString
+        null, // channelsJsonString
         false, // enableConversationConstrainedDecoding
       )
 

kotlin/java/com/google/ai/edge/litertlm/Config.kt

@@ -15,6 +15,25 @@
  */
 package com.google.ai.edge.litertlm
 
+import com.google.gson.JsonObject
+
+/**
+ * Definition of a channel for responses, e.g. thinking channel.
+ *
+ * @property channelName The channel name. Text from this channel will be written to
+ *   [Message.channels] with the [channelName] as the key.
+ * @property start A string that marks the start of the channel.
+ * @property end A string that marks the end of the channel.
+ */
+data class Channel(val channelName: String, val start: String, val end: String) {
+  internal fun toJson() =
+    JsonObject().apply {
+      addProperty("channel_name", channelName)
+      addProperty("start", start)
+      addProperty("end", end)
+    }
+}
+
 /**
  * Backend for the LiteRT-LM engine.
  *
@@ -96,6 +115,7 @@ data class ConversationConfig(
   val tools: List<ToolProvider> = listOf(),
   val samplerConfig: SamplerConfig? = null,
   val automaticToolCalling: Boolean = true,
+  val channels: List<Channel>? = null,
 )
 
 /**

kotlin/java/com/google/ai/edge/litertlm/Engine.kt

@@ -126,13 +126,27 @@ class Engine(val engineConfig: EngineConfig) : AutoCloseable {
           }
         }
 
+      // Convert the channels to a JSON array, if provided.
+      // If `channels` is null, the `Conversation` uses the default from the
+      // `LlmMetadata` or the model type.
+      // If channels is empty, channels will be disabled.
+      val channelsJson: JsonArray? =
+        conversationConfig.channels?.let { channels ->
+          JsonArray().apply {
+            for (channel in channels) {
+              this.add(channel.toJson())
+            }
+          }
+        }
+
       @OptIn(ExperimentalApi::class) // opt-in experimental flags
       return Conversation(
         LiteRtLmJni.nativeCreateConversation(
           handle!!, // Using !! is okay. Checked initialization already.
           conversationConfig.samplerConfig,
           messagesJson.toString(),
           toolManager.getToolsDescription().toString(),
+          channelsJson?.toString(),
           ExperimentalFlags.enableConversationConstrainedDecoding,
         ),
         toolManager,

kotlin/java/com/google/ai/edge/litertlm/LiteRtLmJni.kt

@@ -186,6 +186,8 @@ internal object LiteRtLmJni {
    * @param systemMessageJsonString The system instruction to be used in the conversation.
    * @param toolsDescriptionJsonString A json string of a list of tool definitions (Open API json).
    *   could be used.
+   * @param channelsJsonString A json string of a list of channel definitions. If null, use the
+   *   default from the model or engine. If empty, channels will be disabled.
    * @param enableConversationConstrainedDecoding Whether to enable conversation constrained
    *   decoding.
    * @return A pointer to the native conversation instance.
@@ -195,6 +197,7 @@ internal object LiteRtLmJni {
     samplerConfig: SamplerConfig?,
     messageJsonString: String,
     toolsDescriptionJsonString: String,
+    channelsJsonString: String?,
     enableConversationConstrainedDecoding: Boolean,
   ): Long
 

kotlin/java/com/google/ai/edge/litertlm/jni/litertlm.cc

@@ -829,7 +829,7 @@ JNI_METHOD(nativeConversationGetBenchmarkInfo)(JNIEnv* env, jclass thiz,
 LITERTLM_JNIEXPORT jlong JNICALL JNI_METHOD(nativeCreateConversation)(
     JNIEnv* env, jclass thiz, jlong engine_pointer, jobject sampler_config_obj,
     jstring messages_json_string, jstring tools_description_json_string,
-    jboolean enable_constrained_decoding) {
+    jstring channels_json_string, jboolean enable_constrained_decoding) {
   Engine* engine = reinterpret_cast<Engine*>(engine_pointer);
 
   // Create a native SessionConfig
@@ -864,13 +864,37 @@ LITERTLM_JNIEXPORT jlong JNICALL JNI_METHOD(nativeCreateConversation)(
     return 0;
   }
 
-  // Create the conversation
-  auto conversation_config =
+  // Create a ConversationConfig::Builder
+  auto conversation_config_builder =
       ConversationConfig::Builder()
           .SetSessionConfig(session_config)
           .SetPreface(json_preface)
-          .SetEnableConstrainedDecoding(enable_constrained_decoding)
-          .Build(*engine);
+          .SetEnableConstrainedDecoding(enable_constrained_decoding);
+
+  // Set the channels, if provided.
+  // If channels is nullptr, the Conversation will use the channels defined in
+  // the LlmMetadata or the default channels for the model type.
+  // If channels is an empty array, channels will be disabled.
+  if (channels_json_string != nullptr) {
+    const char* channels_chars =
+        env->GetStringUTFChars(channels_json_string, nullptr);
+    std::string channels_json_str(channels_chars);
+    env->ReleaseStringUTFChars(channels_json_string, channels_chars);
+    auto channels_json = nlohmann::ordered_json::parse(channels_json_str);
+
+    std::vector<litert::lm::Channel> channels;
+    if (channels_json.is_array()) {
+      for (const auto& channel_item : channels_json) {
+        channels.push_back({channel_item["channel_name"].get<std::string>(),
+                            channel_item["start"].get<std::string>(),
+                            channel_item["end"].get<std::string>()});
+      }
+    }
+    conversation_config_builder.SetChannels(channels);
+  }
+
+  // Build the conversation
+  auto conversation_config = conversation_config_builder.Build(*engine);
 
   if (!conversation_config.ok()) {
     ThrowLiteRtLmJniException(env, "Failed to create conversation config: " +