[prompt] [integration-tests] Fix intput/output token counts for OpenAILLMClient (#370)

Author: EugeneTheDev

TESTING.md

@@ -94,6 +94,8 @@ Integration tests that interact with LLM services require API tokens to be set a
 - `OPEN_ROUTER_API_TEST_KEY` - Required for tests using OpenRouter
 
 You need to set these environment variables before running the integration tests that use the corresponding LLM clients.
+To simplify development, you can also create `env.properties` file (already gitignored) using [env.template.propertes](./integration-tests/env.template.properties) as a template.
+Then properties specified there would be automatically applied as environment variables when you run any test task.
 
 ## Running Ollama tests
 

examples/src/main/kotlin/ai/koog/agents/example/ApiKeyService.kt

@@ -9,4 +9,7 @@ internal object ApiKeyService {
 
     val googleApiKey: String
         get() = System.getenv("GOOGLE_API_KEY") ?: throw IllegalArgumentException("GOOGLE_API_KEY env is not set")
+
+    val openRouterApiKey: String
+        get() = System.getenv("OPENROUTER_API_KEY") ?: throw IllegalArgumentException("OPENROUTER_API_KEY env is not set")
 }

integration-tests/build.gradle.kts

@@ -4,6 +4,7 @@ version = rootProject.version
 plugins {
     id("ai.kotlin.multiplatform")
     alias(libs.plugins.kotlin.serialization)
+    id("ai.koog.gradle.plugins.credentialsresolver")
 }
 
 kotlin {
@@ -36,6 +37,16 @@ kotlin {
     }
 }
 
+val envs = credentialsResolver.resolve(
+    layout.projectDirectory.file(provider { "env.properties" })
+)
+
+tasks.withType<Test> {
+    doFirst {
+        environment(envs.get())
+    }
+}
+
 dokka {
     dokkaSourceSets.configureEach {
         suppress.set(true)

integration-tests/env.template.properties

@@ -0,0 +1,10 @@
+# Copy this file without "template".
+# Add your envs here:
+# FOO=bar
+# Or, if you want to integrate with 1password, you can put reference, e.g. op://foo/bar:
+# FOO=op://foo/bar
+
+OPEN_AI_API_TEST_KEY=
+ANTHROPIC_API_TEST_KEY=
+GEMINI_API_TEST_KEY=
+OPEN_ROUTER_API_TEST_KEY=

integration-tests/src/jvmTest/kotlin/ai/koog/integration/tests/SingleLLMPromptExecutorIntegrationTest.kt

@@ -122,10 +122,15 @@ class SingleLLMPromptExecutorIntegrationTest {
             assertNotNull(response, "Response should not be null")
             assertTrue(response.isNotEmpty(), "Response should not be empty")
             assertTrue(response.first() is Message.Assistant, "Response should be an Assistant message")
+
+            val message = response.first() as Message.Assistant
             assertTrue(
-                (response.first() as Message.Assistant).content.contains("Paris", ignoreCase = true),
+                message.content.contains("Paris", ignoreCase = true),
                 "Response should contain 'Paris'"
             )
+            assertNotNull(message.metaInfo.inputTokensCount, "Input tokens count should not be null")
+            assertNotNull(message.metaInfo.outputTokensCount, "Output tokens count should not be null")
+            assertNotNull(message.metaInfo.totalTokensCount, "Total tokens count should not be null")
         }
     }
 

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

@@ -136,8 +136,8 @@ internal data class OpenAIChoice(
 
 @Serializable
 internal data class OpenAIUsage(
-    val inputTokens: Int? = null,
-    val outputTokens: Int? = null,
+    val promptTokens: Int? = null,
+    val completionTokens: Int? = null,
     val totalTokens: Int
 )
 

prompt/prompt-executor/prompt-executor-clients/prompt-executor-openai-client/src/commonMain/kotlin/ai/koog/prompt/executor/clients/openai/OpenAILLMClient.kt

@@ -326,7 +326,7 @@ public open class OpenAILLMClient(
         return withContext(Dispatchers.SuitableForIO) {
             val response = httpClient.post(settings.chatCompletionsPath) {
                 setBody(request)
-            }
+           }
 
             if (response.status.isSuccess()) {
                 response.body<OpenAIResponse>()
@@ -450,8 +450,8 @@ public open class OpenAILLMClient(
 
         // Extract token count from the response
         val totalTokensCount = response.usage?.totalTokens
-        val inputTokensCount = response.usage?.inputTokens
-        val outputTokensCount = response.usage?.outputTokens
+        val inputTokensCount = response.usage?.promptTokens
+        val outputTokensCount = response.usage?.completionTokens
 
         val metaInfo = ResponseMetaInfo.create(
             clock,