Skip to content

Commit 436762e

Browse files
sdubovsdubov
committed
Fix comments from CR
- Fixed comments from CR; - Fix tests; - Add event body field and logic to convert them into attributes; - Add more tests for new entities.
1 parent 1163587 commit 436762e

File tree

19 files changed

+405
-139
lines changed

19 files changed

+405
-139
lines changed

‎agents/agents-features/agents-features-opentelemetry/src/jvmMain/kotlin/ai/koog/agents/features/opentelemetry/attribute/EventAttributes.kt

Lines changed: 0 additions & 67 deletions
This file was deleted.

‎agents/agents-features/agents-features-opentelemetry/src/jvmMain/kotlin/ai/koog/agents/features/opentelemetry/attribute/attributes.kt

Lines changed: 27 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -19,10 +19,7 @@ internal fun List<Attribute>.toSdkAttributes() : Attributes {
1919
return object : Attributes {
2020

2121
@Suppress("UNCHECKED_CAST")
22-
override fun <T : Any?> get(key: AttributeKey<T?>): T? {
23-
return sdkAttributesMap[key] as T?
24-
25-
}
22+
override fun <T : Any?> get(key: AttributeKey<T?>): T? = sdkAttributesMap[key] as T?
2623

2724
override fun forEach(consumer: BiConsumer<in AttributeKey<*>, in Any>) {
2825
sdkAttributesMap.forEach { attribute ->
@@ -34,20 +31,9 @@ internal fun List<Attribute>.toSdkAttributes() : Attributes {
3431

3532
override fun isEmpty(): Boolean = sdkAttributesMap.isEmpty()
3633

37-
override fun asMap(): Map<AttributeKey<*>, Any> {
38-
return sdkAttributesMap
39-
}
40-
41-
override fun toBuilder(): AttributesBuilder? {
42-
val builder = Attributes.builder()
34+
override fun asMap(): Map<AttributeKey<*>, Any> = sdkAttributesMap
4335

44-
sdkAttributesMap.forEach { (key, value) ->
45-
@Suppress("UNCHECKED_CAST")
46-
builder.put(key as AttributeKey<Any>, value)
47-
}
48-
49-
return builder
50-
}
36+
override fun toBuilder(): AttributesBuilder = Attributes.builder().also { builder -> builder.addAttributes(sdkAttributesMap) }
5137
}
5238
}
5339

@@ -57,19 +43,33 @@ private fun Attribute.toSdkAttribute(): Pair<AttributeKey<*>, Any> {
5743
val value = this.value
5844

5945
return when (value) {
60-
is String -> Pair(AttributeKey.stringKey(key), value)
61-
is Boolean -> Pair(AttributeKey.booleanKey(key), value)
62-
is Int -> Pair(AttributeKey.longKey(key), value.toLong())
63-
is Long -> Pair(AttributeKey.longKey(key), value)
64-
is Double -> Pair(AttributeKey.doubleKey(key), value)
46+
is CharSequence,
47+
is Char -> {
48+
Pair(AttributeKey.stringKey(key), value)
49+
}
50+
is Boolean -> {
51+
Pair(AttributeKey.booleanKey(key), value)
52+
}
53+
is Int -> {
54+
Pair(AttributeKey.longKey(key), value.toLong())
55+
}
56+
is Long -> {
57+
Pair(AttributeKey.longKey(key), value)
58+
}
59+
is Double -> {
60+
Pair(AttributeKey.doubleKey(key), value)
61+
}
62+
is Float -> {
63+
Pair(AttributeKey.doubleKey(key), value.toDouble())
64+
}
6565
is List<*> -> {
66-
if (value.all { it is String }) {
66+
if (value.all { it is CharSequence || it is Char }) {
6767
Pair(AttributeKey.stringArrayKey(key), value)
6868
}
6969
else if (value.all { it is Boolean }) {
7070
Pair(AttributeKey.booleanArrayKey(key), value)
7171
}
72-
else if (value.all { it is Int }) {
72+
else if (value.all { it is Int}) {
7373
Pair(AttributeKey.longArrayKey(key), value.map { (it as Int).toLong() })
7474
}
7575
else if (value.all { it is Long }) {
@@ -78,6 +78,9 @@ private fun Attribute.toSdkAttribute(): Pair<AttributeKey<*>, Any> {
7878
else if (value.all { it is Double }) {
7979
Pair(AttributeKey.doubleArrayKey(key), value)
8080
}
81+
else if (value.all { it is Float }) {
82+
Pair(AttributeKey.doubleArrayKey(key), value.map { (it as Float).toDouble() })
83+
}
8184
else {
8285
error("Attribute '$key' has unsupported type for List values: ${value.firstOrNull()?.let { it::class.simpleName} }")
8386
}

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

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

33
import ai.koog.agents.features.opentelemetry.attribute.Attribute
44
import ai.koog.agents.features.opentelemetry.attribute.CommonAttributes
5-
import ai.koog.agents.features.opentelemetry.attribute.EventAttributes
65
import ai.koog.prompt.llm.LLMProvider
76
import ai.koog.prompt.message.Message
87

@@ -16,27 +15,29 @@ internal data class AssistantMessageEvent(
1615

1716
override val attributes: List<Attribute> = buildList {
1817
add(CommonAttributes.System(provider))
18+
}
1919

20+
override val bodyFields: List<EventBodyField> = buildList {
2021
if (message.role != Message.Role.Assistant) {
21-
add(EventAttributes.Body.Role(role = message.role))
22+
add(EventBodyFields.Role(role = message.role))
2223
}
2324

2425
when (message) {
2526
is Message.Assistant -> {
26-
add(EventAttributes.Body.Content(content = message.content))
27+
add(EventBodyFields.Content(content = message.content))
2728
}
2829

2930
is Message.Tool.Call -> {
30-
add(EventAttributes.Body.ToolCalls(tools = listOf(message), verbose = verbose))
31+
add(EventBodyFields.ToolCalls(tools = listOf(message), verbose = verbose))
3132
}
3233

3334
is Message.Tool.Result -> {
34-
add(EventAttributes.Body.ToolCalls(tools = listOf(message), verbose = verbose))
35+
add(EventBodyFields.ToolCalls(tools = listOf(message), verbose = verbose))
3536
}
3637

3738
else -> error(
3839
"Expected message types: ${Message.Tool.Call::class.simpleName}, " +
39-
"${Message.Tool.Result::class.simpleName}, but received: ${message::class.simpleName}"
40+
"${Message.Tool.Result::class.simpleName}, but received: ${message::class.simpleName}"
4041
)
4142
}
4243
}

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

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

33
import ai.koog.agents.features.opentelemetry.attribute.Attribute
44
import ai.koog.agents.features.opentelemetry.attribute.CommonAttributes
5-
import ai.koog.agents.features.opentelemetry.attribute.EventAttributes
65
import ai.koog.prompt.llm.LLMProvider
76
import ai.koog.prompt.message.Message
87

@@ -16,24 +15,26 @@ internal class ChoiceEvent(
1615

1716
override val attributes: List<Attribute> = buildList {
1817
add(CommonAttributes.System(provider))
18+
}
1919

20+
override val bodyFields: List<EventBodyField> = buildList {
2021
when (message) {
2122
is Message.Assistant -> {
22-
add(EventAttributes.Body.Message(
23+
add(EventBodyFields.Message(
2324
role = message.role.takeIf { role -> role != Message.Role.Assistant },
2425
content = message.content
2526
))
2627

2728
message.finishReason?.let { reason ->
28-
add(EventAttributes.Body.FinishReason(reason))
29+
add(EventBodyFields.FinishReason(reason))
2930
}
3031

31-
add(EventAttributes.Body.Index(0))
32+
add(EventBodyFields.Index(0))
3233

3334
}
3435
is Message.Tool.Call -> {
35-
add(EventAttributes.Body.ToolCalls(tools = listOf(message), verbose = verbose))
36-
add(EventAttributes.Body.Index(0))
36+
add(EventBodyFields.ToolCalls(tools = listOf(message), verbose = verbose))
37+
add(EventBodyFields.Index(0))
3738
}
3839
}
3940
}

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

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
package ai.koog.agents.features.opentelemetry.event
2+
3+
import ai.koog.agents.features.opentelemetry.attribute.Attribute
4+
import ai.koog.agents.features.opentelemetry.attribute.CustomAttribute
5+
import io.github.oshai.kotlinlogging.KotlinLogging
6+
7+
/**
8+
* Represents an abstract field to be included in an event's body. Each field is characterized
9+
* by a unique key and an associated value, and optionally marked as verbose for additional significance.
10+
*/
11+
internal abstract class EventBodyField {
12+
13+
companion object {
14+
private val logger = KotlinLogging.logger { }
15+
}
16+
17+
abstract val key: String
18+
19+
abstract val value: Any
20+
21+
open val verbose: Boolean
22+
get() = false
23+
24+
internal fun toAttribute(): Attribute {
25+
return CustomAttribute(key, getAttributeValue(key, value), verbose)
26+
}
27+
28+
private fun getAttributeValue(key: Any, value: Any): Any {
29+
return when (value) {
30+
// Types supported by Attributes
31+
is CharSequence, is Char -> {
32+
"\"${value}\""
33+
}
34+
is Boolean,
35+
is Int, is Long,
36+
is Double, is Float -> {
37+
value
38+
}
39+
is List<*> -> {
40+
if (value.all { it is CharSequence } || value.all { it is Char } ||
41+
value.all { it is Boolean } ||
42+
value.all { it is Int } || value.all { it is Long } ||
43+
value.all { it is Double } || value.all { it is Float } ) {
44+
value
45+
}
46+
else {
47+
48+
// Not supported by Attributes
49+
logger.debug { "${key}: Unsupported type for event body: ${value::class.simpleName}. use toString()" }
50+
value.filterNotNull().joinToString(prefix = "[", postfix = "]") { getAttributeValue(key, it).toString() }
51+
}
52+
}
53+
54+
// Types not supported by Attributes
55+
is Map<*, *> -> {
56+
value.entries
57+
.filter { it.key != null && it.value != null }
58+
.joinToString(prefix = "{", postfix = "}") { entry -> "\"${entry.key}\": ${getAttributeValue(entry.key!!, entry.value!!)}" }
59+
}
60+
61+
else -> {
62+
logger.debug { "${key}: Unsupported type for event body: ${value::class.simpleName}. use toString()" }
63+
value.toString()
64+
}
65+
}
66+
}
67+
}

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

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
package ai.koog.agents.features.opentelemetry.event
2+
3+
internal object EventBodyFields {
4+
5+
data class ToolCalls(
6+
private val tools: List<ai.koog.prompt.message.Message.Tool>,
7+
override val verbose: Boolean = false
8+
) : EventBodyField() {
9+
override val key: String = "tool_calls"
10+
override val value: List<Map<String, Any>>
11+
get() {
12+
return tools.map { tool ->
13+
buildMap {
14+
val functionMap = buildMap {
15+
put("name", tool.tool)
16+
if (verbose) {
17+
put("arguments", tool.content)
18+
}
19+
}
20+
21+
put("function", functionMap)
22+
put("id", tool.id ?: "")
23+
put("type", "function")
24+
}
25+
}
26+
}
27+
}
28+
29+
data class Content(private val content: String) : EventBodyField() {
30+
override val key: String = "content"
31+
override val value: String = content
32+
}
33+
34+
data class Role(private val role: ai.koog.prompt.message.Message.Role) : EventBodyField() {
35+
override val key: String = "role"
36+
override val value: String = role.name.lowercase()
37+
}
38+
39+
data class Index(private val index: Int) : EventBodyField() {
40+
override val key: String = "index"
41+
override val value: Int = index
42+
}
43+
44+
data class FinishReason(private val reason: String) : EventBodyField() {
45+
override val key: String = "finish_reason"
46+
override val value: String = reason
47+
}
48+
49+
data class Message(private val role: ai.koog.prompt.message.Message.Role?, private val content: String) :
50+
EventBodyField() {
51+
override val key: String = "message"
52+
override val value: Map<String, String> = buildMap {
53+
role?.let { role -> put("role", role.name.lowercase()) }
54+
put("content", content)
55+
}
56+
}
57+
58+
data class Id(private val id: String) : EventBodyField() {
59+
override val key: String = "id"
60+
override val value: String = id
61+
}
62+
}

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,5 +11,14 @@ internal interface GenAIAgentEvent {
1111

1212
val attributes: List<Attribute>
1313

14+
/**
15+
* The body field for the event.
16+
*
17+
* Note: Currently, the OpenTelemetry SDK does not support event body fields.
18+
* This field is used to store the body fields.
19+
* Fields are merged with attributes when creating the event.
20+
*/
21+
val bodyFields: List<EventBodyField>
22+
1423
fun String.concatName(other: String): String = "$this.$other"
1524
}

0 commit comments

Comments
 (0)