Skip to content

Commit 60f8ed3

Browse files
skarpovdevskarpovdev
committed
Reduce duplication + clean up
Signed-off-by: Sergey Karpov <[email protected]>
1 parent 1945e7a commit 60f8ed3

File tree

2 files changed

+63
-93
lines changed

2 files changed

+63
-93
lines changed

‎integration-tests/src/jvmTest/java/ai/koog/integration/tests/agent/JavaApiIntegrationTest.java‎

Lines changed: 63 additions & 86 deletions
Original file line numberDiff line numberDiff line change
@@ -48,49 +48,63 @@ public void cleanup() throws Exception {
4848
resourcesToClose.clear();
4949
}
5050

51+
private void assertValidResponse(List<Message.Response> responses) {
52+
assertNotNull(responses);
53+
assertFalse(responses.isEmpty());
54+
assertInstanceOf(Message.Assistant.class, responses.get(0));
55+
String content = JavaInteropUtils.getAssistantContent((Message.Assistant) responses.get(0));
56+
assertFalse(content.isEmpty());
57+
}
58+
59+
private String getAssistantContentOrDefault(Message.Response response, String defaultValue) {
60+
if (response instanceof Message.Assistant) {
61+
return JavaInteropUtils.getAssistantContent((Message.Assistant) response);
62+
}
63+
return defaultValue;
64+
}
65+
66+
private AIAgent<String, String> buildAgentWithCalculator(MultiLLMPromptExecutor executor, LLModel model) {
67+
JavaInteropUtils.CalculatorTools calculator = new JavaInteropUtils.CalculatorTools();
68+
ToolRegistry toolRegistry = JavaInteropUtils.createToolRegistry(calculator);
69+
70+
return AIAgent.builder()
71+
.promptExecutor(executor)
72+
.llmModel(model)
73+
.systemPrompt("You are a calculator assistant. You MUST use the add tool to perform additions. ALWAYS call the tool.")
74+
.toolRegistry(toolRegistry)
75+
.build();
76+
}
77+
5178
@Test
5279
public void integration_testOpenAILLMClient() {
53-
String apiKey = TestCredentials.INSTANCE.readTestOpenAIKeyFromEnv();
54-
OpenAILLMClient client = JavaInteropUtils.createOpenAIClient(apiKey);
80+
OpenAILLMClient client = JavaInteropUtils.createOpenAIClient(TestCredentials.INSTANCE.readTestOpenAIKeyFromEnv());
5581
resourcesToClose.add((AutoCloseable) client);
5682

5783
assertEquals(LLMProvider.OpenAI.INSTANCE, client.llmProvider());
5884

5985
Prompt prompt = JavaInteropUtils.buildSimplePrompt("test-openai", "You are a helpful assistant.", "Say 'Hello from OpenAI'");
6086
List<Message.Response> responses = JavaInteropUtils.executeClientBlocking(client, prompt, OpenAIModels.Chat.GPT4o, Collections.emptyList());
6187

62-
assertNotNull(responses);
63-
assertFalse(responses.isEmpty());
64-
assertInstanceOf(Message.Assistant.class, responses.get(0));
65-
String content = JavaInteropUtils.getAssistantContent((Message.Assistant) responses.get(0));
66-
assertFalse(content.isEmpty());
88+
assertValidResponse(responses);
6789
}
6890

6991
@Test
7092
public void integration_testAnthropicLLMClient() {
71-
String apiKey = TestCredentials.INSTANCE.readTestAnthropicKeyFromEnv();
72-
AnthropicLLMClient client = JavaInteropUtils.createAnthropicClient(apiKey);
93+
AnthropicLLMClient client = JavaInteropUtils.createAnthropicClient(TestCredentials.INSTANCE.readTestAnthropicKeyFromEnv());
7394
resourcesToClose.add((AutoCloseable) client);
7495

7596
assertEquals(LLMProvider.Anthropic.INSTANCE, client.llmProvider());
7697

7798
Prompt prompt = JavaInteropUtils.buildSimplePrompt("test-anthropic", "You are a helpful assistant.", "Say 'Hello from Anthropic'");
7899
List<Message.Response> responses = JavaInteropUtils.executeClientBlocking(client, prompt, AnthropicModels.Haiku_4_5, Collections.emptyList());
79100

80-
assertNotNull(responses);
81-
assertFalse(responses.isEmpty());
82-
assertInstanceOf(Message.Assistant.class, responses.get(0));
83-
String content = JavaInteropUtils.getAssistantContent((Message.Assistant) responses.get(0));
84-
assertFalse(content.isEmpty());
101+
assertValidResponse(responses);
85102
}
86103

87104
@Test
88105
public void integration_testMultiLLMPromptExecutor() {
89-
String openAIKey = TestCredentials.INSTANCE.readTestOpenAIKeyFromEnv();
90-
String anthropicKey = TestCredentials.INSTANCE.readTestAnthropicKeyFromEnv();
91-
92-
OpenAILLMClient openAIClient = JavaInteropUtils.createOpenAIClient(openAIKey);
93-
AnthropicLLMClient anthropicClient = JavaInteropUtils.createAnthropicClient(anthropicKey);
106+
OpenAILLMClient openAIClient = JavaInteropUtils.createOpenAIClient(TestCredentials.INSTANCE.readTestOpenAIKeyFromEnv());
107+
AnthropicLLMClient anthropicClient = JavaInteropUtils.createAnthropicClient(TestCredentials.INSTANCE.readTestAnthropicKeyFromEnv());
94108

95109
resourcesToClose.add((AutoCloseable) openAIClient);
96110
resourcesToClose.add((AutoCloseable) anthropicClient);
@@ -135,18 +149,10 @@ public void integration_testBuilderBasicUsage(LLModel model) {
135149
public void integration_testBuilderWithToolRegistry(LLModel model) {
136150
Models.assumeAvailable(model.getProvider());
137151

138-
MultiLLMPromptExecutor executor = createExecutor(model);
139-
140-
JavaInteropUtils.CalculatorTools calculator = new JavaInteropUtils.CalculatorTools();
141-
142-
ToolRegistry toolRegistry = JavaInteropUtils.createToolRegistry(calculator);
143-
144-
AIAgent<String, String> agent = AIAgent.builder()
145-
.promptExecutor(executor)
146-
.llmModel(model)
147-
.systemPrompt("You are a calculator assistant. You MUST use the add tool to perform additions. ALWAYS call the tool.")
148-
.toolRegistry(toolRegistry)
149-
.build();
152+
AIAgent<String, String> agent = buildAgentWithCalculator(
153+
createExecutor(model),
154+
model
155+
);
150156

151157
String result = runBlocking(continuation ->
152158
agent.run("What is 15 + 27?", continuation)
@@ -162,18 +168,15 @@ public void integration_testBuilderWithToolRegistry(LLModel model) {
162168
public void integration_testEventHandler(LLModel model) {
163169
Models.assumeAvailable(model.getProvider());
164170

165-
MultiLLMPromptExecutor executor = createExecutor(model);
166-
167171
AtomicBoolean agentStarted = new AtomicBoolean(false);
168172
AtomicBoolean agentCompleted = new AtomicBoolean(false);
169173
AtomicInteger llmCallsCount = new AtomicInteger(0);
170174

171175
JavaInteropUtils.CalculatorTools calculator = new JavaInteropUtils.CalculatorTools();
172-
173176
ToolRegistry toolRegistry = JavaInteropUtils.createToolRegistry(calculator);
174177

175178
AIAgent<String, String> agent = AIAgent.builder()
176-
.promptExecutor(executor)
179+
.promptExecutor(createExecutor(model))
177180
.llmModel(model)
178181
.systemPrompt("You are a calculator. Use the add tool when needed.")
179182
.toolRegistry(toolRegistry)
@@ -199,20 +202,17 @@ public void integration_testEventHandler(LLModel model) {
199202
public void integration_testSimpleFunctionalStrategy(LLModel model) {
200203
Models.assumeAvailable(model.getProvider());
201204

202-
MultiLLMPromptExecutor executor = createExecutor(model);
203-
204205
AIAgent<String, String> agent = JavaInteropUtils.buildFunctionalAgent(
205206
JavaInteropUtils.createAgentBuilder()
206-
.promptExecutor(executor)
207+
.promptExecutor(createExecutor(model))
207208
.llmModel(model)
208209
.systemPrompt("You are a helpful assistant.")
209-
.functionalStrategy((context, input) -> {
210-
Message.Response response = JavaInteropUtils.requestLLMBlocking(context, input, true);
211-
if (response instanceof Message.Assistant) {
212-
return JavaInteropUtils.getAssistantContent((Message.Assistant) response);
213-
}
214-
return "Unexpected response type";
215-
})
210+
.functionalStrategy((context, input) ->
211+
getAssistantContentOrDefault(
212+
JavaInteropUtils.requestLLMBlocking(context, input, true),
213+
"Unexpected response type"
214+
)
215+
)
216216
);
217217

218218
String result = JavaInteropUtils.runAgentBlocking(agent, "Say hello");
@@ -227,31 +227,22 @@ public void integration_testSimpleFunctionalStrategy(LLModel model) {
227227
public void integration_testMultiStepFunctionalStrategy(LLModel model) {
228228
Models.assumeAvailable(model.getProvider());
229229

230-
MultiLLMPromptExecutor executor = createExecutor(model);
231-
232230
AIAgent<String, String> agent = JavaInteropUtils.buildFunctionalAgent(
233231
JavaInteropUtils.createAgentBuilder()
234-
.promptExecutor(executor)
232+
.promptExecutor(createExecutor(model))
235233
.llmModel(model)
236234
.systemPrompt("You are a helpful assistant.")
237235
.functionalStrategy((context, input) -> {
238236
Message.Response response1 = JavaInteropUtils.requestLLMBlocking(context, "First step: " + input, true);
239-
240-
String step1Result = "";
241-
if (response1 instanceof Message.Assistant) {
242-
step1Result = JavaInteropUtils.getAssistantContent((Message.Assistant) response1);
243-
}
237+
String step1Result = getAssistantContentOrDefault(response1, "");
244238

245239
Message.Response response2 = JavaInteropUtils.requestLLMBlocking(
246240
context,
247241
"Second step, previous result was: " + step1Result,
248242
true
249243
);
250244

251-
if (response2 instanceof Message.Assistant) {
252-
return JavaInteropUtils.getAssistantContent((Message.Assistant) response2);
253-
}
254-
return "Unexpected response type";
245+
return getAssistantContentOrDefault(response2, "Unexpected response type");
255246
})
256247
);
257248

@@ -266,15 +257,12 @@ public void integration_testMultiStepFunctionalStrategy(LLModel model) {
266257
public void integration_testFunctionalStrategyWithManualToolHandling(LLModel model) {
267258
Models.assumeAvailable(model.getProvider());
268259

269-
MultiLLMPromptExecutor executor = createExecutor(model);
270-
271260
JavaInteropUtils.CalculatorTools calculator = new JavaInteropUtils.CalculatorTools();
272-
273261
ToolRegistry toolRegistry = JavaInteropUtils.createToolRegistry(calculator);
274262

275263
AIAgent<String, String> agent = JavaInteropUtils.buildFunctionalAgent(
276264
JavaInteropUtils.createAgentBuilder()
277-
.promptExecutor(executor)
265+
.promptExecutor(createExecutor(model))
278266
.llmModel(model)
279267
.systemPrompt("You are a calculator. Use the add tool to perform calculations.")
280268
.toolRegistry(toolRegistry)
@@ -285,14 +273,11 @@ public void integration_testFunctionalStrategyWithManualToolHandling(LLModel mod
285273
true
286274
);
287275

288-
int iterations = 0;
289276
int maxIterations = 5;
290-
291-
while (currentResponse instanceof Message.Tool.Call && iterations < maxIterations) {
277+
for (int i = 0; i < maxIterations && currentResponse instanceof Message.Tool.Call; i++) {
292278
Message.Tool.Call toolCall = (Message.Tool.Call) currentResponse;
293279
ReceivedToolResult toolResult = JavaInteropUtils.executeToolBlocking(context, toolCall);
294280
currentResponse = JavaInteropUtils.sendToolResultBlocking(context, toolResult);
295-
iterations++;
296281
}
297282

298283
if (currentResponse instanceof Message.Assistant) {
@@ -373,22 +358,19 @@ public void integration_testSubtask(LLModel model) {
373358
public void integration_testCustomStrategyWithRetry(LLModel model) {
374359
Models.assumeAvailable(model.getProvider());
375360

376-
MultiLLMPromptExecutor executor = createExecutor(model);
377-
378361
AIAgent<String, String> agent = JavaInteropUtils.buildFunctionalAgent(
379362
JavaInteropUtils.createAgentBuilder()
380-
.promptExecutor(executor)
363+
.promptExecutor(createExecutor(model))
381364
.llmModel(model)
382365
.systemPrompt("You are a helpful assistant.")
383366
.functionalStrategy((context, input) -> {
384-
int maxRetries = 3;
385-
for (int i = 0; i < maxRetries; i++) {
386-
Message.Response response = JavaInteropUtils.requestLLMBlocking(context, input, true);
387-
if (response instanceof Message.Assistant) {
388-
String result = JavaInteropUtils.getAssistantContent((Message.Assistant) response);
389-
if (!result.isEmpty()) {
390-
return result;
391-
}
367+
for (int i = 0; i < 3; i++) {
368+
String result = getAssistantContentOrDefault(
369+
JavaInteropUtils.requestLLMBlocking(context, input, true),
370+
""
371+
);
372+
if (!result.isEmpty()) {
373+
return result;
392374
}
393375
}
394376
return "Failed after retries";
@@ -406,11 +388,9 @@ public void integration_testCustomStrategyWithRetry(LLModel model) {
406388
public void integration_testCustomStrategyWithValidation(LLModel model) {
407389
Models.assumeAvailable(model.getProvider());
408390

409-
MultiLLMPromptExecutor executor = createExecutor(model);
410-
411391
AIAgent<String, String> agent = JavaInteropUtils.buildFunctionalAgent(
412392
JavaInteropUtils.createAgentBuilder()
413-
.promptExecutor(executor)
393+
.promptExecutor(createExecutor(model))
414394
.llmModel(model)
415395
.systemPrompt("You are a helpful assistant that generates JSON.")
416396
.functionalStrategy((context, input) -> {
@@ -420,14 +400,11 @@ public void integration_testCustomStrategyWithValidation(LLModel model) {
420400
true
421401
);
422402

423-
if (response instanceof Message.Assistant) {
424-
String content = JavaInteropUtils.getAssistantContent((Message.Assistant) response);
425-
if (content.contains("status") && content.contains("success")) {
426-
return content;
427-
}
428-
return "Validation failed: response doesn't contain expected fields";
403+
String content = getAssistantContentOrDefault(response, "Unexpected response type");
404+
if (content.contains("status") && content.contains("success")) {
405+
return content;
429406
}
430-
return "Unexpected response type";
407+
return "Validation failed: response doesn't contain expected fields";
431408
})
432409
);
433410

‎integration-tests/src/jvmTest/kotlin/ai/koog/integration/tests/utils/JavaInteropUtils.kt‎

Lines changed: 0 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -140,13 +140,6 @@ object JavaInteropUtils {
140140
@LLMDescription(description = "Second number") b: Int
141141
): Int = a + b
142142

143-
@ai.koog.agents.core.tools.annotations.Tool
144-
@LLMDescription(description = "Multiplies two numbers")
145-
fun multiply(
146-
@LLMDescription(description = "First number") a: Int,
147-
@LLMDescription(description = "Second number") b: Int
148-
): Int = a * b
149-
150143
fun getAddTool(): Tool<*, *> {
151144
return createToolRegistry(this).tools.first { it.name == "add" }
152145
}

0 commit comments

Comments
 (0)