Skip to content

Add @Retry annotation #230

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
aozherelyeva merged 3 commits into from
Jun 6, 2025
Merged

Add @Retry annotation #230

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
a9f0076
Refactor error handling for Google API exceptions
aozherelyeva Jun 5, 2025
5edd5e7
Introduce @Retry annotation for flaky tests on the test class level
aozherelyeva Jun 5, 2025
2fa60ab
Replace `executeWithRetry` and `runWithRetry` with `@Retry`
aozherelyeva Jun 6, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
231 changes: 118 additions & 113 deletions .../src/jvmTest/kotlin/ai/koog/integration/tests/MultipleLLMPromptExecutorIntegrationTest.kt
View file
Open in desktop

Large diffs are not rendered by default.

9 changes: 5 additions & 4 deletions ...c/jvmTest/kotlin/ai/koog/integration/tests/MultipleSystemMessagesPromptIntegrationTest.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
package ai.koog.integration.tests

import ai.koog.integration.tests.utils.TestUtils.executeWithRetry
import ai.koog.integration.tests.utils.TestUtils.readTestAnthropicKeyFromEnv
import ai.koog.integration.tests.utils.TestUtils.readTestGoogleAIKeyFromEnv
import ai.koog.integration.tests.utils.TestUtils.readTestOpenAIKeyFromEnv
import ai.koog.integration.tests.utils.annotations.Retry
import ai.koog.prompt.dsl.prompt
import ai.koog.prompt.executor.model.PromptExecutorExt.execute
import ai.koog.prompt.executor.clients.anthropic.AnthropicLLMClient
Expand All @@ -23,6 +23,7 @@ class MultipleSystemMessagesPromptIntegrationTest {
private val anthropicApiKey = readTestAnthropicKeyFromEnv()
private val googleApiKey = readTestGoogleAIKeyFromEnv()

@Retry(3)
@Test
fun integration_testMultipleSystemMessages() = runBlocking {
val openAIClient = OpenAILLMClient(openAIApiKey)
Expand All @@ -46,9 +47,9 @@ class MultipleSystemMessagesPromptIntegrationTest {
val modelAnthropic = AnthropicModels.Haiku_3_5
val modelGemini = GoogleModels.Gemini2_0Flash

val responseOpenAI = executeWithRetry { executor.execute(prompt, modelOpenAI) }
val responseAnthropic = executeWithRetry { executor.execute(prompt, modelAnthropic) }
val responseGemini = executeWithRetry { executor.execute(prompt, modelGemini) }
val responseOpenAI = executor.execute(prompt, modelOpenAI)
val responseAnthropic = executor.execute(prompt, modelAnthropic)
val responseGemini = executor.execute(prompt, modelGemini)

assertTrue(responseOpenAI.content.isNotEmpty(), "OpenAI response should not be empty")
assertTrue(responseAnthropic.content.isNotEmpty(), "Anthropic response should not be empty")
Expand Down
5 changes: 3 additions & 2 deletions integration-tests/src/jvmTest/kotlin/ai/koog/integration/tests/OllamaAgentIntegrationTest.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ import ai.koog.agents.core.dsl.builder.strategy
import ai.koog.agents.core.dsl.extension.*
import ai.koog.agents.core.tools.ToolRegistry
import ai.koog.agents.features.eventHandler.feature.EventHandler
import ai.koog.integration.tests.utils.TestUtils.runWithRetry
import ai.koog.integration.tests.utils.annotations.Retry
import ai.koog.prompt.dsl.prompt
import ai.koog.prompt.executor.model.PromptExecutor
import ai.koog.prompt.llm.OllamaModels
Expand Down Expand Up @@ -174,13 +174,14 @@ class OllamaAgentIntegrationTest {
}
}

@Retry(3)
@Test
fun ollama_testAgentClearContext() = runTest(timeout = 600.seconds) {
val strategy = createTestStrategy()
val toolRegistry = createToolRegistry()
val agent = createAgent(executor, strategy, toolRegistry)

val result = runWithRetry { agent.runAndGetResult("What is the capital of France?") }
val result = agent.runAndGetResult("What is the capital of France?")

assertNotNull(result, "Result should not be empty")
assertTrue(result.isNotEmpty(), "Result should not be empty")
Expand Down
8 changes: 4 additions & 4 deletions ...on-tests/src/jvmTest/kotlin/ai/koog/integration/tests/OllamaSimpleAgentIntegrationTest.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ import ai.koog.agents.core.tools.ToolRegistry
import ai.koog.agents.ext.tool.SayToUser
import ai.koog.agents.features.eventHandler.feature.EventHandler
import ai.koog.agents.features.eventHandler.feature.EventHandlerConfig
import ai.koog.integration.tests.utils.TestUtils.runWithRetry
import ai.koog.integration.tests.utils.annotations.Retry
import kotlinx.coroutines.test.runTest
import org.junit.jupiter.api.extension.ExtendWith
import kotlin.test.AfterTest
Expand Down Expand Up @@ -87,6 +87,7 @@ class OllamaSimpleAgentIntegrationTest {
actualToolCalls.clear()
}

@Retry(3)
@Test
fun ollama_simpleTest() = runTest(timeout = 600.seconds) {
val toolRegistry = ToolRegistry.Companion {
Expand Down Expand Up @@ -119,9 +120,8 @@ class OllamaSimpleAgentIntegrationTest {
installFeatures = { install(EventHandler.Feature, eventHandlerConfig) }
)

runWithRetry {
agent.run("Give me top 10 books of the all time.")
}

agent.run("Give me top 10 books of the all time.")

assertTrue(actualToolCalls.isNotEmpty(), "No tools were called for model")
assertTrue(
Expand Down
13 changes: 6 additions & 7 deletions integration-tests/src/jvmTest/kotlin/ai/koog/integration/tests/SimpleAgentIntegrationTest.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ import ai.koog.integration.tests.utils.TestUtils.CalculatorTool
import ai.koog.integration.tests.utils.TestUtils.readTestAnthropicKeyFromEnv
import ai.koog.integration.tests.utils.TestUtils.readTestGoogleAIKeyFromEnv
import ai.koog.integration.tests.utils.TestUtils.readTestOpenAIKeyFromEnv
import ai.koog.integration.tests.utils.TestUtils.runWithRetry
import ai.koog.integration.tests.utils.annotations.Retry
import ai.koog.prompt.executor.clients.google.GoogleModels
import ai.koog.prompt.executor.clients.openai.OpenAIModels
import ai.koog.prompt.executor.llms.all.simpleAnthropicExecutor
Expand Down Expand Up @@ -116,6 +116,7 @@ class SimpleAgentIntegrationTest {
results.clear()
}

@Retry(3)
@ParameterizedTest
@MethodSource("openAIModels", "anthropicModels", "googleModels")
fun integration_AIAgentShouldNotCallToolsByDefault(model: LLModel) = runBlocking {
Expand All @@ -134,14 +135,13 @@ class SimpleAgentIntegrationTest {
installFeatures = { install(EventHandler.Feature, eventHandlerConfig) }
)

runWithRetry {
agent.run("Repeat what I say: hello, I'm good.")
}
agent.run("Repeat what I say: hello, I'm good.")

// by default, AIAgent has no tools underneath
assertTrue(actualToolCalls.isEmpty(), "No tools should be called for model $model")
}

@Retry(3)
@ParameterizedTest
@MethodSource("openAIModels", "anthropicModels", "googleModels")
fun integration_AIAgentShouldCallCustomTool(model: LLModel) = runBlocking {
Expand Down Expand Up @@ -172,9 +172,8 @@ class SimpleAgentIntegrationTest {
installFeatures = { install(EventHandler.Feature, eventHandlerConfig) }
)

runWithRetry {
agent.run("How much is 3 times 5?")
}

agent.run("How much is 3 times 5?")

assertTrue(actualToolCalls.isNotEmpty(), "No tools were called for model $model")
assertTrue(
Expand Down
Loading
Loading