Skip to content

Commit d2b34f2

Browse files
RizzenRizzen
authored
fixing QWQ model, introduced error handling for error responses in Ollama (#274)
1 parent bedaccc commit d2b34f2

File tree

7 files changed

+162
-60
lines changed

7 files changed

+162
-60
lines changed

‎examples/src/main/kotlin/ai/koog/agents/example/calculator/Calculator.kt‎

Lines changed: 1 addition & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,6 @@ package ai.koog.agents.example.calculator
44

55
import ai.koog.agents.core.agent.AIAgent
66
import ai.koog.agents.core.agent.config.AIAgentConfig
7-
import ai.koog.agents.core.dsl.builder.forwardTo
8-
import ai.koog.agents.core.dsl.builder.strategy
9-
import ai.koog.agents.core.dsl.extension.*
10-
import ai.koog.agents.core.environment.ReceivedToolResult
117
import ai.koog.agents.core.tools.Tool
128
import ai.koog.agents.core.tools.ToolArgs
139
import ai.koog.agents.core.tools.ToolRegistry
@@ -24,10 +20,6 @@ import kotlinx.coroutines.runBlocking
2420
import kotlin.uuid.ExperimentalUuidApi
2521
import kotlin.uuid.Uuid
2622

27-
28-
// Example threshold
29-
private const val MAX_TOKENS_THRESHOLD = 1000
30-
3123
fun main() = runBlocking {
3224
val executor: PromptExecutor = simpleOpenAIExecutor(ApiKeyService.openAIApiKey)
3325

@@ -39,49 +31,6 @@ fun main() = runBlocking {
3931
tools(CalculatorTools().asTools())
4032
}
4133

42-
val strategy = strategy("test") {
43-
val nodeCallLLM by nodeLLMRequestMultiple()
44-
val nodeExecuteToolMultiple by nodeExecuteMultipleTools(parallelTools = true)
45-
val nodeSendToolResultMultiple by nodeLLMSendMultipleToolResults()
46-
val nodeCompressHistory by nodeLLMCompressHistory<List<ReceivedToolResult>>()
47-
48-
edge(nodeStart forwardTo nodeCallLLM)
49-
50-
edge(
51-
(nodeCallLLM forwardTo nodeFinish)
52-
transformed { it.first() }
53-
onAssistantMessage { true }
54-
)
55-
56-
edge(
57-
(nodeCallLLM forwardTo nodeExecuteToolMultiple)
58-
onMultipleToolCalls { true }
59-
)
60-
61-
edge(
62-
(nodeExecuteToolMultiple forwardTo nodeCompressHistory)
63-
onCondition { llm.readSession { prompt.latestTokenUsage > MAX_TOKENS_THRESHOLD } }
64-
)
65-
66-
edge(nodeCompressHistory forwardTo nodeSendToolResultMultiple)
67-
68-
edge(
69-
(nodeExecuteToolMultiple forwardTo nodeSendToolResultMultiple)
70-
onCondition { llm.readSession { prompt.latestTokenUsage <= MAX_TOKENS_THRESHOLD } }
71-
)
72-
73-
edge(
74-
(nodeSendToolResultMultiple forwardTo nodeExecuteToolMultiple)
75-
onMultipleToolCalls { true }
76-
)
77-
78-
edge(
79-
(nodeSendToolResultMultiple forwardTo nodeFinish)
80-
transformed { it.first() }
81-
onAssistantMessage { true }
82-
)
83-
}
84-
8534
// Create agent config with proper prompt
8635
val agentConfig = AIAgentConfig(
8736
prompt = prompt("test") {
@@ -94,7 +43,7 @@ fun main() = runBlocking {
9443
// Create the runner
9544
val agent = AIAgent(
9645
promptExecutor = executor,
97-
strategy = strategy,
46+
strategy = CalculatorStrategy.strategy,
9847
agentConfig = agentConfig,
9948
toolRegistry = toolRegistry
10049
) {

‎examples/src/main/kotlin/ai/koog/agents/example/calculator/CalculatorTools.kt‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,19 @@
11
package ai.koog.agents.example.calculator
22

3+
import ai.koog.agents.core.dsl.builder.forwardTo
4+
import ai.koog.agents.core.dsl.builder.strategy
5+
import ai.koog.agents.core.dsl.extension.nodeExecuteMultipleTools
6+
import ai.koog.agents.core.dsl.extension.nodeLLMCompressHistory
7+
import ai.koog.agents.core.dsl.extension.nodeLLMRequestMultiple
8+
import ai.koog.agents.core.dsl.extension.nodeLLMSendMultipleToolResults
9+
import ai.koog.agents.core.dsl.extension.onAssistantMessage
10+
import ai.koog.agents.core.dsl.extension.onMultipleToolCalls
11+
import ai.koog.agents.core.environment.ReceivedToolResult
312
import ai.koog.agents.core.tools.annotations.LLMDescription
413
import ai.koog.agents.core.tools.annotations.Tool
514
import ai.koog.agents.core.tools.reflect.ToolSet
615

16+
@Suppress("unused")
717
@LLMDescription("Tools for basic calculator operations")
818
class CalculatorTools : ToolSet {
919

@@ -55,3 +65,50 @@ class CalculatorTools : ToolSet {
5565
return (a * b).toString()
5666
}
5767
}
68+
69+
object CalculatorStrategy {
70+
private const val MAX_TOKENS_THRESHOLD = 1000
71+
72+
val strategy = strategy("test") {
73+
val nodeCallLLM by nodeLLMRequestMultiple()
74+
val nodeExecuteToolMultiple by nodeExecuteMultipleTools(parallelTools = true)
75+
val nodeSendToolResultMultiple by nodeLLMSendMultipleToolResults()
76+
val nodeCompressHistory by nodeLLMCompressHistory<List<ReceivedToolResult>>()
77+
78+
edge(nodeStart forwardTo nodeCallLLM)
79+
80+
edge(
81+
(nodeCallLLM forwardTo nodeFinish)
82+
transformed { it.first() }
83+
onAssistantMessage { true }
84+
)
85+
86+
edge(
87+
(nodeCallLLM forwardTo nodeExecuteToolMultiple)
88+
onMultipleToolCalls { true }
89+
)
90+
91+
edge(
92+
(nodeExecuteToolMultiple forwardTo nodeCompressHistory)
93+
onCondition { llm.readSession { prompt.latestTokenUsage > MAX_TOKENS_THRESHOLD } }
94+
)
95+
96+
edge(nodeCompressHistory forwardTo nodeSendToolResultMultiple)
97+
98+
edge(
99+
(nodeExecuteToolMultiple forwardTo nodeSendToolResultMultiple)
100+
onCondition { llm.readSession { prompt.latestTokenUsage <= MAX_TOKENS_THRESHOLD } }
101+
)
102+
103+
edge(
104+
(nodeSendToolResultMultiple forwardTo nodeExecuteToolMultiple)
105+
onMultipleToolCalls { true }
106+
)
107+
108+
edge(
109+
(nodeSendToolResultMultiple forwardTo nodeFinish)
110+
transformed { it.first() }
111+
onAssistantMessage { true }
112+
)
113+
}
114+
}

‎examples/src/main/kotlin/ai/koog/agents/example/calculator/OllamaCalculatorExample.kt‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
package ai.koog.agents.example.calculator
2+
3+
import ai.koog.agents.core.agent.AIAgent
4+
import ai.koog.agents.core.agent.config.AIAgentConfig
5+
import ai.koog.agents.core.tools.Tool
6+
import ai.koog.agents.core.tools.ToolArgs
7+
import ai.koog.agents.core.tools.ToolRegistry
8+
import ai.koog.agents.core.tools.reflect.asTools
9+
import ai.koog.agents.ext.tool.AskUser
10+
import ai.koog.agents.ext.tool.SayToUser
11+
import ai.koog.agents.features.eventHandler.feature.handleEvents
12+
import ai.koog.prompt.dsl.prompt
13+
import ai.koog.prompt.executor.llms.all.simpleOllamaAIExecutor
14+
import ai.koog.prompt.executor.model.PromptExecutor
15+
import ai.koog.prompt.llm.OllamaModels
16+
import kotlinx.coroutines.runBlocking
17+
import kotlin.uuid.ExperimentalUuidApi
18+
import kotlin.uuid.Uuid
19+
20+
@OptIn(ExperimentalUuidApi::class)
21+
fun main() = runBlocking {
22+
val executor: PromptExecutor = simpleOllamaAIExecutor()
23+
24+
// Create tool registry with calculator tools
25+
val toolRegistry = ToolRegistry {
26+
// Special tool, required with this type of agent.
27+
tool(AskUser)
28+
tool(SayToUser)
29+
tools(CalculatorTools().asTools())
30+
}
31+
32+
// Create agent config with proper prompt
33+
val agentConfig = AIAgentConfig(
34+
prompt = prompt("test") {
35+
system("You are a calculator.")
36+
},
37+
model = OllamaModels.Meta.LLAMA_3_2,
38+
maxAgentIterations = 50
39+
)
40+
41+
val agent = AIAgent(
42+
promptExecutor = executor,
43+
strategy = CalculatorStrategy.strategy,
44+
agentConfig = agentConfig,
45+
toolRegistry = toolRegistry
46+
) {
47+
handleEvents {
48+
onToolCall { tool: Tool<*, *>, toolArgs: ToolArgs ->
49+
println("Tool called: tool ${tool.name}, args $toolArgs")
50+
}
51+
52+
onAgentRunError { strategyName: String, sessionUuid: Uuid?, throwable: Throwable ->
53+
println("An error occurred: ${throwable.message}\n${throwable.stackTraceToString()}")
54+
}
55+
56+
onAgentFinished { strategyName: String, result: String? ->
57+
println("Result: $result")
58+
}
59+
}
60+
}
61+
62+
runBlocking {
63+
agent.run("(10 + 20) * (5 + 5) / (2 - 11)")
64+
}
65+
}

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,6 @@
22

33
package ai.koog.integration.tests
44

5-
import ai.koog.integration.tests.tools.AnswerVerificationTool
6-
import ai.koog.integration.tests.tools.GenericParameterTool
7-
import ai.koog.integration.tests.tools.GeographyQueryTool
85
import ai.koog.agents.core.agent.AIAgent
96
import ai.koog.agents.core.agent.config.AIAgentConfig
107
import ai.koog.agents.core.agent.entity.AIAgentStrategy
@@ -13,6 +10,9 @@ import ai.koog.agents.core.dsl.builder.strategy
1310
import ai.koog.agents.core.dsl.extension.*
1411
import ai.koog.agents.core.tools.ToolRegistry
1512
import ai.koog.agents.features.eventHandler.feature.EventHandler
13+
import ai.koog.integration.tests.tools.AnswerVerificationTool
14+
import ai.koog.integration.tests.tools.GenericParameterTool
15+
import ai.koog.integration.tests.tools.GeographyQueryTool
1616
import ai.koog.integration.tests.utils.annotations.Retry
1717
import ai.koog.prompt.dsl.prompt
1818
import ai.koog.prompt.executor.model.PromptExecutor

‎prompt/prompt-executor/prompt-executor-clients/prompt-executor-ollama-client/src/commonMain/kotlin/ai/koog/prompt/executor/ollama/client/OllamaClient.kt‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,7 @@ public class OllamaClient(
8080
): List<Message.Response> {
8181
require(model.provider == LLMProvider.Ollama) { "Model not supported by Ollama" }
8282

83-
val response: OllamaChatResponseDTO = client.post(DEFAULT_MESSAGE_PATH) {
83+
val response = client.post(DEFAULT_MESSAGE_PATH) {
8484
setBody(
8585
OllamaChatRequestDTO(
8686
model = model.id,
@@ -90,13 +90,18 @@ public class OllamaClient(
9090
options = prompt.extractOllamaOptions(),
9191
stream = false,
9292
))
93-
}.body<OllamaChatResponseDTO>()
93+
}
9494

95-
return parseResponse(response, prompt)
95+
if (response.status.isSuccess()) {
96+
return parseResponse(response.body<OllamaChatResponseDTO>())
97+
} else {
98+
val errorResponse = response.body<OllamaErrorResponseDTO>()
99+
logger.error { "Ollama error: ${errorResponse.error}" }
100+
throw RuntimeException("Ollama API error: ${errorResponse.error}")
101+
}
96102
}
97103

98-
99-
private fun parseResponse(response: OllamaChatResponseDTO, prompt: Prompt): List<Message.Response> {
104+
private fun parseResponse(response: OllamaChatResponseDTO): List<Message.Response> {
100105
val messages = response.message ?: return emptyList()
101106
val content = messages.content
102107
val toolCalls = messages.toolCalls ?: emptyList()

‎prompt/prompt-executor/prompt-executor-clients/prompt-executor-ollama-client/src/commonMain/kotlin/ai/koog/prompt/executor/ollama/client/dto/OllamaModels.kt‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,12 @@ internal data class OllamaChatResponseDTO(
8686
@SerialName("eval_count") val evalCount: Int? = null
8787
)
8888

89+
/**
90+
* Error response from the /api/chat endpoint.
91+
*/
92+
@Serializable
93+
internal data class OllamaErrorResponseDTO(val error: String)
94+
8995
/**
9096
* Represents a request to generate an embedding using a specific model.
9197
*

‎prompt/prompt-llm/src/commonMain/kotlin/ai/koog/prompt/llm/OllamaModels.kt‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -172,6 +172,26 @@ public object OllamaModels {
172172
)
173173
)
174174

175+
/**
176+
* Represents the QWQ model with 32 billion parameters.
177+
*
178+
* This predefined instance of `LLModel` is provided by Alibaba and supports the following capabilities:
179+
* - `Temperature`: Allows adjustment of the temperature setting for controlling the randomness in responses.
180+
* - `Schema.JSON.Simple`: Supports tasks requiring JSON schema validation and handling in a simplified manner.
181+
* - `Tools`: Enables interaction with external tools or functionalities within the model's ecosystem.
182+
*
183+
* The model is identified by the unique ID "qwq:32b" and categorized under the Ollama provider.
184+
*/
185+
public val QWQ_32B: LLModel = LLModel(
186+
provider = LLMProvider.Ollama,
187+
id = "qwq:32b",
188+
capabilities = listOf(
189+
LLMCapability.Temperature,
190+
LLMCapability.Schema.JSON.Simple,
191+
LLMCapability.Tools
192+
)
193+
)
194+
175195
/**
176196
* Represents the `QWQ` language model instance provided by Alibaba with specific capabilities.
177197
*

0 commit comments

Comments
 (0)