Open Telemetry. Add index for ChoiceEvent in OpenTelemetry

Author: sdubov

agents/agents-features/agents-features-opentelemetry/src/jvmMain/kotlin/ai/koog/agents/features/opentelemetry/event/ChoiceEvent.kt

@@ -10,6 +10,7 @@ internal class ChoiceEvent(
     provider: LLMProvider,
     private val message: Message.Response,
     private val arguments: JsonObject? = null,
+    val index: Int,
     override val verbose: Boolean = false,
 ) : GenAIAgentEvent {
 
@@ -20,7 +21,7 @@ internal class ChoiceEvent(
     }
 
     override val bodyFields: List<EventBodyField> = buildList {
-        add(EventBodyFields.Index(0))
+        add(EventBodyFields.Index(index))
 
         when (message) {
             is Message.Assistant -> {

agents/agents-features/agents-features-opentelemetry/src/jvmMain/kotlin/ai/koog/agents/features/opentelemetry/feature/OpenTelemetry.kt

@@ -279,7 +279,6 @@ public class OpenTelemetry {
                     when (message) {
                         is Message.User -> UserMessageEvent(provider, message, verbose = config.isVerbose)
                         is Message.System -> SystemMessageEvent(provider, message, verbose = config.isVerbose)
-                        is Message.Assistant -> AssistantMessageEvent(provider, message, verbose = config.isVerbose)
                         is Message.Tool.Result -> {
                             ToolMessageEvent(
                                 provider = provider,
@@ -325,7 +324,7 @@ public class OpenTelemetry {
                         when (message) {
                             is Message.Assistant -> add(AssistantMessageEvent(provider, message, verbose = config.isVerbose))
                             is Message.Tool.Call -> add(
-                                ChoiceEvent(provider, message, arguments = message.contentJson, verbose = config.isVerbose)
+                                ChoiceEvent(provider, message, arguments = message.contentJson, index = 0, verbose = config.isVerbose)
                             )
                         }
                     }

agents/agents-features/agents-features-opentelemetry/src/jvmTest/kotlin/ai/koog/agents/features/opentelemetry/event/ChoiceEventTest.kt

@@ -24,6 +24,7 @@ class ChoiceEventTest {
         val choiceEvent = ChoiceEvent(
             provider = llmProvider,
             message = expectedMessage,
+            index = 0,
             verbose = false,
         )
 
@@ -44,6 +45,7 @@ class ChoiceEventTest {
         val choiceEvent = ChoiceEvent(
             provider = llmProvider,
             message = expectedMessage,
+            index = 0,
             verbose = true,
         )
 
@@ -67,6 +69,7 @@ class ChoiceEventTest {
         val choiceEvent = ChoiceEvent(
             provider = MockLLMProvider(),
             message = expectedMessage,
+            index = 0,
             verbose = false,
         )
 
@@ -86,6 +89,7 @@ class ChoiceEventTest {
         val choiceEvent = ChoiceEvent(
             provider = MockLLMProvider(),
             message = expectedMessage,
+            index = 0,
             verbose = false,
         )
 
@@ -106,6 +110,7 @@ class ChoiceEventTest {
         val choiceEvent = ChoiceEvent(
             provider = MockLLMProvider(),
             message = expectedMessage,
+            index = 0,
             verbose = false,
         )
 
@@ -125,6 +130,7 @@ class ChoiceEventTest {
         val choiceEvent = ChoiceEvent(
             provider = MockLLMProvider(),
             message = expectedMessage,
+            index = 0,
             verbose = true,
         )
 
@@ -148,6 +154,7 @@ class ChoiceEventTest {
         val choiceEvent = ChoiceEvent(
             provider = MockLLMProvider(),
             message = expectedMessage,
+            index = 0,
             verbose = true,
         )
 
@@ -177,6 +184,7 @@ class ChoiceEventTest {
             provider = MockLLMProvider(),
             message = expectedMessage,
             arguments = args,
+            index = 0,
             verbose = true,
         )
 
@@ -203,6 +211,7 @@ class ChoiceEventTest {
             provider = MockLLMProvider(),
             message = expectedMessage,
             arguments = args,
+            index = 0,
             verbose = false,
         )
 
@@ -224,6 +233,7 @@ class ChoiceEventTest {
             provider = MockLLMProvider(),
             message = expectedMessage,
             arguments = args,
+            index = 0,
             verbose = true,
         )
 
@@ -246,6 +256,7 @@ class ChoiceEventTest {
             provider = MockLLMProvider(),
             message = expectedMessage,
             arguments = args,
+            index = 0,
             verbose = false,
         )