Skip to content

Commit aa4a2e8

Browse files
sdubovsdubov
committed
Add tool arguments for gen_ai.choice and gen_ai.assistant.message events (#462)
1 parent c0ca25d commit aa4a2e8

File tree

9 files changed

+221
-10
lines changed

9 files changed

+221
-10
lines changed

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

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,12 @@ import ai.koog.agents.features.opentelemetry.attribute.Attribute
44
import ai.koog.agents.features.opentelemetry.attribute.CommonAttributes
55
import ai.koog.prompt.llm.LLMProvider
66
import ai.koog.prompt.message.Message
7+
import kotlinx.serialization.json.JsonObject
78

8-
internal data class AssistantMessageEvent(
9-
private val provider: LLMProvider,
9+
internal class AssistantMessageEvent(
10+
provider: LLMProvider,
1011
private val message: Message.Response,
12+
private val arguments: JsonObject? = null,
1113
override val verbose: Boolean = false
1214
) : GenAIAgentEvent {
1315

@@ -26,6 +28,7 @@ internal data class AssistantMessageEvent(
2628
is Message.Assistant -> {
2729
if (verbose) {
2830
add(EventBodyFields.Content(content = message.content))
31+
arguments?.let { add(EventBodyFields.Arguments(it)) }
2932
}
3033
}
3134

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,12 @@ import ai.koog.agents.features.opentelemetry.attribute.Attribute
44
import ai.koog.agents.features.opentelemetry.attribute.CommonAttributes
55
import ai.koog.prompt.llm.LLMProvider
66
import ai.koog.prompt.message.Message
7+
import kotlinx.serialization.json.JsonObject
78

89
internal class ChoiceEvent(
910
provider: LLMProvider,
1011
private val message: Message.Response,
12+
private val arguments: JsonObject? = null,
1113
override val verbose: Boolean = false,
1214
) : GenAIAgentEvent {
1315

@@ -33,6 +35,8 @@ internal class ChoiceEvent(
3335
content = message.content
3436
)
3537
)
38+
39+
arguments?.let { add(EventBodyFields.Arguments(it)) }
3640
}
3741
}
3842

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

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
package ai.koog.agents.features.opentelemetry.event
22

3+
import kotlinx.serialization.json.JsonObject
4+
35
internal object EventBodyFields {
46

57
data class ToolCalls(
@@ -23,6 +25,11 @@ internal object EventBodyFields {
2325
}
2426
}
2527

28+
data class Arguments(private val arguments: JsonObject) : EventBodyField() {
29+
override val key: String = "arguments"
30+
override val value: String = arguments.toString()
31+
}
32+
2633
data class Content(private val content: String) : EventBodyField() {
2734
override val key: String = "content"
2835
override val value: String = content
@@ -43,8 +50,11 @@ internal object EventBodyFields {
4350
override val value: String = reason
4451
}
4552

46-
data class Message(private val role: ai.koog.prompt.message.Message.Role?, private val content: String) :
47-
EventBodyField() {
53+
data class Message(
54+
private val role: ai.koog.prompt.message.Message.Role?,
55+
private val content: String
56+
) : EventBodyField() {
57+
4858
override val key: String = "message"
4959
override val value: Map<String, String> = buildMap {
5060
role?.let { role -> put("role", role.name.lowercase()) }

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,8 @@ import ai.koog.agents.features.opentelemetry.attribute.CommonAttributes
55
import ai.koog.prompt.llm.LLMProvider
66
import ai.koog.prompt.message.Message
77

8-
internal data class SystemMessageEvent(
9-
private val provider: LLMProvider,
8+
internal class SystemMessageEvent(
9+
provider: LLMProvider,
1010
private val message: Message.System,
1111
override val verbose: Boolean = false
1212
) : GenAIAgentEvent {

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ import ai.koog.agents.features.opentelemetry.attribute.CommonAttributes
55
import ai.koog.prompt.llm.LLMProvider
66

77
internal class ToolMessageEvent(
8-
private val provider: LLMProvider,
8+
provider: LLMProvider,
99
private val toolCallId: String?,
1010
private val content: String,
1111
override val verbose: Boolean = false

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,8 @@ import ai.koog.agents.features.opentelemetry.attribute.CommonAttributes
55
import ai.koog.prompt.llm.LLMProvider
66
import ai.koog.prompt.message.Message
77

8-
internal data class UserMessageEvent(
9-
private val provider: LLMProvider,
8+
internal class UserMessageEvent(
9+
provider: LLMProvider,
1010
private val message: Message.User,
1111
override val verbose: Boolean = false
1212
) : GenAIAgentEvent {

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -305,7 +305,9 @@ public class OpenTelemetry {
305305
eventContext.responses.map { message ->
306306
when (message) {
307307
is Message.Assistant -> add(AssistantMessageEvent(provider, message, verbose = config.isVerbose))
308-
is Message.Tool.Call -> add(ChoiceEvent(provider, message, verbose = config.isVerbose))
308+
is Message.Tool.Call -> add(
309+
ChoiceEvent(provider, message, arguments = message.contentJson, verbose = config.isVerbose)
310+
)
309311
}
310312
}
311313

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

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@ import ai.koog.agents.features.opentelemetry.mock.MockLLMProvider
55
import ai.koog.prompt.message.Message
66
import ai.koog.prompt.message.ResponseMetaInfo
77
import kotlinx.datetime.Clock
8+
import kotlinx.serialization.json.buildJsonObject
9+
import kotlinx.serialization.json.put
810
import kotlin.test.Test
911
import kotlin.test.assertContentEquals
1012
import kotlin.test.assertEquals
@@ -135,6 +137,97 @@ class AssistantMessageEventTest {
135137

136138
//endregion Body Fields
137139

140+
//region Arguments Tests
141+
142+
@Test
143+
fun `test assistant message verbose true with arguments`() {
144+
val expectedContent = "Test message"
145+
val expectedMessage = createTestAssistantMessage(expectedContent)
146+
val args = buildJsonObject {
147+
put("string", "value")
148+
put("integer", 42)
149+
}
150+
151+
val assistantMessageEvent = AssistantMessageEvent(
152+
provider = MockLLMProvider(),
153+
message = expectedMessage,
154+
arguments = args,
155+
verbose = true,
156+
)
157+
158+
val expectedBodyFields = listOf(
159+
EventBodyFields.Content(content = expectedContent),
160+
EventBodyFields.Arguments(args)
161+
)
162+
163+
assertEquals(expectedBodyFields.size, assistantMessageEvent.bodyFields.size)
164+
assertContentEquals(expectedBodyFields, assistantMessageEvent.bodyFields)
165+
}
166+
167+
@Test
168+
fun `test assistant message verbose false with arguments`() {
169+
val expectedContent = "Test message"
170+
val expectedMessage = createTestAssistantMessage(expectedContent)
171+
val args = buildJsonObject { put("x", "y") }
172+
173+
val assistantMessageEvent = AssistantMessageEvent(
174+
provider = MockLLMProvider(),
175+
message = expectedMessage,
176+
arguments = args,
177+
verbose = false,
178+
)
179+
180+
assertTrue(
181+
assistantMessageEvent.bodyFields.isEmpty(),
182+
"No body fields should be added with verbose set to 'false' even if arguments are provided"
183+
)
184+
}
185+
186+
@Test
187+
fun `test tool call message verbose true ignores arguments`() {
188+
val expectedContent = "Test message"
189+
val expectedMessage = createTestToolCallMessage("test-id", "test-tool", expectedContent)
190+
val args = buildJsonObject { put("ignored", true) }
191+
192+
val assistantMessageEvent = AssistantMessageEvent(
193+
provider = MockLLMProvider(),
194+
message = expectedMessage,
195+
arguments = args,
196+
verbose = true,
197+
)
198+
199+
val expectedBodyFields = listOf(
200+
EventBodyFields.Role(role = Message.Role.Tool),
201+
EventBodyFields.ToolCalls(tools = listOf(expectedMessage))
202+
)
203+
204+
assertEquals(expectedBodyFields.size, assistantMessageEvent.bodyFields.size)
205+
assertContentEquals(expectedBodyFields, assistantMessageEvent.bodyFields)
206+
}
207+
208+
@Test
209+
fun `test tool call message verbose false ignores arguments`() {
210+
val expectedContent = "Test message"
211+
val expectedMessage = createTestToolCallMessage("test-id", "test-tool", expectedContent)
212+
val args = buildJsonObject { put("ignored", true) }
213+
214+
val assistantMessageEvent = AssistantMessageEvent(
215+
provider = MockLLMProvider(),
216+
message = expectedMessage,
217+
arguments = args,
218+
verbose = false,
219+
)
220+
221+
val expectedBodyFields = listOf(
222+
EventBodyFields.Role(role = Message.Role.Tool)
223+
)
224+
225+
assertEquals(expectedBodyFields.size, assistantMessageEvent.bodyFields.size)
226+
assertContentEquals(expectedBodyFields, assistantMessageEvent.bodyFields)
227+
}
228+
229+
//endregion Arguments Tests
230+
138231
//region Private Methods
139232

140233
private fun createTestAssistantMessage(content: String): Message.Response = Message.Assistant(

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

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@ import ai.koog.agents.features.opentelemetry.mock.MockLLMProvider
55
import ai.koog.prompt.message.Message
66
import ai.koog.prompt.message.ResponseMetaInfo
77
import kotlinx.datetime.Clock
8+
import kotlinx.serialization.json.buildJsonObject
9+
import kotlinx.serialization.json.put
810
import kotlin.test.Test
911
import kotlin.test.assertContentEquals
1012
import kotlin.test.assertEquals
@@ -160,6 +162,103 @@ class ChoiceEventTest {
160162

161163
//endregion Body Fields
162164

165+
//region Arguments Tests
166+
167+
@Test
168+
fun `test assistant message verbose true with arguments`() {
169+
val expectedContent = "Test message"
170+
val expectedMessage = createTestAssistantMessage(expectedContent)
171+
val args = buildJsonObject {
172+
put("string", "value")
173+
put("integer", 42)
174+
}
175+
176+
val choiceEvent = ChoiceEvent(
177+
provider = MockLLMProvider(),
178+
message = expectedMessage,
179+
arguments = args,
180+
verbose = true,
181+
)
182+
183+
val expectedBodyFields = listOf(
184+
EventBodyFields.Index(0),
185+
EventBodyFields.Message(
186+
role = null,
187+
content = expectedContent
188+
),
189+
EventBodyFields.Arguments(args)
190+
)
191+
192+
assertEquals(expectedBodyFields.size, choiceEvent.bodyFields.size)
193+
assertContentEquals(expectedBodyFields, choiceEvent.bodyFields)
194+
}
195+
196+
@Test
197+
fun `test assistant message verbose false with arguments`() {
198+
val expectedContent = "Test message"
199+
val expectedMessage = createTestAssistantMessage(expectedContent)
200+
val args = buildJsonObject { put("x", "y") }
201+
202+
val choiceEvent = ChoiceEvent(
203+
provider = MockLLMProvider(),
204+
message = expectedMessage,
205+
arguments = args,
206+
verbose = false,
207+
)
208+
209+
val expectedBodyFields = listOf(
210+
EventBodyFields.Index(0)
211+
)
212+
213+
assertEquals(expectedBodyFields.size, choiceEvent.bodyFields.size)
214+
assertContentEquals(expectedBodyFields, choiceEvent.bodyFields)
215+
}
216+
217+
@Test
218+
fun `test tool call message verbose true ignores arguments`() {
219+
val expectedContent = "Test message"
220+
val expectedMessage = createTestToolCallMessage("test-id", "test-tool", expectedContent)
221+
val args = buildJsonObject { put("ignored", true) }
222+
223+
val choiceEvent = ChoiceEvent(
224+
provider = MockLLMProvider(),
225+
message = expectedMessage,
226+
arguments = args,
227+
verbose = true,
228+
)
229+
230+
val expectedBodyFields = listOf(
231+
EventBodyFields.Index(0),
232+
EventBodyFields.ToolCalls(tools = listOf(expectedMessage))
233+
)
234+
235+
assertEquals(expectedBodyFields.size, choiceEvent.bodyFields.size)
236+
assertContentEquals(expectedBodyFields, choiceEvent.bodyFields)
237+
}
238+
239+
@Test
240+
fun `test tool call message verbose false ignores arguments`() {
241+
val expectedContent = "Test message"
242+
val expectedMessage = createTestToolCallMessage("test-id", "test-tool", expectedContent)
243+
val args = buildJsonObject { put("ignored", true) }
244+
245+
val choiceEvent = ChoiceEvent(
246+
provider = MockLLMProvider(),
247+
message = expectedMessage,
248+
arguments = args,
249+
verbose = false,
250+
)
251+
252+
val expectedBodyFields = listOf(
253+
EventBodyFields.Index(0)
254+
)
255+
256+
assertEquals(expectedBodyFields.size, choiceEvent.bodyFields.size)
257+
assertContentEquals(expectedBodyFields, choiceEvent.bodyFields)
258+
}
259+
260+
//endregion Arguments Tests
261+
163262
//region Private Methods
164263

165264
private fun createTestAssistantMessage(content: String, finishReason: String? = null): Message.Assistant =

0 commit comments

Comments
 (0)