Skip to content

Commit 0328497

Browse files
tiginamariatiginamaria
committed
Fix some tests
1 parent 0dcb2ca commit 0328497

33 files changed

Lines changed: 258 additions & 294 deletions

File tree

‎agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/environment/ContextualAgentEnvironment.kt‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ package ai.koog.agents.core.environment
33
import ai.koog.agents.core.agent.context.AIAgentContext
44
import ai.koog.agents.core.agent.execution.AgentExecutionInfo
55
import ai.koog.agents.core.annotation.InternalAgentsApi
6-
import ai.koog.agents.core.feature.model.toAgentError
76
import ai.koog.prompt.message.MessagePart
87
import ai.koog.serialization.JSONObject
98
import ai.koog.serialization.kotlinx.toKoogJSONObject
@@ -58,7 +57,7 @@ public class ContextualAgentEnvironment(
5857
toolDescription = toolDescription,
5958
toolArgs = toolArgs,
6059
message = message,
61-
error = e.toAgentError(),
60+
error = e,
6261
context = context
6362
)
6463
return ReceivedToolResult(
@@ -67,7 +66,7 @@ public class ContextualAgentEnvironment(
6766
toolArgs = toolArgs,
6867
toolDescription = null,
6968
content = message,
70-
resultKind = ToolResultKind.ValidationError(e.toAgentError()),
69+
resultKind = ToolResultKind.ValidationError(e),
7170
result = null
7271
)
7372
}

‎agents/agents-core/src/commonTest/kotlin/ai/koog/agents/core/CalculatorPromptExecutor.kt‎

Lines changed: 18 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import ai.koog.prompt.dsl.Prompt
66
import ai.koog.prompt.executor.model.PromptExecutor
77
import ai.koog.prompt.llm.LLModel
88
import ai.koog.prompt.message.Message
9+
import ai.koog.prompt.message.MessagePart
910
import ai.koog.prompt.message.ResponseMetaInfo
1011
import ai.koog.prompt.streaming.StreamFrame
1112
import ai.koog.prompt.streaming.toStreamFrames
@@ -30,27 +31,31 @@ object CalculatorChatExecutor : PromptExecutor() {
3031
override fun now(): Instant = Instant.parse("2023-01-01T00:00:00Z")
3132
}
3233

33-
override suspend fun execute(prompt: Prompt, model: LLModel, tools: List<ToolDescriptor>): List<Message.Assistant> {
34-
val input = prompt.messages.filterIsInstance<Message.User>().joinToString("\n") { it.content }
34+
override suspend fun execute(prompt: Prompt, model: LLModel, tools: List<ToolDescriptor>): Message.Assistant {
35+
val input = prompt.messages.filterIsInstance<Message.User>()
36+
.joinToString("\n") { msg -> msg.parts.filterIsInstance<MessagePart.Text>().joinToString("\n") { it.text } }
3537
val numbers = input.split(Regex("[^0-9.]")).filter { it.isNotEmpty() }.map { it.toFloat() }
36-
val result = when {
38+
return when {
3739
plusAliases.any { it in input } && tools.contains(CalculatorTools.PlusTool.descriptor) -> {
38-
Message.Tool.Call(
39-
id = "1",
40-
tool = CalculatorTools.PlusTool.name,
41-
content = json.encodeToString(
42-
buildJsonObject {
43-
put("a", numbers[0])
44-
put("b", numbers[1])
45-
}
40+
Message.Assistant(
41+
parts = listOf(
42+
MessagePart.Tool.Call(
43+
id = "1",
44+
tool = CalculatorTools.PlusTool.name,
45+
args = json.encodeToString(
46+
buildJsonObject {
47+
put("a", numbers[0])
48+
put("b", numbers[1])
49+
}
50+
)
51+
)
4652
),
4753
metaInfo = ResponseMetaInfo.create(testClock)
4854
)
4955
}
5056

5157
else -> Message.Assistant("Unknown operation", metaInfo = ResponseMetaInfo.create(testClock))
5258
}
53-
return listOf(result)
5459
}
5560

5661
override fun executeStreaming(
@@ -60,7 +65,7 @@ object CalculatorChatExecutor : PromptExecutor() {
6065
): Flow<StreamFrame> =
6166
flow {
6267
try {
63-
execute(prompt, model, tools).toStreamFrames().forEach { emit(it) }
68+
execute(prompt, model, tools).toStreamFrames().forEach { frame -> emit(frame) }
6469
} catch (t: CancellationException) {
6570
throw t
6671
} catch (t: Throwable) {

‎agents/agents-core/src/commonTest/kotlin/ai/koog/agents/core/agent/AIAgentGenericTypesTest.kt‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import ai.koog.agents.testing.tools.getMockExecutor
77
import ai.koog.prompt.dsl.prompt
88
import ai.koog.prompt.executor.ollama.client.OllamaModels
99
import ai.koog.prompt.message.Message
10+
import ai.koog.prompt.message.MessagePart
1011
import ai.koog.serialization.kotlinx.KotlinxSerializer
1112
import kotlinx.coroutines.test.runTest
1213
import kotlin.test.Test
@@ -29,7 +30,7 @@ class AIAgentGenericTypesTest {
2930

3031
val customStrategy = strategy<CustomInput, CustomOutput>("custom-strategy") {
3132
val processInput = { input: CustomInput -> input.query }
32-
val processOutput = { output: Message.Assistant -> CustomOutput(result = output.content, confidence = 0.95) }
33+
val processOutput = { output: Message.Assistant -> CustomOutput(result = output.parts.filterIsInstance<MessagePart.Text>().first().text, confidence = 0.95) }
3334

3435
val callLLM by nodeLLMRequest()
3536

@@ -61,8 +62,8 @@ class AIAgentGenericTypesTest {
6162
val convertToString = { input: Int -> "Is $input an even number?" }
6263

6364
val parseResponse = { output: Message.Assistant ->
64-
output.content.contains("yes", ignoreCase = true) ||
65-
output.content.contains("even", ignoreCase = true)
65+
val text = output.parts.filterIsInstance<MessagePart.Text>().firstOrNull()?.text ?: ""
66+
text.contains("yes", ignoreCase = true) || text.contains("even", ignoreCase = true)
6667
}
6768

6869
val callLLM by nodeLLMRequest()

‎agents/agents-core/src/commonTest/kotlin/ai/koog/agents/core/agent/FunctionalAIAgentTest.kt‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import ai.koog.prompt.executor.clients.anthropic.AnthropicModels
1717
import ai.koog.prompt.executor.clients.google.GoogleModels
1818
import ai.koog.prompt.executor.clients.openai.OpenAIModels
1919
import ai.koog.prompt.executor.ollama.client.OllamaModels
20+
import ai.koog.prompt.message.MessagePart
2021
import ai.koog.serialization.kotlinx.KotlinxSerializer
2122
import kotlinx.coroutines.test.runTest
2223
import kotlinx.serialization.Serializable
@@ -56,13 +57,13 @@ class FunctionalAIAgentTest {
5657
strategy = functionalStrategy { inputParam ->
5758
var responses = requestLLM(inputParam)
5859

59-
while (responses.containsToolCalls()) {
60+
while (responses.parts.any { it is MessagePart.Tool.Call }) {
6061
val tools = extractToolCalls(responses)
6162
val results = executeMultipleTools(tools)
6263
responses = sendMultipleToolResults(results)
6364
}
6465

65-
responses.single().asAssistantMessage().content
66+
responses.parts.filterIsInstance<MessagePart.Text>().first().text
6667
},
6768
llmModel = OllamaModels.Meta.LLAMA_3_2,
6869
toolRegistry = testToolRegistry
@@ -102,7 +103,7 @@ class FunctionalAIAgentTest {
102103
appendPrompt { user(inputParam) }
103104
this.requestLLM()
104105
}
105-
resp.content
106+
resp.parts.filterIsInstance<MessagePart.Text>().first().text
106107
},
107108
llmModel = OllamaModels.Meta.LLAMA_3_2,
108109
toolRegistry = testToolRegistry,
@@ -140,13 +141,13 @@ class FunctionalAIAgentTest {
140141
strategy = functionalStrategy { inputParam: String ->
141142
var responses = requestLLM(inputParam)
142143

143-
while (responses.containsToolCalls()) {
144+
while (responses.parts.any { it is MessagePart.Tool.Call }) {
144145
val tools = extractToolCalls(responses)
145146
val results = executeMultipleTools(tools)
146147
responses = sendMultipleToolResults(results)
147148
}
148149

149-
responses.single().asAssistantMessage().content
150+
responses.parts.filterIsInstance<MessagePart.Text>().first().text
150151
}
151152
) {
152153
install(EventHandler) {

‎agents/agents-core/src/commonTest/kotlin/ai/koog/agents/core/dsl/extension/TestLLMExecutor.kt‎

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,15 +6,20 @@ import ai.koog.prompt.dsl.Prompt
66
import ai.koog.prompt.executor.model.PromptExecutor
77
import ai.koog.prompt.llm.LLModel
88
import ai.koog.prompt.message.Message
9+
import ai.koog.prompt.message.MessagePart
910
import ai.koog.prompt.message.ResponseMetaInfo
1011
import ai.koog.prompt.streaming.StreamFrame
1112
import ai.koog.prompt.streaming.toStreamFrames
13+
import io.github.oshai.kotlinlogging.KotlinLogging
1214
import kotlinx.coroutines.flow.Flow
1315
import kotlinx.coroutines.flow.flow
1416
import kotlin.time.Clock
1517
import kotlin.time.Instant
1618

1719
class TestLLMExecutor : PromptExecutor() {
20+
21+
private val logger = KotlinLogging.logger {}
22+
1823
companion object {
1924
val testClock: Clock = object : Clock {
2025
override fun now(): Instant = Instant.parse("2023-01-01T00:00:00Z")
@@ -36,8 +41,8 @@ class TestLLMExecutor : PromptExecutor() {
3641
messages.clear()
3742
}
3843

39-
override suspend fun execute(prompt: Prompt, model: LLModel, tools: List<ToolDescriptor>): List<Message.Assistant> {
40-
return listOf(handlePrompt(prompt))
44+
override suspend fun execute(prompt: Prompt, model: LLModel, tools: List<ToolDescriptor>): Message.Assistant {
45+
return handlePrompt(prompt)
4146
}
4247

4348
override fun executeStreaming(
@@ -56,13 +61,19 @@ class TestLLMExecutor : PromptExecutor() {
5661
}
5762

5863
private fun handlePrompt(prompt: Prompt): Message.Assistant {
59-
prompt.messages.forEach { println("[DEBUG_LOG] Message: ${it.content}") }
64+
prompt.messages.forEach { logger.debug { "Message: $it" } }
6065

6166
// Store all messages for later inspection
6267
messages.addAll(prompt.messages)
6368

6469
// For compression test, return a TLDR summary
65-
if (prompt.messages.any { it.content.contains("Create a comprehensive summary of this conversation") }) {
70+
if (prompt.messages.any {
71+
it.parts.any { part ->
72+
part is MessagePart.Text &&
73+
part.text.contains("Create a comprehensive summary of this conversation")
74+
}
75+
}
76+
) {
6677
tldrCount++
6778
val tldrResponse = Message.Assistant(
6879
"TLDR #$tldrCount: Summary of conversation history",

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ internal class ChoiceEvent(
2828
addBodyField(EventBodyFields.FinishReason(GenAIAttributes.Response.FinishReasonType.ToolCalls.id))
2929
} else {
3030
message.finishReason?.let { reason ->
31-
addBodyField(EventBodyFields.FinishReason(reason.name))
31+
addBodyField(EventBodyFields.FinishReason(reason))
3232
}
3333
val textContent = message.parts.filterIsInstance<MessagePart.Text>().joinToString("\n") { it.text }
3434
addBodyField(EventBodyFields.Message(role = message.role, content = textContent))

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

Lines changed: 17 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import ai.koog.agents.features.opentelemetry.span.GenAIAgentSpan
99
import ai.koog.agents.features.opentelemetry.span.SpanEndStatus
1010
import ai.koog.http.client.KoogHttpClientException
1111
import ai.koog.prompt.message.Message
12+
import ai.koog.prompt.message.MessagePart
1213
import io.opentelemetry.api.trace.Span
1314
import io.opentelemetry.api.trace.SpanBuilder
1415
import io.opentelemetry.api.trace.StatusCode
@@ -88,43 +89,39 @@ internal fun List<Message>.systemMessages(): List<Message.System> =
8889
filterIsInstance<Message.System>()
8990

9091
/**
91-
* Returns the last [Message.Response] in this list, or `null` if no response has been produced yet.
92+
* Returns the last [Message.Assistant] in this list, or `null` if no response has been produced yet.
9293
*/
93-
internal fun List<Message>.lastResponse(): Message.Response? =
94-
filterIsInstance<Message.Response>().lastOrNull()
94+
internal fun List<Message>.lastAssistant(): Message.Assistant? =
95+
filterIsInstance<Message.Assistant>().lastOrNull()
9596

9697
/**
97-
* Sum of `metaInfo.inputTokensCount` across all [Message.Response] entries, for the `gen_ai.usage.input_tokens` attribute.
98+
* Sum of `metaInfo.inputTokensCount` across all [Message.Assistant] entries, for the `gen_ai.usage.input_tokens` attribute.
9899
*/
99100
internal fun List<Message>.sumInputTokens(): Int =
100-
filterIsInstance<Message.Response>().sumOf { it.metaInfo.inputTokensCount ?: 0 }
101+
filterIsInstance<Message.Assistant>().sumOf { it.metaInfo.inputTokensCount ?: 0 }
101102

102103
/**
103-
* Sum of `metaInfo.outputTokensCount` across all [Message.Response] entries, for the `gen_ai.usage.output_tokens` attribute.
104+
* Sum of `metaInfo.outputTokensCount` across all [Message.Assistant] entries, for the `gen_ai.usage.output_tokens` attribute.
104105
*/
105106
internal fun List<Message>.sumOutputTokens(): Int =
106-
filterIsInstance<Message.Response>().sumOf { it.metaInfo.outputTokensCount ?: 0 }
107+
filterIsInstance<Message.Assistant>().sumOf { it.metaInfo.outputTokensCount ?: 0 }
107108

108109
/**
109110
* Merges the per-response `metaInfo.metadata` maps into a single flat map, for the `gen_ai.response.metadata` attribute.
110111
* Later responses overwrite earlier ones on key collision.
111112
*/
112-
internal fun List<Message>.mergedResponseMetadata(): Map<String, JsonElement> =
113-
filterIsInstance<Message.Response>()
113+
internal fun List<Message>.mergedAssistantMetadata(): Map<String, JsonElement> =
114+
filterIsInstance<Message.Assistant>()
114115
.mapNotNull { it.metaInfo.metadata }
115116
.fold(mutableMapOf()) { acc, m -> acc.apply { putAll(m) } }
116117

117118
/**
118-
* Maps a [Message.Response] to its `FinishReasonType` for the `gen_ai.response.finish_reasons` attribute.
119-
* Exhaustive by design (no `else`) so a new [Message.Response] subtype surfaces as a compile error.
119+
* Maps a [Message.Assistant] to its `FinishReasonType` for the `gen_ai.response.finish_reasons` attribute.
120+
* Exhaustive by design (no `else`) so a new [Message.Assistant] subtype surfaces as a compile error.
120121
*/
121-
internal fun Message.Response.toFinishReason(): GenAIAttributes.Response.FinishReasonType =
122-
when (this) {
123-
is Message.Assistant,
124-
is Message.Reasoning -> {
125-
GenAIAttributes.Response.FinishReasonType.Stop
126-
}
127-
is Message.Tool.Call -> {
128-
GenAIAttributes.Response.FinishReasonType.ToolCalls
129-
}
122+
internal fun Message.Assistant.toFinishReason(): GenAIAttributes.Response.FinishReasonType =
123+
if (parts.any { it is MessagePart.Tool.Call }) {
124+
GenAIAttributes.Response.FinishReasonType.ToolCalls
125+
} else {
126+
GenAIAttributes.Response.FinishReasonType.Stop
130127
}

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ import ai.koog.agents.features.opentelemetry.event.ModerationResponseEvent
2020
import ai.koog.agents.features.opentelemetry.event.SystemMessageEvent
2121
import ai.koog.agents.features.opentelemetry.event.ToolMessageEvent
2222
import ai.koog.agents.features.opentelemetry.event.UserMessageEvent
23-
import ai.koog.agents.features.opentelemetry.extension.lastResponse
23+
import ai.koog.agents.features.opentelemetry.extension.lastAssistant
2424
import ai.koog.agents.features.opentelemetry.extension.toFinishReason
2525
import ai.koog.agents.features.opentelemetry.integration.SpanAdapter
2626
import ai.koog.agents.features.opentelemetry.integration.mcp.McpMethod
@@ -358,7 +358,7 @@ public class OpenTelemetry {
358358
spanType = SpanType.INVOKE_AGENT
359359
) ?: return@intercept
360360

361-
eventContext.context.llm.prompt.messages.lastResponse()?.let { response ->
361+
eventContext.context.llm.prompt.messages.lastAssistant()?.let { response ->
362362
invokeAgentSpan.addAttribute(
363363
GenAIAttributes.Response.FinishReasons(reasons = listOf(response.toFinishReason()))
364364
)
@@ -664,7 +664,7 @@ public class OpenTelemetry {
664664
spanAdapter?.onBeforeSpanFinished(inferenceSpan)
665665
endInferenceSpan(
666666
span = inferenceSpan,
667-
messages = emptyList(),
667+
message = null,
668668
model = eventContext.model,
669669
verbose = config.isVerbose,
670670
error = eventContext.error

0 commit comments

Comments
 (0)