Skip to content

Commit 3603dc0

Browse files
ptitjesOloloshechkin
ptitjes
authored and
Ololoshechkin
committed
Enhance the Ollama model definitions (#149)
1 parent ed30412 commit 3603dc0

File tree

6 files changed

+43
-131
lines changed

6 files changed

+43
-131
lines changed

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

Lines changed: 9 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import ai.koog.prompt.executor.clients.LLMClient
77
import ai.koog.prompt.executor.clients.LLMEmbeddingProvider
88
import ai.koog.prompt.executor.ollama.client.dto.*
99
import ai.koog.prompt.llm.LLMCapability
10+
import ai.koog.prompt.llm.LLMProvider
1011
import ai.koog.prompt.llm.LLModel
1112
import ai.koog.prompt.message.Message
1213
import io.ktor.client.*
@@ -30,13 +31,11 @@ import kotlinx.serialization.json.Json
3031
* @property baseUrl The base URL of the Ollama API server.
3132
* @property baseClient The HTTP client used for making requests.
3233
* @property timeoutConfig Timeout configuration for HTTP requests.
33-
* @property enableDynamicModels Whether to enable dynamic model resolution.
3434
*/
3535
public class OllamaClient(
3636
private val baseUrl: String = "http://localhost:11434",
3737
baseClient: HttpClient = HttpClient(engineFactoryProvider()),
3838
timeoutConfig: ConnectionTimeoutConfig = ConnectionTimeoutConfig(),
39-
private val enableDynamicModels: Boolean = true
4039
): LLMClient, LLMEmbeddingProvider {
4140

4241
private val ollamaJson = Json {
@@ -65,17 +64,13 @@ public class OllamaClient(
6564
model: LLModel,
6665
tools: List<ToolDescriptor>
6766
): List<Message.Response> {
68-
val ollamaModelName = if (enableDynamicModels) {
69-
modelResolver.resolveToOllamaName(model)
70-
} else {
71-
model.toOllamaModelId()
72-
}
67+
require(model.provider == LLMProvider.Ollama) { "Model not supported by Ollama" }
7368

7469
val response = client.post("$baseUrl/api/chat") {
7570
contentType(ContentType.Application.Json)
7671
setBody(
7772
OllamaChatRequestDTO(
78-
model = ollamaModelName,
73+
model = model.id,
7974
messages = prompt.toOllamaChatMessages(),
8075
stream = false,
8176
tools = if (tools.isNotEmpty()) tools.map { it.toOllamaTool() } else null
@@ -97,17 +92,13 @@ public class OllamaClient(
9792
prompt: Prompt,
9893
model: LLModel
9994
): Flow<String> = flow {
100-
val ollamaModelName = if (enableDynamicModels) {
101-
modelResolver.resolveToOllamaName(model)
102-
} else {
103-
model.toOllamaModelId()
104-
}
95+
require(model.provider == LLMProvider.Ollama) { "Model not supported by Ollama" }
10596

10697
val response = client.post("$baseUrl/api/chat") {
10798
contentType(ContentType.Application.Json)
10899
setBody(
109100
OllamaChatRequestDTO(
110-
model = ollamaModelName,
101+
model = model.id,
111102
messages = prompt.toOllamaChatMessages(),
112103
stream = true,
113104
)
@@ -127,7 +118,7 @@ public class OllamaClient(
127118
emit(content)
128119
}
129120
}
130-
} catch (e: Exception) {
121+
} catch (_: Exception) {
131122
// Skip malformed JSON lines
132123
continue
133124
}
@@ -143,19 +134,15 @@ public class OllamaClient(
143134
* @throws IllegalArgumentException if the model does not have the Embed capability.
144135
*/
145136
override suspend fun embed(text: String, model: LLModel): List<Double> {
137+
require(model.provider == LLMProvider.Ollama) { "Model not supported by Ollama" }
138+
146139
if (!model.capabilities.contains(LLMCapability.Embed)) {
147140
throw IllegalArgumentException("Model ${model.id} does not have the Embed capability")
148141
}
149142

150-
val ollamaModelName = if (enableDynamicModels) {
151-
modelResolver.resolveToOllamaName(model)
152-
} else {
153-
model.id
154-
}
155-
156143
val response = client.post("$baseUrl/api/embeddings") {
157144
contentType(ContentType.Application.Json)
158-
setBody(EmbeddingRequest(model = ollamaModelName, prompt = text))
145+
setBody(EmbeddingRequest(model = model.id, prompt = text))
159146
}
160147

161148
val embeddingResponse = response.body<EmbeddingResponse>()

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

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -91,9 +91,9 @@ public class OllamaModelManager(
9191
val capabilities = detectCapabilities(matchedModel)
9292

9393
return LLModel(
94-
provider = LLMProvider.Meta, // Default to Meta for now
95-
id = "ollama-${matchedModel.name}",
96-
capabilities = capabilities
94+
provider = LLMProvider.Ollama,
95+
id = matchedModel.name,
96+
capabilities = capabilities,
9797
)
9898
}
9999

@@ -122,8 +122,6 @@ public class OllamaModelManager(
122122
return null
123123
}
124124

125-
126-
127125
/**
128126
* Detects capabilities based on model information - purely dynamic.
129127
*/
@@ -189,9 +187,9 @@ public class OllamaModelManager(
189187
val capabilities = getCapabilitiesForFamily(family)
190188

191189
return LLModel(
192-
provider = LLMProvider.Meta, // Default provider
193-
id = "ollama-dynamic-$modelName",
194-
capabilities = capabilities
190+
provider = LLMProvider.Ollama,
191+
id = modelName,
192+
capabilities = capabilities,
195193
)
196194
}
197195

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

Lines changed: 0 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -10,60 +10,6 @@ import io.github.oshai.kotlinlogging.KotlinLogging
1010
public class OllamaModelResolver(
1111
private val modelManager: OllamaModelManager
1212
) {
13-
private val logger = KotlinLogging.logger { }
14-
15-
/**
16-
* Only the absolute minimum static mappings for backward compatibility.
17-
* These are for the existing hardcoded models in the codebase.
18-
*/
19-
private val backwardCompatibilityMappings = mapOf(
20-
"meta-llama-3-2" to "llama3.2",
21-
"alibaba-qwq" to "qwq",
22-
"alibaba-qwen-coder-2-5-32b" to "qwen2.5-coder:32b"
23-
)
24-
25-
/**
26-
* Resolves an LLModel to the actual Ollama model name - fully dynamic.
27-
*/
28-
public suspend fun resolveToOllamaName(model: LLModel): String {
29-
// Check backward compatibility mappings first (only for existing hardcoded models)
30-
backwardCompatibilityMappings[model.id]?.let {
31-
logger.debug { "Resolved ${model.id} to $it via backward compatibility mapping" }
32-
return it
33-
}
34-
35-
// If the model ID starts with "ollama-", extract the actual model name
36-
if (model.id.startsWith("ollama-")) {
37-
val extractedName = model.id.removePrefix("ollama-")
38-
if (extractedName.startsWith("dynamic-")) {
39-
return extractedName.removePrefix("dynamic-")
40-
}
41-
return extractedName
42-
}
43-
44-
// Try to find the model in available models by exact match
45-
val availableModels = modelManager.getAvailableModels()
46-
47-
// First try exact match
48-
availableModels.find { it.name.equals(model.id, ignoreCase = true) }?.let {
49-
logger.debug { "Found exact match for ${model.id} as ${it.name}" }
50-
return it.name
51-
}
52-
53-
// Try to find by partial match
54-
availableModels.find {
55-
it.name.contains(model.id, ignoreCase = true) ||
56-
it.model.contains(model.id, ignoreCase = true)
57-
}?.let {
58-
logger.debug { "Found partial match for ${model.id} as ${it.name}" }
59-
return it.name
60-
}
61-
62-
// As a last resort, use the model ID directly
63-
logger.warn { "No mapping found for model ${model.id}, using ID directly" }
64-
return model.id
65-
}
66-
6713
/**
6814
* Attempts to resolve a model name and create an LLModel if it doesn't exist.
6915
*/

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

Lines changed: 0 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -3,42 +3,10 @@ package ai.koog.prompt.executor.ollama.client.dto
33
import ai.koog.agents.core.tools.ToolDescriptor
44
import ai.koog.prompt.dsl.Prompt
55
import ai.koog.prompt.executor.ollama.tools.json.toJSONSchema
6-
import ai.koog.prompt.llm.LLModel
7-
import ai.koog.prompt.llm.OllamaModels
86
import ai.koog.prompt.message.Message
97
import kotlinx.serialization.encodeToString
108
import kotlinx.serialization.json.Json
119

12-
/**
13-
* Converts LLModel to Ollama model ID string.
14-
* This function provides backward compatibility for hardcoded models.
15-
* For dynamic model resolution, use OllamaModelResolver instead.
16-
*/
17-
internal fun LLModel.toOllamaModelId(): String = when (this.id) {
18-
OllamaModels.Meta.LLAMA_3_2.id -> "llama3.2"
19-
OllamaModels.Alibaba.QWQ.id -> "qwq"
20-
OllamaModels.Alibaba.QWEN_CODER_2_5_32B.id -> "qwen2.5-coder:32b"
21-
else -> {
22-
// Try to extract model name from dynamic model IDs
23-
when {
24-
this.id.startsWith("ollama-") -> {
25-
val extractedName = this.id.removePrefix("ollama-")
26-
if (extractedName.startsWith("dynamic-")) {
27-
extractedName.removePrefix("dynamic-")
28-
} else {
29-
extractedName
30-
}
31-
}
32-
33-
else -> {
34-
// Log warning and use ID directly as fallback
35-
println("Warning: Unknown model ID ${this.id}, using ID directly. Consider using OllamaModelResolver for dynamic model support.")
36-
this.id
37-
}
38-
}
39-
}
40-
}
41-
4210
/**
4311
* Converts a Prompt to a list of ChatMessage objects for the Ollama API.
4412
*/

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

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,7 @@ public sealed class LLMProvider(public val id: String, public val display: Strin
7070
*/
7171
@Serializable
7272
public data object Alibaba : LLMProvider("alibaba", "Alibaba")
73-
73+
7474
/**
7575
* Represents the OpenRouter provider within the available set of large language model providers.
7676
*
@@ -83,4 +83,17 @@ public sealed class LLMProvider(public val id: String, public val display: Strin
8383
*/
8484
@Serializable
8585
public data object OpenRouter : LLMProvider("openrouter", "OpenRouter")
86+
87+
/**
88+
* Represents the Ollama provider within the available set of large language model providers.
89+
*
90+
* Ollama is identified by its unique ID ("ollama") and display name ("Ollama").
91+
* It extends the `LLMProvider` sealed class, which serves as a base class for all supported language model providers.
92+
*
93+
* This data object adheres to the structure and serialization requirements defined by the parent class.
94+
* It is part of the available LLM provider hierarchy, which is used to configure and identify specific
95+
* providers for large language model functionalities and capabilities.
96+
*/
97+
@Serializable
98+
public data object Ollama : LLMProvider("ollama", "Ollama")
8699
}

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

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,20 @@
11
package ai.koog.prompt.llm
22

33
/**
4-
* Represents a collection of predefined Large Language Models (LLM) categorized by providers.
5-
* Each provider contains specific models with configurations such as unique identifiers and capabilities.
4+
* Represents a collection of predefined Large Language Models (LLM) categorized by makers.
5+
* Each maker contains specific models with configurations such as unique identifiers and capabilities.
66
*/
77
public object OllamaModels {
88
/**
9-
* The `Meta` object represents the configuration for the Meta-provided large language models (LLMs).
9+
* The `Meta` object represents the configuration for the Meta large language models (LLMs).
1010
* It contains the predefined model specifications for Meta's LLMs, including their identifiers
1111
* and supported capabilities.
1212
*/
1313
public object Meta {
1414
/**
1515
* Represents the LLAMA version 3.2 model provided by Meta.
1616
*
17-
* This variable defines an instance of the `LLModel` class with the Meta provider, a unique identifier "meta-llama-3-2",
17+
* This variable defines an instance of the `LLModel` class with the Ollama provider, a unique identifier "llama3.2",
1818
* and a set of capabilities. The supported capabilities include:
1919
* - Temperature adjustment.
2020
* - JSON Schema-based tasks (Simple Schema).
@@ -24,8 +24,8 @@ public object OllamaModels {
2424
* that require dynamic behavior adjustments, schema adherence, and tool-based interactions.
2525
*/
2626
public val LLAMA_3_2: LLModel = LLModel(
27-
provider = LLMProvider.Meta,
28-
id = "meta-llama-3-2",
27+
provider = LLMProvider.Ollama,
28+
id = "llama3.2",
2929
capabilities = listOf(
3030
LLMCapability.Temperature,
3131
LLMCapability.Schema.JSON.Simple,
@@ -35,16 +35,16 @@ public object OllamaModels {
3535
}
3636

3737
/**
38-
* Represents an object that contains predefined Large Language Models (LLMs) provided by Alibaba.
38+
* Represents an object that contains predefined Large Language Models (LLMs) made by Alibaba.
3939
*
4040
* The `Alibaba` object provides access to multiple LLM instances, each with specific identifiers and capabilities.
41-
* These models are configured with Alibaba as the provider and are characterized by their unique capabilities.
41+
* These models are configured with Ollama as the provider and are characterized by their unique capabilities.
4242
*/
4343
public object Alibaba {
4444
/**
4545
* Represents the `QWQ` language model instance provided by Alibaba with specific capabilities.
4646
*
47-
* The model is identified by its unique `id` "alibaba-qwq". It belongs to the Alibaba provider
47+
* The model is identified by its unique `id` "qwq". It belongs to the Ollama provider
4848
* and supports multiple advanced capabilities:
4949
* - Temperature Adjustment: Enables control over the randomness of the model's output.
5050
* - JSON Schema (Simple): Supports tasks structured through simple JSON schemas.
@@ -54,8 +54,8 @@ public object OllamaModels {
5454
* require these capabilities for varied and advanced tasks.
5555
*/
5656
public val QWQ: LLModel = LLModel(
57-
provider = LLMProvider.Alibaba,
58-
id = "alibaba-qwq",
57+
provider = LLMProvider.Ollama,
58+
id = "qwq",
5959
capabilities = listOf(
6060
LLMCapability.Temperature,
6161
LLMCapability.Schema.JSON.Simple,
@@ -71,11 +71,11 @@ public object OllamaModels {
7171
* - `Schema.JSON.Simple`: Supports tasks requiring JSON schema validation and handling in a simplified manner.
7272
* - `Tools`: Enables interaction with external tools or functionalities within the model's ecosystem.
7373
*
74-
* The model is identified by the unique ID "alibaba-qwen-coder-2-5-32b" and categorized under the Alibaba provider.
74+
* The model is identified by the unique ID "qwen2.5-coder:32b" and categorized under the Ollama provider.
7575
*/
7676
public val QWEN_CODER_2_5_32B: LLModel = LLModel(
77-
provider = LLMProvider.Alibaba,
78-
id = "alibaba-qwen-coder-2-5-32b",
77+
provider = LLMProvider.Ollama,
78+
id = "qwen2.5-coder:32b",
7979
capabilities = listOf(
8080
LLMCapability.Temperature,
8181
LLMCapability.Schema.JSON.Simple,

0 commit comments

Comments
 (0)