feat(kotlin): add maxOutputToken API to ConversationConfig and sendMessage

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

@@ -91,7 +91,6 @@ fun benchmark(
         null, // loraPath
         null, // audioLoraPath
         false, // prefillPrefaceOnInit
-        -1, // maxOutputToken
       )
 
     Conversation(conversationHandle).use { conversation ->

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

@@ -120,8 +120,6 @@ data class EngineConfig(
  * @property prefillPrefaceOnInit Whether to prefill the preface on initialization. Defaults to
  *   false. Note that this will make createConversation() take longer to finish, so you may want to
  *   call it in a background thread.
- * @property maxOutputToken The maximum number of output tokens per decode step. When `null`, use
- *   the default value from the model or the engine.
  */
 data class ConversationConfig
 @JvmOverloads
@@ -135,14 +133,7 @@ constructor(
   val extraContext: Map<String, Any> = emptyMap(),
   val loraConfig: LoraConfig? = null,
   val prefillPrefaceOnInit: Boolean = false,
-  val maxOutputToken: Int? = null,
-) {
-  init {
-    require(maxOutputToken == null || maxOutputToken > 0) {
-      "maxOutputToken must be positive or null (use the default from model or engine)."
-    }
-  }
-}
+)
 
 /**
  * Configuration for the sampling process.

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

@@ -87,17 +87,12 @@ class Conversation(
    *
    * @param message The message to send to the model.
    * @param extraContext Optional context used for prompt template rendering.
-   * @param maxOutputToken Optional override for the maximum number of output tokens per decode step.
    * @return The model's response message.
    * @throws IllegalStateException if the conversation is not alive, if the native layer returns an
    *   invalid response, or if the tool call limit is exceeded.
    * @throws LiteRtLmJniException if an error occurs during the native call.
    */
-  fun sendMessage(
-    message: Message,
-    extraContext: Map<String, Any> = emptyMap(),
-    maxOutputToken: Int? = null,
-  ): Message {
+  fun sendMessage(message: Message, extraContext: Map<String, Any> = emptyMap()): Message {
     checkIsAlive()
 
     var currentMessageJson = message.toJson()
@@ -111,7 +106,6 @@ class Conversation(
           currentMessageJson.toString(),
           extraContextJsonString,
           visualTokenBudget,
-          maxOutputToken ?: -1,
         )
       val responseJsonObject = JsonParser.parseString(responseJsonString).asJsonObject
 
@@ -140,18 +134,13 @@ class Conversation(
    *
    * @param contents The list of contents to send to the model.
    * @param extraContext Optional context used for prompt template rendering.
-   * @param maxOutputToken Optional override for the maximum number of output tokens per decode step.
    * @return The model's response message.
    * @throws IllegalStateException if the conversation is not alive, if the native layer returns an
    *   invalid response, or if the tool call limit is exceeded.
    * @throws LiteRtLmJniException if an error occurs during the native call.
    */
-  fun sendMessage(
-    contents: Contents,
-    extraContext: Map<String, Any> = emptyMap(),
-    maxOutputToken: Int? = null,
-  ): Message {
-    return sendMessage(Message.user(contents), extraContext, maxOutputToken)
+  fun sendMessage(contents: Contents, extraContext: Map<String, Any> = emptyMap()): Message {
+    return sendMessage(Message.user(contents), extraContext)
   }
 
   /**
@@ -164,17 +153,13 @@ class Conversation(
    *
    * @param text The text to send to the model.
    * @param extraContext Optional context used for prompt template rendering.
-   * @param maxOutputToken Optional override for the maximum number of output tokens per decode step.
    * @return The model's response message.
    * @throws IllegalStateException if the conversation is not alive, if the native layer returns an
    *   invalid response, or if the tool call limit is exceeded.
    * @throws LiteRtLmJniException if an error occurs during the native call.
    */
-  fun sendMessage(
-    text: String,
-    extraContext: Map<String, Any> = emptyMap(),
-    maxOutputToken: Int? = null,
-  ): Message = sendMessage(Contents.of(text), extraContext, maxOutputToken)
+  fun sendMessage(text: String, extraContext: Map<String, Any> = emptyMap()): Message =
+    sendMessage(Contents.of(text), extraContext)
 
   /**
    * Send a message to the model and returns the response async with a callback.
@@ -187,29 +172,26 @@ class Conversation(
    * @param message The message to send to the model.
    * @param callback The callback to receive the streaming responses.
    * @param extraContext Optional context used for prompt template rendering.
-   * @param maxOutputToken Optional override for the maximum number of output tokens per decode step.
    * @throws IllegalStateException if the conversation has already been closed or the content is
    *   empty.
    */
   fun sendMessageAsync(
     message: Message,
     callback: MessageCallback,
     extraContext: Map<String, Any> = emptyMap(),
-    maxOutputToken: Int? = null,
   ) {
     checkIsAlive()
 
     val extraContextJsonString = extraContext.toJsonObject().toString()
     val visualTokenBudget = @OptIn(ExperimentalApi::class) ExperimentalFlags.visualTokenBudget
 
-    val jniCallback = JniMessageCallbackImpl(callback, maxOutputToken)
+    val jniCallback = JniMessageCallbackImpl(callback)
     LiteRtLmJni.nativeSendMessageAsync(
       handle,
       message.toJson().toString(),
       extraContextJsonString,
       jniCallback,
       visualTokenBudget,
-      maxOutputToken ?: -1,
     )
   }
 
@@ -224,16 +206,14 @@ class Conversation(
    * @param contents The list of contents to send to the model.
    * @param callback The callback to receive the streaming responses.
    * @param extraContext Optional context used for prompt template rendering.
-   * @param maxOutputToken Optional override for the maximum number of output tokens per decode step.
    * @throws IllegalStateException if the conversation has already been closed or the content is
    *   empty.
    */
   fun sendMessageAsync(
     contents: Contents,
     callback: MessageCallback,
     extraContext: Map<String, Any> = emptyMap(),
-    maxOutputToken: Int? = null,
-  ) = sendMessageAsync(Message.user(contents), callback, extraContext, maxOutputToken)
+  ) = sendMessageAsync(Message.user(contents), callback, extraContext)
 
   /**
    * Send a text to the model and returns the response async with a callback.
@@ -246,16 +226,14 @@ class Conversation(
    * @param text The text to send to the model.
    * @param callback The callback to receive the streaming responses.
    * @param extraContext Optional context used for prompt template rendering.
-   * @param maxOutputToken Optional override for the maximum number of output tokens per decode step.
    * @throws IllegalStateException if the conversation has already been closed or the content is
    *   empty.
    */
   fun sendMessageAsync(
     text: String,
     callback: MessageCallback,
     extraContext: Map<String, Any> = emptyMap(),
-    maxOutputToken: Int? = null,
-  ) = sendMessageAsync(Contents.of(text), callback, extraContext, maxOutputToken)
+  ) = sendMessageAsync(Contents.of(text), callback, extraContext)
 
   /**
    * Sends a message to the model and returns the response async as a [Flow].
@@ -267,15 +245,13 @@ class Conversation(
    *
    * @param message The message to send to the model.
    * @param extraContext Optional context used for prompt template rendering.
-   * @param maxOutputToken Optional override for the maximum number of output tokens per decode step.
    * @return A Flow of messages representing the model's response.
    * @throws IllegalStateException if the conversation has already been closed or the content is
    *   empty.
    */
   fun sendMessageAsync(
     message: Message,
     extraContext: Map<String, Any> = emptyMap(),
-    maxOutputToken: Int? = null,
   ): Flow<Message> = callbackFlow {
     sendMessageAsync(
       message,
@@ -293,7 +269,6 @@ class Conversation(
         }
       },
       extraContext,
-      maxOutputToken,
     )
     awaitClose {}
   }
@@ -308,16 +283,14 @@ class Conversation(
    *
    * @param contents The list of contents to send to the model.
    * @param extraContext Optional context used for prompt template rendering.
-   * @param maxOutputToken Optional override for the maximum number of output tokens per decode step.
    * @return A Flow of messages representing the model's response.
    * @throws IllegalStateException if the conversation has already been closed or the content is
    *   empty.
    */
   fun sendMessageAsync(
     contents: Contents,
     extraContext: Map<String, Any> = emptyMap(),
-    maxOutputToken: Int? = null,
-  ): Flow<Message> = sendMessageAsync(Message.user(contents), extraContext, maxOutputToken)
+  ): Flow<Message> = sendMessageAsync(Message.user(contents), extraContext)
 
   /**
    * Sends a text to the model and returns the response async as a [Flow].
@@ -329,16 +302,12 @@ class Conversation(
    *
    * @param text The text to send to the model.
    * @param extraContext Optional context used for prompt template rendering.
-   * @param maxOutputToken Optional override for the maximum number of output tokens per decode step.
    * @return A Flow of messages representing the model's response.
    * @throws IllegalStateException if the conversation has already been closed or the content is
    *   empty.
    */
-  fun sendMessageAsync(
-    text: String,
-    extraContext: Map<String, Any> = emptyMap(),
-    maxOutputToken: Int? = null,
-  ): Flow<Message> = sendMessageAsync(Contents.of(text), extraContext, maxOutputToken)
+  fun sendMessageAsync(text: String, extraContext: Map<String, Any> = emptyMap()): Flow<Message> =
+    sendMessageAsync(Contents.of(text), extraContext)
 
   private fun handleToolCalls(toolCallsJsonObject: JsonObject): JsonObject {
     val toolCallsJSONArray = toolCallsJsonObject.getAsJsonArray("tool_calls")
@@ -368,10 +337,8 @@ class Conversation(
     }
   }
 
-  private inner class JniMessageCallbackImpl(
-    private val callback: MessageCallback,
-    private val maxOutputToken: Int? = null,
-  ) : LiteRtLmJni.JniMessageCallback {
+  private inner class JniMessageCallbackImpl(private val callback: MessageCallback) :
+    LiteRtLmJni.JniMessageCallback {
 
     /** The tool response to be returned back */
     private var pendingToolResponseJSONMessage: JsonObject? = null
@@ -410,7 +377,6 @@ class Conversation(
           "{}",
           this@JniMessageCallbackImpl,
           @OptIn(ExperimentalApi::class) ExperimentalFlags.visualTokenBudget,
-          maxOutputToken ?: -1,
         )
         pendingToolResponseJSONMessage = null // Clear after sending
       } else {

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

@@ -159,7 +159,6 @@ class Engine(val engineConfig: EngineConfig) : AutoCloseable {
           conversationConfig.loraConfig?.loraPath,
           conversationConfig.loraConfig?.audioLoraPath,
           conversationConfig.prefillPrefaceOnInit,
-          conversationConfig.maxOutputToken ?: -1,
         ),
         toolManager,
         conversationConfig.automaticToolCalling,

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

@@ -206,7 +206,6 @@ internal object LiteRtLmJni {
    *   decoding.
    * @param filterChannelContentFromKvCache Whether to filter channel content from the KV cache.
    * @param prefillPrefaceOnInit Whether to prefill the preface when initializing the conversation.
-   * @param maxOutputToken The maximum number of output tokens. When non-positive, use the default.
    * @return A pointer to the native conversation instance.
    */
   external fun nativeCreateConversation(
@@ -222,7 +221,6 @@ internal object LiteRtLmJni {
     loraPath: String?,
     audioLoraPath: String?,
     prefillPrefaceOnInit: Boolean,
-    maxOutputToken: Int,
   ): Long
 
   /**
@@ -244,15 +242,13 @@ internal object LiteRtLmJni {
    * @param callback The callback to receive the streaming responses.
    * @param visualTokenBudget The visual token budget. Only supported by Gemma4 currently. Null for
    *   default.
-   * @param maxOutputToken The maximum number of output tokens. When non-positive, use the default.
    */
   external fun nativeSendMessageAsync(
     conversationPointer: Long,
     messageJsonString: String,
     extraContextJsonString: String,
     callback: JniMessageCallback,
     visualTokenBudget: Int?,
-    maxOutputToken: Int,
   )
 
   /**
@@ -264,15 +260,13 @@ internal object LiteRtLmJni {
    *   format.
    * @param visualTokenBudget The visual token budget. Only supported by Gemma4 currently. Null for
    *   default.
-   * @param maxOutputToken The maximum number of output tokens. When non-positive, use the default.
    * @return The response message in JSON string format.
    */
   external fun nativeSendMessage(
     conversationPointer: Long,
     messageJsonString: String,
     extraContextJsonString: String,
     visualTokenBudget: Int?,
-    maxOutputToken: Int,
   ): String
 
   /**

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

@@ -898,15 +898,11 @@ LITERTLM_JNIEXPORT jlong JNICALL JNI_METHOD(nativeCreateConversation)(
     jboolean enable_constrained_decoding,
     jboolean filter_channel_content_from_kv_cache,
     jstring overwrite_prompt_template, jstring lora_path_str,
-    jstring audio_lora_path_str, jboolean prefill_preface_on_init,
-    jint max_output_token) {
+    jstring audio_lora_path_str, jboolean prefill_preface_on_init) {
   Engine* engine = reinterpret_cast<Engine*>(engine_pointer);
 
   // Create a native SessionConfig
   auto session_config = SessionConfig::CreateDefault();
-  if (max_output_token > 0) {
-    session_config.SetMaxOutputTokens(max_output_token);
-  }
   if (sampler_config_obj != nullptr) {
     session_config.GetMutableSamplerParams() =
         CreateSamplerParamsFromJni(env, sampler_config_obj);
@@ -1048,7 +1044,7 @@ LITERTLM_JNIEXPORT void JNICALL JNI_METHOD(nativeDeleteConversation)(
 LITERTLM_JNIEXPORT void JNICALL JNI_METHOD(nativeSendMessageAsync)(
     JNIEnv* env, jclass thiz, jlong conversation_pointer,
     jstring messageJSONString, jstring extraContextJsonString, jobject callback,
-    jobject visual_token_budget, jint max_output_token) {
+    jobject visual_token_budget) {
   JavaVM* jvm = nullptr;
   if (env->GetJavaVM(&jvm) != JNI_OK) {
     ThrowLiteRtLmJniException(env, "Failed to get JavaVM");
@@ -1063,9 +1059,6 @@ LITERTLM_JNIEXPORT void JNICALL JNI_METHOD(nativeSendMessageAsync)(
   env->ReleaseStringUTFChars(messageJSONString, json_chars);
 
   litert::lm::OptionalArgs optional_args;
-  if (max_output_token > 0) {
-    optional_args.max_output_tokens = max_output_token;
-  }
   nlohmann::ordered_json extra_context =
       GetExtraContextJson(env, extraContextJsonString);
   if (!extra_context.is_null() && !extra_context.empty()) {
@@ -1160,7 +1153,7 @@ LITERTLM_JNIEXPORT void JNICALL JNI_METHOD(nativeSendMessageAsync)(
 LITERTLM_JNIEXPORT jstring JNICALL JNI_METHOD(nativeSendMessage)(
     JNIEnv* env, jclass thiz, jlong conversation_pointer,
     jstring messageJSONString, jstring extraContextJsonString,
-    jobject visual_token_budget, jint max_output_token) {
+    jobject visual_token_budget) {
   Conversation* conversation =
       reinterpret_cast<Conversation*>(conversation_pointer);
 
@@ -1169,9 +1162,6 @@ LITERTLM_JNIEXPORT jstring JNICALL JNI_METHOD(nativeSendMessage)(
   env->ReleaseStringUTFChars(messageJSONString, json_chars);
 
   litert::lm::OptionalArgs optional_args;
-  if (max_output_token > 0) {
-    optional_args.max_output_tokens = max_output_token;
-  }
   nlohmann::ordered_json extra_context =
       GetExtraContextJson(env, extraContextJsonString);
   if (!extra_context.is_null() && !extra_context.empty()) {