diff --git a/kotlin/java/com/google/ai/edge/litertlm/Benchmark.kt b/kotlin/java/com/google/ai/edge/litertlm/Benchmark.kt index 2f60484bd..4142c2324 100644 --- a/kotlin/java/com/google/ai/edge/litertlm/Benchmark.kt +++ b/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 -> diff --git a/kotlin/java/com/google/ai/edge/litertlm/Config.kt b/kotlin/java/com/google/ai/edge/litertlm/Config.kt index f96bbcae0..683fca144 100644 --- a/kotlin/java/com/google/ai/edge/litertlm/Config.kt +++ b/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 = 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. diff --git a/kotlin/java/com/google/ai/edge/litertlm/Conversation.kt b/kotlin/java/com/google/ai/edge/litertlm/Conversation.kt index 8f92088e9..f23bafb56 100644 --- a/kotlin/java/com/google/ai/edge/litertlm/Conversation.kt +++ b/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 = emptyMap(), - maxOutputToken: Int? = null, - ): Message { + fun sendMessage(message: Message, extraContext: Map = 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 = emptyMap(), - maxOutputToken: Int? = null, - ): Message { - return sendMessage(Message.user(contents), extraContext, maxOutputToken) + fun sendMessage(contents: Contents, extraContext: Map = 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 = emptyMap(), - maxOutputToken: Int? = null, - ): Message = sendMessage(Contents.of(text), extraContext, maxOutputToken) + fun sendMessage(text: String, extraContext: Map = emptyMap()): Message = + sendMessage(Contents.of(text), extraContext) /** * Send a message to the model and returns the response async with a callback. @@ -187,7 +172,6 @@ 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. */ @@ -195,21 +179,19 @@ class Conversation( message: Message, callback: MessageCallback, extraContext: Map = 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,7 +206,6 @@ 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. */ @@ -232,8 +213,7 @@ class Conversation( contents: Contents, callback: MessageCallback, extraContext: Map = 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,7 +226,6 @@ 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. */ @@ -254,8 +233,7 @@ class Conversation( text: String, callback: MessageCallback, extraContext: Map = 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,7 +245,6 @@ 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. @@ -275,7 +252,6 @@ class Conversation( fun sendMessageAsync( message: Message, extraContext: Map = emptyMap(), - maxOutputToken: Int? = null, ): Flow = callbackFlow { sendMessageAsync( message, @@ -293,7 +269,6 @@ class Conversation( } }, extraContext, - maxOutputToken, ) awaitClose {} } @@ -308,7 +283,6 @@ 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. @@ -316,8 +290,7 @@ class Conversation( fun sendMessageAsync( contents: Contents, extraContext: Map = emptyMap(), - maxOutputToken: Int? = null, - ): Flow = sendMessageAsync(Message.user(contents), extraContext, maxOutputToken) + ): Flow = 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 = emptyMap(), - maxOutputToken: Int? = null, - ): Flow = sendMessageAsync(Contents.of(text), extraContext, maxOutputToken) + fun sendMessageAsync(text: String, extraContext: Map = emptyMap()): Flow = + 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 { diff --git a/kotlin/java/com/google/ai/edge/litertlm/Engine.kt b/kotlin/java/com/google/ai/edge/litertlm/Engine.kt index 8311b5f64..526ace3b0 100644 --- a/kotlin/java/com/google/ai/edge/litertlm/Engine.kt +++ b/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, diff --git a/kotlin/java/com/google/ai/edge/litertlm/LiteRtLmJni.kt b/kotlin/java/com/google/ai/edge/litertlm/LiteRtLmJni.kt index 6f59421b8..ce0857921 100644 --- a/kotlin/java/com/google/ai/edge/litertlm/LiteRtLmJni.kt +++ b/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,7 +242,6 @@ 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, @@ -252,7 +249,6 @@ internal object LiteRtLmJni { extraContextJsonString: String, callback: JniMessageCallback, visualTokenBudget: Int?, - maxOutputToken: Int, ) /** @@ -264,7 +260,6 @@ 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( @@ -272,7 +267,6 @@ internal object LiteRtLmJni { messageJsonString: String, extraContextJsonString: String, visualTokenBudget: Int?, - maxOutputToken: Int, ): String /** diff --git a/kotlin/java/com/google/ai/edge/litertlm/jni/litertlm.cc b/kotlin/java/com/google/ai/edge/litertlm/jni/litertlm.cc index d5a6b1dc5..e93ac56b0 100644 --- a/kotlin/java/com/google/ai/edge/litertlm/jni/litertlm.cc +++ b/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_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_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()) {