Skip to content

Commit d75311b

Browse files
vova-jbvova-jb
committed
KG-675: fix Opentelemetry feature failure for failed llm requests
1 parent d5bc207 commit d75311b

File tree

16 files changed

+367
-61
lines changed

16 files changed

+367
-61
lines changed

‎agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/feature/ContextualPromptExecutor.kt‎

Lines changed: 32 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import ai.koog.prompt.message.LLMChoice
1010
import ai.koog.prompt.message.Message
1111
import ai.koog.prompt.streaming.StreamFrame
1212
import io.github.oshai.kotlinlogging.KotlinLogging
13+
import kotlinx.coroutines.CancellationException
1314
import kotlinx.coroutines.flow.Flow
1415
import kotlinx.coroutines.flow.catch
1516
import kotlinx.coroutines.flow.onCompletion
@@ -41,22 +42,37 @@ public class ContextualPromptExecutor(
4142
logger.debug { "Executing LLM call (event id: $eventId, prompt: $prompt, tools: [${tools.joinToString { it.name }}])" }
4243
context.pipeline.onLLMCallStarting(eventId, context.executionInfo, context.runId, prompt, model, tools, context)
4344

44-
val responses = executor.execute(prompt, model, tools)
45-
46-
logger.trace { "Finished LLM call (event id: $eventId) with responses: [${responses.joinToString { "${it.role}: ${it.content}" }}]" }
47-
context.pipeline.onLLMCallCompleted(
48-
eventId,
49-
context.executionInfo,
50-
context.runId,
51-
prompt,
52-
model,
53-
tools,
54-
responses,
55-
null,
56-
context
57-
)
58-
59-
return responses
45+
return try {
46+
val responses = executor.execute(prompt, model, tools)
47+
logger.trace { "Finished LLM call (event id: $eventId) with responses: [${responses.joinToString { "${it.role}: ${it.content}" }}]" }
48+
context.pipeline.onLLMCallCompleted(
49+
eventId,
50+
context.executionInfo,
51+
context.runId,
52+
prompt,
53+
model,
54+
tools,
55+
responses,
56+
null,
57+
context
58+
)
59+
responses
60+
} catch (e: CancellationException) {
61+
throw e
62+
} catch (e: Exception) {
63+
logger.trace(e) { "Failed LLM call (event id: $eventId) with responses" }
64+
context.pipeline.onLLMCallFailed(
65+
eventId,
66+
context.executionInfo,
67+
context.runId,
68+
prompt,
69+
model,
70+
tools,
71+
context,
72+
e
73+
)
74+
throw e
75+
}
6076
}
6177

6278
/**

‎agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/feature/handler/llm/LLMCallEventContext.kt‎

Lines changed: 39 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,13 @@ import ai.koog.prompt.message.Message
1313
/**
1414
* Represents the context for handling LLM-specific events within the framework.
1515
*/
16-
public interface LLMCallEventContext : AgentLifecycleEventContext
16+
public interface LLMCallEventContext : AgentLifecycleEventContext {
17+
public val runId: String
18+
public val prompt: Prompt
19+
public val model: LLModel
20+
public val tools: List<ToolDescriptor>
21+
public val context: AIAgentContext
22+
}
1723

1824
/**
1925
* Represents the context for handling a before LLM call event.
@@ -27,11 +33,33 @@ public interface LLMCallEventContext : AgentLifecycleEventContext
2733
public data class LLMCallStartingContext(
2834
override val eventId: String,
2935
override val executionInfo: AgentExecutionInfo,
30-
val runId: String,
31-
val prompt: Prompt,
32-
val model: LLModel,
33-
val tools: List<ToolDescriptor>,
34-
val context: AIAgentContext
36+
override val runId: String,
37+
override val prompt: Prompt,
38+
override val model: LLModel,
39+
override val tools: List<ToolDescriptor>,
40+
override val context: AIAgentContext
41+
) : LLMCallEventContext {
42+
override val eventType: AgentLifecycleEventType = AgentLifecycleEventType.LLMCallStarting
43+
}
44+
45+
/**
46+
* Represents the context for handling a after LLM call failed.
47+
*
48+
* @property executionInfo The execution information containing parentId and current execution path;
49+
* @property runId The unique identifier for this LLM call session.
50+
* @property prompt The prompt that will be sent to the language model.
51+
* @property model The language model instance being used.
52+
* @property tools The list of tool descriptors available for the LLM call.
53+
*/
54+
public data class LLMCallFailedContext(
55+
override val eventId: String,
56+
override val executionInfo: AgentExecutionInfo,
57+
override val runId: String,
58+
override val prompt: Prompt,
59+
override val model: LLModel,
60+
override val tools: List<ToolDescriptor>,
61+
override val context: AIAgentContext,
62+
val error: Throwable
3563
) : LLMCallEventContext {
3664
override val eventType: AgentLifecycleEventType = AgentLifecycleEventType.LLMCallStarting
3765
}
@@ -50,13 +78,13 @@ public data class LLMCallStartingContext(
5078
public data class LLMCallCompletedContext(
5179
override val eventId: String,
5280
override val executionInfo: AgentExecutionInfo,
53-
val runId: String,
54-
val prompt: Prompt,
55-
val model: LLModel,
56-
val tools: List<ToolDescriptor>,
81+
override val runId: String,
82+
override val prompt: Prompt,
83+
override val model: LLModel,
84+
override val tools: List<ToolDescriptor>,
5785
val responses: List<Message.Response>,
5886
val moderationResponse: ModerationResult?,
59-
val context: AIAgentContext
87+
override val context: AIAgentContext
6088
) : LLMCallEventContext {
6189
override val eventType: AgentLifecycleEventType = AgentLifecycleEventType.LLMCallCompleted
6290
}

‎agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/feature/handler/llm/LLMCallEventHandler.kt‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,21 @@ public class LLMCallEventHandler {
2929
*/
3030
public var llmCallCompletedHandler: LLMCallCompletedHandler =
3131
LLMCallCompletedHandler { _ -> }
32+
33+
/**
34+
* A handler invoked when an error occurs during a call to a language model (LLM).
35+
*
36+
* This variable represents a custom implementation of the `LLMCallFailedHandler` functional interface,
37+
* allowing post-processing or custom logic to be performed when an error occurs during LLM processing.
38+
*
39+
* The handler receives various pieces of information about the LLM call, including the original prompt,
40+
* the tools used, the model invoked, the responses returned by the model, a unique run identifier, and
41+
* the error that occurred.
42+
*
43+
* Customize this handler to implement specific behavior required when LLM processing fails.
44+
*/
45+
public var llmCallFailedHandler: LLMCallFailedHandler =
46+
LLMCallFailedHandler { _ -> }
3247
}
3348

3449
/**
@@ -62,3 +77,18 @@ public fun interface LLMCallCompletedHandler {
6277
*/
6378
public suspend fun handle(eventContext: LLMCallCompletedContext)
6479
}
80+
81+
/**
82+
* Represents a functional interface for handling operations or logic that should occur after a call
83+
* to a large language model (LLM) was made and failed. The implementation of this interface provides a mechanism
84+
* to perform custom logic or processing based on the provided inputs, such as the prompt, tools,
85+
* model, and generated responses.
86+
*/
87+
public fun interface LLMCallFailedHandler {
88+
/**
89+
* Handles the post-processing of a prompt and its associated data after a language model call.
90+
*
91+
* @param eventContext The context for the event
92+
*/
93+
public suspend fun handle(eventContext: LLMCallFailedContext)
94+
}

‎agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/feature/pipeline/AIAgentPipeline.kt‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import ai.koog.agents.core.feature.handler.agent.AgentEnvironmentTransformingCon
2020
import ai.koog.agents.core.feature.handler.agent.AgentExecutionFailedContext
2121
import ai.koog.agents.core.feature.handler.agent.AgentStartingContext
2222
import ai.koog.agents.core.feature.handler.llm.LLMCallCompletedContext
23+
import ai.koog.agents.core.feature.handler.llm.LLMCallFailedContext
2324
import ai.koog.agents.core.feature.handler.llm.LLMCallStartingContext
2425
import ai.koog.agents.core.feature.handler.strategy.StrategyCompletedContext
2526
import ai.koog.agents.core.feature.handler.strategy.StrategyStartingContext
@@ -289,6 +290,28 @@ public expect abstract class AIAgentPipeline(agentConfig: AIAgentConfig, clock:
289290
context: AIAgentContext
290291
)
291292

293+
/**
294+
* Notifies all registered LLM handlers after a language model call fails with an exception.
295+
*
296+
* @param eventId The unique identifier for the event group.
297+
* @param executionInfo The execution information for the LLM call event
298+
* @param runId Identifier for the current run.
299+
* @param prompt The prompt that was sent to the language model
300+
* @param model The language model instance that processed the request
301+
* @param tools The list of tool descriptors that were available for the LLM call
302+
* @param error The [Throwable] exception instance that was thrown during the call
303+
*/
304+
public override suspend fun onLLMCallFailed(
305+
eventId: String,
306+
executionInfo: AgentExecutionInfo,
307+
runId: String,
308+
prompt: Prompt,
309+
model: LLModel,
310+
tools: List<ToolDescriptor>,
311+
context: AIAgentContext,
312+
error: Throwable
313+
)
314+
292315
//endregion Trigger LLM Call Handlers
293316

294317
//region Trigger Tool Call Handlers
@@ -658,6 +681,18 @@ public expect abstract class AIAgentPipeline(agentConfig: AIAgentConfig, clock:
658681
handle: suspend (eventContext: LLMCallCompletedContext) -> Unit
659682
)
660683

684+
/**
685+
* Intercepts failed calls to a Large Language Model (LLM) and provides a mechanism to handle such cases.
686+
*
687+
* @param feature The AI agent feature associated with the LLM call.
688+
* @param handle A suspendable function to handle the event when an LLM call fails. The function receives an
689+
* instance of [LLMCallFailedContext] that contains contextual information about the failure.
690+
*/
691+
public override fun interceptLLMCallFailed(
692+
feature: AIAgentFeature<*, *>,
693+
handle: suspend (eventContext: LLMCallFailedContext) -> Unit
694+
)
695+
661696
/**
662697
* Intercepts streaming operations before they begin to modify or log the streaming request.
663698
*

‎agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/feature/pipeline/AIAgentPipelineAPI.kt‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import ai.koog.agents.core.feature.handler.agent.AgentEnvironmentTransformingCon
2020
import ai.koog.agents.core.feature.handler.agent.AgentExecutionFailedContext
2121
import ai.koog.agents.core.feature.handler.agent.AgentStartingContext
2222
import ai.koog.agents.core.feature.handler.llm.LLMCallCompletedContext
23+
import ai.koog.agents.core.feature.handler.llm.LLMCallFailedContext
2324
import ai.koog.agents.core.feature.handler.llm.LLMCallStartingContext
2425
import ai.koog.agents.core.feature.handler.strategy.StrategyCompletedContext
2526
import ai.koog.agents.core.feature.handler.strategy.StrategyStartingContext
@@ -136,6 +137,17 @@ public interface AIAgentPipelineAPI {
136137
context: AIAgentContext
137138
)
138139

140+
public suspend fun onLLMCallFailed(
141+
eventId: String,
142+
executionInfo: AgentExecutionInfo,
143+
runId: String,
144+
prompt: Prompt,
145+
model: LLModel,
146+
tools: List<ToolDescriptor>,
147+
context: AIAgentContext,
148+
error: Throwable
149+
)
150+
139151
public suspend fun onLLMCallCompleted(
140152
eventId: String,
141153
executionInfo: AgentExecutionInfo,
@@ -284,6 +296,11 @@ public interface AIAgentPipelineAPI {
284296
handle: suspend (eventContext: LLMCallCompletedContext) -> Unit
285297
)
286298

299+
public fun interceptLLMCallFailed(
300+
feature: AIAgentFeature<*, *>,
301+
handle: suspend (eventContext: LLMCallFailedContext) -> Unit
302+
)
303+
287304
public fun interceptLLMStreamingStarting(
288305
feature: AIAgentFeature<*, *>,
289306
handle: suspend (eventContext: LLMStreamingStartingContext) -> Unit

‎agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/feature/pipeline/AIAgentPipelineImpl.kt‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,8 @@ import ai.koog.agents.core.feature.handler.agent.AgentStartingHandler
3232
import ai.koog.agents.core.feature.handler.llm.LLMCallCompletedContext
3333
import ai.koog.agents.core.feature.handler.llm.LLMCallCompletedHandler
3434
import ai.koog.agents.core.feature.handler.llm.LLMCallEventHandler
35+
import ai.koog.agents.core.feature.handler.llm.LLMCallFailedContext
36+
import ai.koog.agents.core.feature.handler.llm.LLMCallFailedHandler
3537
import ai.koog.agents.core.feature.handler.llm.LLMCallStartingContext
3638
import ai.koog.agents.core.feature.handler.llm.LLMCallStartingHandler
3739
import ai.koog.agents.core.feature.handler.strategy.StrategyCompletedContext
@@ -290,6 +292,30 @@ public class AIAgentPipelineImpl(
290292
llmCallEventHandlers.values.forEach { handler -> handler.llmCallStartingHandler.handle(eventContext) }
291293
}
292294

295+
public override suspend fun onLLMCallFailed(
296+
eventId: String,
297+
executionInfo: AgentExecutionInfo,
298+
runId: String,
299+
prompt: Prompt,
300+
model: LLModel,
301+
tools: List<ToolDescriptor>,
302+
context: AIAgentContext,
303+
error: Throwable
304+
) {
305+
val eventContext =
306+
LLMCallFailedContext(
307+
eventId,
308+
executionInfo,
309+
runId,
310+
prompt,
311+
model,
312+
tools,
313+
context,
314+
error
315+
)
316+
llmCallEventHandlers.values.forEach { handler -> handler.llmCallFailedHandler.handle(eventContext) }
317+
}
318+
293319
public override suspend fun onLLMCallCompleted(
294320
eventId: String,
295321
executionInfo: AgentExecutionInfo,
@@ -589,6 +615,17 @@ public class AIAgentPipelineImpl(
589615
)
590616
}
591617

618+
public override fun interceptLLMCallFailed(
619+
feature: AIAgentFeature<*, *>,
620+
handle: suspend (eventContext: LLMCallFailedContext) -> Unit
621+
) {
622+
val handler = llmCallEventHandlers.getOrPut(feature.key) { LLMCallEventHandler() }
623+
624+
handler.llmCallFailedHandler = LLMCallFailedHandler(
625+
function = createConditionalHandler(feature, handle)
626+
)
627+
}
628+
592629
public override fun interceptLLMStreamingStarting(
593630
feature: AIAgentFeature<*, *>,
594631
handle: suspend (eventContext: LLMStreamingStartingContext) -> Unit

‎agents/agents-features/agents-features-opentelemetry/build.gradle.kts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,8 @@ kotlin {
3636
dependencies {
3737
implementation(kotlin("test"))
3838
implementation(libs.kotlinx.coroutines.test)
39+
implementation(libs.ktor.client.core)
40+
implementation(libs.ktor.client.cio)
3941
}
4042
}
4143

‎agents/agents-features/agents-features-opentelemetry/src/jvmMain/kotlin/ai/koog/agents/features/opentelemetry/extension/SpanExt.kt‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,12 @@ package ai.koog.agents.features.opentelemetry.extension
22

33
import ai.koog.agents.core.feature.model.AIAgentError
44
import ai.koog.agents.features.opentelemetry.attribute.Attribute
5+
import ai.koog.agents.features.opentelemetry.attribute.CommonAttributes
56
import ai.koog.agents.features.opentelemetry.attribute.toSdkAttributes
67
import ai.koog.agents.features.opentelemetry.event.GenAIAgentEvent
8+
import ai.koog.agents.features.opentelemetry.span.GenAIAgentSpan
79
import ai.koog.agents.features.opentelemetry.span.SpanEndStatus
10+
import ai.koog.http.client.KoogHttpClientException
811
import io.opentelemetry.api.trace.Span
912
import io.opentelemetry.api.trace.SpanBuilder
1013
import io.opentelemetry.api.trace.StatusCode
@@ -63,3 +66,22 @@ internal fun AIAgentError?.toSpanEndStatus(): SpanEndStatus =
6366
} else {
6467
SpanEndStatus(code = StatusCode.ERROR, description = this.message)
6568
}
69+
70+
/**
71+
* We want to add an actual error type as an error attribute.
72+
*/
73+
internal fun GenAIAgentSpan.addCommonErrorAttributes(error: Throwable?) {
74+
error?.let {
75+
addAttribute(CommonAttributes.Error.Type(it.extractActualErrorType()))
76+
}
77+
}
78+
79+
private fun Throwable.extractActualErrorType(): String {
80+
val cause = this.cause
81+
return when {
82+
this is KoogHttpClientException -> "KoogHttpClientException-$clientName-httpCode=$statusCode"
83+
cause is KoogHttpClientException -> "KoogHttpClientException-${cause.clientName}-httpCode=${cause.statusCode}"
84+
cause != null -> "${this.javaClass.simpleName}-${cause.javaClass.typeName}"
85+
else -> this.javaClass.typeName
86+
}
87+
}

0 commit comments

Comments
 (0)