Add validation for maxTokens and update tests

Author: aozherelyeva

prompt/prompt-executor/prompt-executor-clients/prompt-executor-anthropic-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/anthropic/DataModel.kt

@@ -34,6 +34,13 @@ public data class AnthropicMessageRequest(
     val stream: Boolean = false,
     val toolChoice: AnthropicToolChoice? = null,
 ) {
+    init {
+        require(maxTokens > 0) { "maxTokens must be greater than 0, but was $maxTokens" }
+        if (temperature != null) {
+            require(temperature >= 0) { "temperature must be greater than 0, but was $temperature" }
+        }
+    }
+
     /**
      * Companion object with default values for request
      */

prompt/prompt-executor/prompt-executor-clients/prompt-executor-anthropic-client/src/jvmTest/kotlin/ai/koog/prompt/executor/clients/anthropic/AnthropicModelsTest.kt

@@ -3,6 +3,8 @@ package ai.koog.prompt.executor.clients.anthropic
 import ai.koog.prompt.executor.clients.list
 import ai.koog.prompt.llm.LLMProvider
 import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertFailsWith
 import kotlin.test.assertSame
 
 class AnthropicModelsTest {
@@ -19,4 +21,38 @@ class AnthropicModelsTest {
             )
         }
     }
+
+    @Test
+    fun `AnthropicMessageRequest should use custom maxTokens when provided`() {
+        val customMaxTokens = 4000
+        val request = AnthropicMessageRequest(
+            model = AnthropicModels.Opus_3.id,
+            messages = emptyList(),
+            maxTokens = customMaxTokens
+        )
+
+        assertEquals(customMaxTokens, request.maxTokens)
+    }
+
+    @Test
+    fun `AnthropicMessageRequest should use default maxTokens when not provided`() {
+        val request = AnthropicMessageRequest(
+            model = AnthropicModels.Opus_3.id,
+            messages = emptyList()
+        )
+
+        assertEquals(AnthropicMessageRequest.MAX_TOKENS_DEFAULT, request.maxTokens)
+    }
+
+    @Test
+    fun `AnthropicMessageRequest should reject zero maxTokens`() {
+        val exception = assertFailsWith<IllegalArgumentException> {
+            AnthropicMessageRequest(
+                model = AnthropicModels.Opus_3.id,
+                messages = emptyList(),
+                maxTokens = 0
+            )
+        }
+        assertEquals("maxTokens must be greater than 0, but was 0", exception.message)
+    }
 }

prompt/prompt-executor/prompt-executor-clients/prompt-executor-bedrock-client/src/jvmMain/kotlin/ai/koog/prompt/executor/clients/bedrock/modelfamilies/ai21/JambaDataModel.kt

@@ -6,20 +6,27 @@ import kotlinx.serialization.json.JsonObject
 
 @Serializable
 internal data class JambaRequest(
-    public val model: String,
-    public val messages: List<JambaMessage>,
-    @SerialName("max_tokens")
-    public val maxTokens: Int? = MAX_TOKENS_DEFAULT,
-    public val temperature: Double? = null,
-    @SerialName("top_p")
-    public val topP: Double? = null,
-    public val stop: List<String>? = null,
-    public val n: Int? = null,
-    public val stream: Boolean? = null,
-    public val tools: List<JambaTool>? = null,
-    @SerialName("response_format")
-    public val responseFormat: JambaResponseFormat? = null
+    val model: String,
+    val messages: List<JambaMessage>,
+    @SerialName("max_tokens") val maxTokens: Int? = MAX_TOKENS_DEFAULT,
+    val temperature: Double? = null,
+    @SerialName("top_p") val topP: Double? = null,
+    val stop: List<String>? = null,
+    val n: Int? = null,
+    val stream: Boolean? = null,
+    val tools: List<JambaTool>? = null,
+    @SerialName("response_format") val responseFormat: JambaResponseFormat? = null
 ) {
+    init {
+        require(maxTokens != null && maxTokens > 0) { "maxTokens must be greater than 0, but was $maxTokens" }
+        if (temperature != null) {
+            require(temperature >= 0) { "temperature must be greater than 0, but was $temperature" }
+        }
+        if (topP != null) {
+            require(topP in 0.0..1.0) { "topP must be between 0 and 1, but was $topP" }
+        }
+    }
+
     /**
      * Companion object with default values for request
      */
@@ -33,90 +40,90 @@ internal data class JambaRequest(
 
 @Serializable
 internal data class JambaMessage(
-    public val role: String,
-    public val content: String? = null,
+    val role: String,
+    val content: String? = null,
     @SerialName("tool_calls")
-    public val toolCalls: List<JambaToolCall>? = null,
+    val toolCalls: List<JambaToolCall>? = null,
     @SerialName("tool_call_id")
-    public val toolCallId: String? = null
+    val toolCallId: String? = null
 )
 
 @Serializable
 internal data class JambaTool(
-    public val type: String = "function",
-    public val function: JambaFunction
+    val type: String = "function",
+    val function: JambaFunction
 )
 
 @Serializable
 internal data class JambaFunction(
-    public val name: String,
-    public val description: String,
-    public val parameters: JsonObject
+    val name: String,
+    val description: String,
+    val parameters: JsonObject
 )
 
 @Serializable
 internal data class JambaToolCall(
-    public val id: String,
-    public val type: String = "function",
-    public val function: JambaFunctionCall
+    val id: String,
+    val type: String = "function",
+    val function: JambaFunctionCall
 )
 
 @Serializable
 internal data class JambaFunctionCall(
-    public val name: String,
-    public val arguments: String
+    val name: String,
+    val arguments: String
 )
 
 @Serializable
 internal data class JambaResponseFormat(
-    public val type: String
+    val type: String
 )
 
 @Serializable
 internal data class JambaResponse(
-    public val id: String,
-    public val model: String,
-    public val choices: List<JambaChoice>,
-    public val usage: JambaUsage? = null
+    val id: String,
+    val model: String,
+    val choices: List<JambaChoice>,
+    val usage: JambaUsage? = null
 )
 
 @Serializable
 internal data class JambaChoice(
-    public val index: Int,
-    public val message: JambaMessage,
+    val index: Int,
+    val message: JambaMessage,
     @SerialName("finish_reason")
-    public val finishReason: String? = null
+    val finishReason: String? = null
 )
 
 @Serializable
 internal data class JambaUsage(
     @SerialName("prompt_tokens")
-    public val promptTokens: Int,
+    val promptTokens: Int,
     @SerialName("completion_tokens")
-    public val completionTokens: Int,
+    val completionTokens: Int,
     @SerialName("total_tokens")
-    public val totalTokens: Int
+    val totalTokens: Int
 )
 
 @Serializable
 internal data class JambaStreamResponse(
-    public val id: String,
-    public val choices: List<JambaStreamChoice>,
-    public val usage: JambaUsage? = null
+    val id: String,
+    val choices: List<JambaStreamChoice>,
+    val usage: JambaUsage? = null
 )
 
 @Serializable
 internal data class JambaStreamChoice(
-    public val index: Int,
-    public val delta: JambaStreamDelta,
+    val index: Int,
+    val delta: JambaStreamDelta,
     @SerialName("finish_reason")
-    public val finishReason: String? = null
+    val finishReason: String? = null
 )
 
 @Serializable
 internal data class JambaStreamDelta(
-    public val role: String? = null,
-    public val content: String? = null,
+    val role: String? = null,
+    val content: String? = null,
     @SerialName("tool_calls")
-    public val toolCalls: List<JambaToolCall>? = null
+    val toolCalls: List<JambaToolCall>? = null
 )

prompt/prompt-executor/prompt-executor-clients/prompt-executor-bedrock-client/src/jvmTest/kotlin/ai/koog/prompt/executor/clients/bedrock/modelfamilies/ai21/BedrockAI21JambaSerializationTest.kt

@@ -5,6 +5,7 @@ import ai.koog.agents.core.tools.ToolParameterDescriptor
 import ai.koog.agents.core.tools.ToolParameterType
 import ai.koog.prompt.dsl.Prompt
 import ai.koog.prompt.executor.clients.bedrock.BedrockModels
+import ai.koog.prompt.executor.clients.bedrock.modelfamilies.ai21.JambaRequest.Companion.MAX_TOKENS_DEFAULT
 import ai.koog.prompt.llm.LLMCapability
 import ai.koog.prompt.llm.LLMProvider
 import ai.koog.prompt.llm.LLModel
@@ -43,7 +44,7 @@ class BedrockAI21JambaSerializationTest {
 
         assertNotNull(request)
         assertEquals(model.id, request.model)
-        assertEquals(4096, request.maxTokens)
+        assertEquals(MAX_TOKENS_DEFAULT, request.maxTokens)
         assertEquals(temperature, request.temperature)
 
         assertEquals(2, request.messages.size)
@@ -55,6 +56,22 @@ class BedrockAI21JambaSerializationTest {
         assertEquals(userMessage, request.messages[1].content)
     }
 
+    @Test
+    fun `createJambaRequest with custom maxTokens`() {
+        val maxTokens = 1000
+
+        val prompt = Prompt.build("test", params = LLMParams(maxTokens = maxTokens)) {
+            system(systemMessage)
+            user(userMessage)
+        }
+
+        val request = BedrockAI21JambaSerialization.createJambaRequest(prompt, model, emptyList())
+
+        assertNotNull(request)
+        assertEquals(model.id, request.model)
+        assertEquals(MAX_TOKENS_DEFAULT, request.maxTokens)
+    }
+
     @Test
     fun `createJambaRequest with conversation history`() {
         val userNewMessage = "Hello, who are you?"

prompt/prompt-executor/prompt-executor-clients/prompt-executor-bedrock-client/src/jvmTest/kotlin/ai/koog/prompt/executor/clients/bedrock/modelfamilies/ai21/JambaDataModelsTest.kt

@@ -1,10 +1,14 @@
 package ai.koog.prompt.executor.clients.bedrock.modelfamilies.ai21
 
+import ai.koog.prompt.executor.clients.anthropic.AnthropicMessageRequest
+import ai.koog.prompt.executor.clients.anthropic.AnthropicModels
+import ai.koog.prompt.executor.clients.bedrock.modelfamilies.ai21.JambaRequest.Companion.MAX_TOKENS_DEFAULT
 import kotlinx.serialization.json.Json
 import kotlinx.serialization.json.buildJsonObject
 import kotlinx.serialization.json.put
 import kotlin.test.Test
 import kotlin.test.assertEquals
+import kotlin.test.assertFailsWith
 import kotlin.test.assertNotNull
 import kotlin.test.assertNull
 
@@ -51,6 +55,30 @@ class JambaDataModelsTest {
         assert(serialized.contains("0.7")) { "Serialized JSON should contain the temperature value: $serialized" }
     }
 
+    @Test
+    fun `JambaRequest serialization with default maxTokens`() {
+        val request = JambaRequest(
+            model = "ai21.jamba-1-5-large-v1:0",
+            messages = listOf(
+                JambaMessage(role = "user", content = "Tell me about Paris")
+            ),
+            temperature = 0.7
+        )
+        assertEquals(MAX_TOKENS_DEFAULT, request.maxTokens)
+    }
+
+    @Test
+    fun `JambaRequest serialization with maxTokens less than 1`() {
+        val exception = assertFailsWith<IllegalArgumentException> {
+            AnthropicMessageRequest(
+                model = AnthropicModels.Opus_3.id,
+                messages = emptyList(),
+                maxTokens = 0
+            )
+        }
+        assertEquals("maxTokens must be greater than 0, but was 0", exception.message)
+    }
+
     @Test
     fun `JambaRequest serialization with null fields`() {
         val request = JambaRequest(
@@ -121,7 +149,7 @@ class JambaDataModelsTest {
         assertEquals(1, request.messages.size)
         assertEquals("user", request.messages[0].role)
         assertEquals("Tell me about Paris", request.messages[0].content)
-        assertNull(request.maxTokens)
+        assertEquals(MAX_TOKENS_DEFAULT, request.maxTokens)
         assertNull(request.temperature)
     }
 

prompt/prompt-executor/prompt-executor-clients/prompt-executor-bedrock-client/src/jvmTest/kotlin/ai/koog/prompt/executor/clients/bedrock/modelfamilies/amazon/BedrockAmazonNovaSerializationTest.kt

@@ -2,6 +2,7 @@ package ai.koog.prompt.executor.clients.bedrock.modelfamilies.amazon
 
 import ai.koog.prompt.dsl.Prompt
 import ai.koog.prompt.executor.clients.bedrock.BedrockModels
+import ai.koog.prompt.executor.clients.bedrock.modelfamilies.amazon.NovaInferenceConfig.Companion.MAX_TOKENS_DEFAULT
 import ai.koog.prompt.llm.LLMCapability
 import ai.koog.prompt.llm.LLMProvider
 import ai.koog.prompt.llm.LLModel
@@ -50,10 +51,23 @@ class BedrockAmazonNovaSerializationTest {
         assertEquals(userMessage, request.messages[0].content[0].text)
 
         assertNotNull(request.inferenceConfig)
-        assertEquals(4096, request.inferenceConfig.maxTokens)
+        assertEquals(MAX_TOKENS_DEFAULT, request.inferenceConfig.maxTokens)
         assertEquals(temperature, request.inferenceConfig.temperature)
     }
 
+    @Test
+    fun `createNovaRequest with default maxTokens`() {
+        val maxTokens = 1000
+
+        val prompt = Prompt.build("test", params = LLMParams(maxTokens = maxTokens)) {
+            system(systemMessage)
+            user(userMessage)
+        }
+
+        val request = BedrockAmazonNovaSerialization.createNovaRequest(prompt, model)
+        assertEquals(maxTokens, request.inferenceConfig!!.maxTokens)
+    }
+
     @Test
     fun `createNovaRequest with conversation history`() {
         val prompt = Prompt.build("test") {

prompt/prompt-executor/prompt-executor-clients/prompt-executor-bedrock-client/src/jvmTest/kotlin/ai/koog/prompt/executor/clients/bedrock/modelfamilies/anthropic/BedrockAnthropicClaudeSerializationTest.kt

@@ -5,6 +5,7 @@ import ai.koog.agents.core.tools.ToolParameterDescriptor
 import ai.koog.agents.core.tools.ToolParameterType
 import ai.koog.prompt.dsl.Prompt
 import ai.koog.prompt.executor.clients.anthropic.AnthropicContent
+import ai.koog.prompt.executor.clients.anthropic.AnthropicMessageRequest
 import ai.koog.prompt.executor.clients.anthropic.AnthropicToolChoice
 import ai.koog.prompt.executor.clients.bedrock.BedrockModels
 import ai.koog.prompt.message.Message
@@ -47,7 +48,7 @@ class BedrockAnthropicClaudeSerializationTest {
 
         assertNotNull(request)
         assertEquals(model.id, request.model)
-        assertEquals(4096, request.maxTokens)
+        assertEquals(AnthropicMessageRequest.MAX_TOKENS_DEFAULT, request.maxTokens)
         assertEquals(temperature, request.temperature)
 
         assertNotNull(request.system)