Skip to content

Commit 33d7b62

Browse files
sdubovsdubov
committed
KG-172. Add all messages to the inference span for Open Telemetry
1 parent 1d78f42 commit 33d7b62

File tree

8 files changed

+166
-107
lines changed

8 files changed

+166
-107
lines changed

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

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

3-
import ai.koog.agents.core.tools.ToolResult
43
import ai.koog.agents.features.opentelemetry.attribute.Attribute
54
import ai.koog.agents.features.opentelemetry.attribute.CommonAttributes
65
import ai.koog.prompt.llm.LLMProvider
76

87
internal class ToolMessageEvent(
98
private val provider: LLMProvider,
109
private val toolCallId: String?,
11-
private val toolResult: ToolResult,
10+
private val content: String,
1211
override val verbose: Boolean = false
1312
) : GenAIAgentEvent {
1413

@@ -21,7 +20,7 @@ internal class ToolMessageEvent(
2120
override val bodyFields: List<EventBodyField> = buildList {
2221
// Content
2322
if (verbose) {
24-
add(EventBodyFields.Content(content = toolResult.toStringDefault()))
23+
add(EventBodyFields.Content(content = content))
2524
}
2625

2726
// Id

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

Lines changed: 32 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@ import ai.koog.agents.features.opentelemetry.attribute.CommonAttributes
1010
import ai.koog.agents.features.opentelemetry.attribute.SpanAttributes
1111
import ai.koog.agents.features.opentelemetry.event.AssistantMessageEvent
1212
import ai.koog.agents.features.opentelemetry.event.ChoiceEvent
13-
import ai.koog.agents.features.opentelemetry.event.GenAIAgentEvent
1413
import ai.koog.agents.features.opentelemetry.event.ModerationResponseEvent
1514
import ai.koog.agents.features.opentelemetry.event.SystemMessageEvent
1615
import ai.koog.agents.features.opentelemetry.event.ToolMessageEvent
@@ -261,27 +260,24 @@ public class OpenTelemetry {
261260
spanProcessor.startSpan(inferenceSpan)
262261

263262
// Add events to the InferenceSpan after the span is created
264-
val lastMessage = eventContext.prompt.messages.lastOrNull()
265-
266-
val events: List<GenAIAgentEvent> = lastMessage?.let { message ->
267-
buildList {
268-
when (message) {
269-
is Message.User -> add(UserMessageEvent(provider, message, verbose = config.isVerbose))
270-
is Message.System -> add(SystemMessageEvent(provider, message, verbose = config.isVerbose))
271-
is Message.Assistant -> add(
272-
AssistantMessageEvent(
273-
provider,
274-
message,
275-
verbose = config.isVerbose
276-
)
263+
val eventsFromMessages = eventContext.prompt.messages.mapNotNull { message ->
264+
when (message) {
265+
is Message.User -> UserMessageEvent(provider, message, verbose = config.isVerbose)
266+
is Message.System -> SystemMessageEvent(provider, message, verbose = config.isVerbose)
267+
is Message.Assistant -> AssistantMessageEvent(provider, message, verbose = config.isVerbose)
268+
is Message.Tool.Result -> {
269+
ToolMessageEvent(
270+
provider = provider,
271+
toolCallId = message.id,
272+
content = message.content,
273+
verbose = config.isVerbose
277274
)
278-
279-
else -> {}
280275
}
276+
else -> null
281277
}
282-
} ?: emptyList()
278+
}
283279

284-
inferenceSpan.addEvents(events)
280+
inferenceSpan.addEvents(eventsFromMessages)
285281
}
286282

287283
pipeline.interceptAfterLLMCall(interceptContext) { eventContext ->
@@ -305,22 +301,20 @@ public class OpenTelemetry {
305301
val provider = eventContext.model.provider
306302

307303
// Add events to the InferenceSpan before finishing the span
308-
val lastMessage = eventContext.responses.lastOrNull()
309-
310-
val moderationResult = eventContext.moderationResponse
311-
312-
val events: List<GenAIAgentEvent> = lastMessage?.let { message ->
313-
buildList {
304+
val eventsToAdd = buildList {
305+
eventContext.responses.map { message ->
314306
when (message) {
315-
is Message.Assistant -> add(ChoiceEvent(provider, message, config.isVerbose))
316-
else -> {}
307+
is Message.Assistant -> add(AssistantMessageEvent(provider, message, verbose = config.isVerbose))
308+
is Message.Tool.Call -> add(ChoiceEvent(provider, message, verbose = config.isVerbose))
317309
}
318310
}
319-
} ?: moderationResult?.let {
320-
buildList { add(ModerationResponseEvent(provider, it, config.isVerbose)) }
321-
} ?: emptyList()
322311

323-
inferenceSpan.addEvents(events)
312+
eventContext.moderationResponse?.let { response ->
313+
ModerationResponseEvent(provider, response, config.isVerbose)
314+
}
315+
}
316+
317+
inferenceSpan.addEvents(eventsToAdd)
324318

325319
// Stop InferenceSpan
326320
spanProcessor.endSpan(inferenceSpanId)
@@ -364,38 +358,13 @@ public class OpenTelemetry {
364358
val nodeInfoElement = currentCoroutineContext().getNodeInfoElement()
365359
?: error("Unable to create tool call span due to missing node info in context")
366360

367-
val agentId = agentRunInfoElement.agentId
368-
val runId = agentRunInfoElement.runId
369-
val provider = agentRunInfoElement.agentConfig.model.provider
370-
val nodeName = nodeInfoElement.nodeName
371-
val toolName = eventContext.tool.name
372-
373361
val executeToolSpanId = ExecuteToolSpan.createId(
374-
agentId = agentId,
375-
runId = runId,
376-
nodeName = nodeName,
377-
toolName = toolName
362+
agentId = agentRunInfoElement.agentId,
363+
runId = agentRunInfoElement.runId,
364+
nodeName = nodeInfoElement.nodeName,
365+
toolName = eventContext.tool.name
378366
)
379367

380-
// Add events to the ExecuteToolSpan before finishing the span
381-
val events = buildList {
382-
val toolResult = eventContext.result
383-
logger.debug { "Last tool result message from prompt: $toolResult" }
384-
385-
if (toolResult != null) {
386-
add(
387-
ToolMessageEvent(
388-
provider = provider,
389-
toolCallId = eventContext.toolCallId,
390-
toolResult = toolResult,
391-
verbose = config.isVerbose
392-
)
393-
)
394-
}
395-
}
396-
397-
spanProcessor.addEventsToSpan(spanId = executeToolSpanId, events = events)
398-
399368
// End the ExecuteToolSpan span
400369
spanProcessor.endSpan(executeToolSpanId)
401370
}
@@ -409,16 +378,11 @@ public class OpenTelemetry {
409378
val nodeInfoElement = currentCoroutineContext().getNodeInfoElement()
410379
?: error("Unable to create tool call span due to missing node info in context")
411380

412-
val agentId = agentRunInfoElement.agentId
413-
val runId = agentRunInfoElement.runId
414-
val nodeName = nodeInfoElement.nodeName
415-
val toolName = eventContext.tool.name
416-
417381
val executeToolSpanId = ExecuteToolSpan.createId(
418-
agentId = agentId,
419-
runId = runId,
420-
nodeName = nodeName,
421-
toolName = toolName
382+
agentId = agentRunInfoElement.agentId,
383+
runId = agentRunInfoElement.runId,
384+
nodeName = nodeInfoElement.nodeName,
385+
toolName = eventContext.tool.name
422386
)
423387

424388
// End the ExecuteToolSpan span

‎agents/agents-features/agents-features-opentelemetry/src/jvmMain/kotlin/ai/koog/agents/features/opentelemetry/span/GenAIAgentSpan.kt

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,10 +41,20 @@ internal abstract class GenAIAgentSpan(
4141

4242
abstract val attributes: List<Attribute>
4343

44+
val events: Set<GenAIAgentEvent>
45+
get() = _events
46+
47+
private val _events: MutableSet<GenAIAgentEvent> = mutableSetOf()
48+
4449
fun addEvents(events: List<GenAIAgentEvent>) {
4550
events.forEach { event ->
4651
logger.debug { "Adding event '${event.name}' to span '$spanId'" }
4752

53+
if (_events.contains(event)) {
54+
logger.warn { "Event '${event.name}' already added to span '$spanId'" }
55+
return@forEach
56+
}
57+
4858
// The 'opentelemetry-java' SDK does not have support for event body fields at the moment.
4959
// Pass body fields as attributes until an API is updated.
5060
val attributes = buildList {
@@ -56,5 +66,6 @@ internal abstract class GenAIAgentSpan(
5666

5767
span.addEvent(event.name, attributes.toSdkAttributes())
5868
}
69+
_events.addAll(events)
5970
}
6071
}

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -30,9 +30,9 @@ internal object OpenTelemetryTestAPI {
3030
): AIAgent<String, String> {
3131
val agentConfig = AIAgentConfig(
3232
prompt = prompt(promptId ?: "Test prompt", clock = clock, params = LLMParams(temperature = temperature)) {
33-
system(systemPrompt ?: "Test system message")
34-
user(userPrompt ?: "Test user message")
35-
assistant(assistantPrompt ?: "Test assistant response")
33+
systemPrompt?.let { system(systemPrompt) }
34+
userPrompt?.let { user(userPrompt) }
35+
assistantPrompt?.let { assistant(assistantPrompt) }
3636
},
3737
model = model ?: OpenAIModels.Chat.GPT4o,
3838
maxAgentIterations = 10,

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

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ class ToolMessageEventTest {
2020
val toolMessageEvent = ToolMessageEvent(
2121
provider = llmProvider,
2222
toolCallId = toolCallId,
23-
toolResult = toolResult,
23+
content = toolResult.toStringDefault(),
2424
verbose = false,
2525
)
2626

@@ -41,7 +41,7 @@ class ToolMessageEventTest {
4141
val toolMessageEvent = ToolMessageEvent(
4242
provider = llmProvider,
4343
toolCallId = toolCallId,
44-
toolResult = toolResult,
44+
content = toolResult.toStringDefault(),
4545
verbose = true,
4646
)
4747

@@ -65,7 +65,7 @@ class ToolMessageEventTest {
6565
val toolMessageEvent = ToolMessageEvent(
6666
provider = MockLLMProvider(),
6767
toolCallId = toolCallId,
68-
toolResult = toolResult,
68+
content = toolResult.toStringDefault(),
6969
verbose = false,
7070
)
7171

@@ -85,7 +85,7 @@ class ToolMessageEventTest {
8585
val toolMessageEvent = ToolMessageEvent(
8686
provider = MockLLMProvider(),
8787
toolCallId = toolCallId,
88-
toolResult = toolResult,
88+
content = toolResult.toStringDefault(),
8989
verbose = false,
9090
)
9191

@@ -103,7 +103,7 @@ class ToolMessageEventTest {
103103
val toolMessageEvent = ToolMessageEvent(
104104
provider = MockLLMProvider(),
105105
toolCallId = toolCallId,
106-
toolResult = toolResult,
106+
content = toolResult.toStringDefault(),
107107
verbose = true,
108108
)
109109

@@ -124,7 +124,7 @@ class ToolMessageEventTest {
124124
val toolMessageEvent = ToolMessageEvent(
125125
provider = MockLLMProvider(),
126126
toolCallId = toolCallId,
127-
toolResult = toolResult,
127+
content = toolResult.toStringDefault(),
128128
verbose = true,
129129
)
130130

0 commit comments

Comments
 (0)