feat: correx-managed model lifecycle slice 1 — config + manager factory + boot/shutdown

[[models]] config (ModelConfig/ModelsSettings) parsed by ConfigLoader;
InfrastructureModule.createModelManager + modelConfigToDescriptor; Main.kt
managed boot path spawns the default llama-server when [[models]] is present
(static [[providers]] path preserved when absent) and kills it on shutdown.

Plan: docs/plans/2026-05-31-model-lifecycle-management.md (slice 1 of 5).
This commit is contained in:
2026-06-01 11:03:05 +04:00
parent 5beb866036
commit e45a626cc4
10 changed files with 576 additions and 39 deletions
@@ -47,6 +47,7 @@ import com.correx.core.toolintent.rules.PathContainmentRule
import com.correx.infrastructure.InfrastructureModule import com.correx.infrastructure.InfrastructureModule
import com.correx.infrastructure.inference.DefaultProviderRegistry import com.correx.infrastructure.inference.DefaultProviderRegistry
import com.correx.infrastructure.inference.FirstAvailableRoutingStrategy import com.correx.infrastructure.inference.FirstAvailableRoutingStrategy
import com.correx.core.inference.InferenceProvider
import com.correx.infrastructure.inference.llama.cpp.LlamaCppInferenceProvider import com.correx.infrastructure.inference.llama.cpp.LlamaCppInferenceProvider
import com.correx.infrastructure.tools.FileEditConfig import com.correx.infrastructure.tools.FileEditConfig
import com.correx.infrastructure.tools.FileReadConfig import com.correx.infrastructure.tools.FileReadConfig
@@ -71,15 +72,48 @@ fun main() {
logStoresInfo(eventStore, artifactStore) logStoresInfo(eventStore, artifactStore)
val correxConfig = ConfigLoader.load() val correxConfig = ConfigLoader.load()
val providers = buildProviders(correxConfig) val eventDispatcher = EventDispatcher(eventStore)
require(providers.isNotEmpty()) { "At least one provider must be configured" }
val infraRegistry = InfrastructureModule.createProviderRegistry(providers)
val llamaProvider = providers.first()
val modelId = llamaProvider.id.value
val modelPath = "(from config)"
val llamaUrl = "(from config)"
logModelInfo(modelId, modelPath, llamaUrl, infraRegistry) // Managed path: [[models]] present → correx spawns llama-server at boot
// Static path: [[models]] absent → legacy [[providers]] / env-var behavior
val firstProvider: InferenceProvider
val infraRegistry: DefaultProviderRegistry
if (correxConfig.models.isNotEmpty()) {
val settings = correxConfig.modelsSettings
val modelManager = InfrastructureModule.createModelManager(settings, eventStore, eventDispatcher)
val targetId = settings.defaultModel ?: correxConfig.models.first().id
val targetModelConfig = correxConfig.models.firstOrNull { it.id == targetId }
?: correxConfig.models.first()
log.info("Managed path: loading model '{}' from '{}'", targetModelConfig.id, targetModelConfig.modelPath)
val descriptor = InfrastructureModule.modelConfigToDescriptor(targetModelConfig)
val managedProvider = runBlocking { modelManager.load(descriptor) }
firstProvider = managedProvider
infraRegistry = InfrastructureModule.createProviderRegistry(listOf(managedProvider))
// Also register any static providers from [[providers]] (non-managed remotes)
correxConfig.providers.mapNotNull { providerConfig ->
when (providerConfig.type.lowercase()) {
"llamacpp" -> buildProviderFromConfig(providerConfig)
else -> { log.error("Unknown provider type: {}", providerConfig.type); null }
}
}.forEach { infraRegistry.register(it) }
// Shutdown hook to kill the managed llama-server
Runtime.getRuntime().addShutdownHook(Thread {
log.info("Shutdown: unloading managed model '{}'", targetModelConfig.id)
runBlocking {
runCatching { modelManager.unload(targetModelConfig.id) }
.onFailure { log.warn("Error during managed model unload: {}", it.message) }
}
})
} else {
val staticProviders = buildProviders(correxConfig)
require(staticProviders.isNotEmpty()) { "At least one provider must be configured" }
firstProvider = staticProviders.first()
infraRegistry = InfrastructureModule.createProviderRegistry(staticProviders)
}
val modelId = firstProvider.id.value
logModelInfo(modelId, infraRegistry)
val repositories = buildRepositories(eventStore) val repositories = buildRepositories(eventStore)
val approvalEngine = DefaultApprovalEngine() val approvalEngine = DefaultApprovalEngine()
@@ -107,7 +141,7 @@ fun main() {
) )
val toolExecutor = InfrastructureModule.createToolExecutor( val toolExecutor = InfrastructureModule.createToolExecutor(
registry = toolRegistry, registry = toolRegistry,
eventDispatcher = EventDispatcher(eventStore), eventDispatcher = eventDispatcher,
workDir = sandboxRoot, workDir = sandboxRoot,
artifactStore = artifactStore, artifactStore = artifactStore,
) )
@@ -156,7 +190,7 @@ fun main() {
engines = engines, engines = engines,
retryCoordinator = DefaultRetryCoordinator(eventStore), retryCoordinator = DefaultRetryCoordinator(eventStore),
artifactStore = artifactStore, artifactStore = artifactStore,
tokenizer = llamaProvider.tokenizer, tokenizer = firstProvider.tokenizer,
) )
val defaultOrchestrationConfig = OrchestrationConfig( val defaultOrchestrationConfig = OrchestrationConfig(
sandboxRoot = sandboxRoot, sandboxRoot = sandboxRoot,
@@ -173,7 +207,7 @@ fun main() {
eventStore = eventStore, eventStore = eventStore,
inferenceRouter = inferenceRouter, inferenceRouter = inferenceRouter,
config = RouterConfig(), config = RouterConfig(),
tokenizer = llamaProvider.tokenizer, tokenizer = firstProvider.tokenizer,
embedder = embedder, embedder = embedder,
l3MemoryStore = l3MemoryStore, l3MemoryStore = l3MemoryStore,
) )
@@ -206,14 +240,10 @@ fun main() {
private fun logModelInfo( private fun logModelInfo(
modelId: String, modelId: String,
modelPath: String,
llamaUrl: String,
infraRegistry: DefaultProviderRegistry, infraRegistry: DefaultProviderRegistry,
) { ) {
log.info("==========MODEL INFO==========") log.info("==========MODEL INFO==========")
log.info(" model id : {}", modelId) log.info(" model id : {}", modelId)
log.info(" model path : {}", modelPath)
log.info(" llama url : {}", llamaUrl)
log.info(" providers : {} registered", infraRegistry.listAll().size) log.info(" providers : {} registered", infraRegistry.listAll().size)
} }
@@ -35,7 +35,9 @@ object ConfigLoader {
var currentSection = "" var currentSection = ""
val sections = mutableMapOf<String, MutableMap<String, Any>>() val sections = mutableMapOf<String, MutableMap<String, Any>>()
val providers = mutableListOf<MutableMap<String, Any>>() val providers = mutableListOf<MutableMap<String, Any>>()
val models = mutableListOf<MutableMap<String, Any>>()
var currentProvider: MutableMap<String, Any>? = null var currentProvider: MutableMap<String, Any>? = null
var currentModel: MutableMap<String, Any>? = null
for ((lineNum, line) in lines.withIndex()) { for ((lineNum, line) in lines.withIndex()) {
val trimmed = line.trim() val trimmed = line.trim()
@@ -46,18 +48,28 @@ object ConfigLoader {
} }
trimmed == "[[providers]]" -> { trimmed == "[[providers]]" -> {
// Start a new provider entry // Start a new provider entry
if (currentProvider != null) { if (currentProvider != null) providers.add(currentProvider)
providers.add(currentProvider) if (currentModel != null) { models.add(currentModel); currentModel = null }
}
currentProvider = mutableMapOf() currentProvider = mutableMapOf()
currentSection = "" currentSection = ""
} }
trimmed == "[[models]]" -> {
// Start a new model entry
if (currentModel != null) models.add(currentModel)
if (currentProvider != null) { providers.add(currentProvider); currentProvider = null }
currentModel = mutableMapOf()
currentSection = ""
}
trimmed.startsWith("[") && !trimmed.startsWith("[[") && trimmed.endsWith("]") -> { trimmed.startsWith("[") && !trimmed.startsWith("[[") && trimmed.endsWith("]") -> {
// Parse section headers like [server] or [tools.shell] // Parse section headers like [server] or [tools.shell]
if (currentProvider != null) { if (currentProvider != null) {
providers.add(currentProvider) providers.add(currentProvider)
currentProvider = null currentProvider = null
} }
if (currentModel != null) {
models.add(currentModel)
currentModel = null
}
currentSection = trimmed.substring(1, trimmed.length - 1).trim() currentSection = trimmed.substring(1, trimmed.length - 1).trim()
sections.putIfAbsent(currentSection, mutableMapOf()) sections.putIfAbsent(currentSection, mutableMapOf())
} }
@@ -69,22 +81,21 @@ object ConfigLoader {
val valueStr = trimmed.substring(eqIndex + 1).trim() val valueStr = trimmed.substring(eqIndex + 1).trim()
val parsedValue = parseValue(valueStr, lineNum + 1) val parsedValue = parseValue(valueStr, lineNum + 1)
if (currentProvider != null) { when {
currentProvider[key] = parsedValue currentProvider != null -> currentProvider[key] = parsedValue
} else if (currentSection.isNotEmpty()) { currentModel != null -> currentModel[key] = parsedValue
sections[currentSection]?.put(key, parsedValue) currentSection.isNotEmpty() -> sections[currentSection]?.put(key, parsedValue)
} }
} }
} }
} }
} }
// Don't forget the last provider if file ends with [[providers]] block // Don't forget the last array-of-table entry if file ends with one
if (currentProvider != null) { if (currentProvider != null) providers.add(currentProvider)
providers.add(currentProvider) if (currentModel != null) models.add(currentModel)
}
return buildConfig(sections, providers) return buildConfig(sections, providers, models)
} }
private fun parseValue(valueStr: String, lineNum: Int): Any { private fun parseValue(valueStr: String, lineNum: Int): Any {
@@ -200,6 +211,7 @@ object ConfigLoader {
private fun buildConfig( private fun buildConfig(
sections: Map<String, Map<String, Any>>, sections: Map<String, Map<String, Any>>,
providersList: List<Map<String, Any>> = emptyList(), providersList: List<Map<String, Any>> = emptyList(),
modelsList: List<Map<String, Any>> = emptyList(),
): CorrexConfig { ): CorrexConfig {
val serverSection = sections["server"] ?: emptyMap() val serverSection = sections["server"] ?: emptyMap()
val tuiSection = sections["tui"] ?: emptyMap() val tuiSection = sections["tui"] ?: emptyMap()
@@ -211,6 +223,7 @@ object ConfigLoader {
val toolsFileEditSection = sections["tools.file_edit"] ?: emptyMap() val toolsFileEditSection = sections["tools.file_edit"] ?: emptyMap()
val routerEmbedderSection = sections["router.embedder"] ?: emptyMap() val routerEmbedderSection = sections["router.embedder"] ?: emptyMap()
val routerL3Section = sections["router.l3"] ?: emptyMap() val routerL3Section = sections["router.l3"] ?: emptyMap()
val modelsSection = sections["models"] ?: emptyMap()
val server = ServerConfig( val server = ServerConfig(
host = asString(serverSection["host"], "localhost"), host = asString(serverSection["host"], "localhost"),
@@ -347,6 +360,32 @@ object ConfigLoader {
l3 = l3, l3 = l3,
) )
val models = modelsList.mapNotNull { modelMap ->
val id = asString(modelMap["id"]) ?: return@mapNotNull null
val modelPath = asString(modelMap["model_path"]) ?: return@mapNotNull null
val contextSize = asInt(modelMap["context_size"], 8192)
@Suppress("UNCHECKED_CAST")
val params = (modelMap["params"] as? Map<String, Any>)
?.mapValues { it.value.toString() }
?: emptyMap()
val capabilities = parseCapabilities(modelMap["capabilities"])
ModelConfig(
id = id,
modelPath = modelPath,
contextSize = contextSize,
params = params,
capabilities = capabilities,
)
}
val modelsSettings = ModelsSettings(
defaultModel = asStringOrNull(modelsSection["default_model"]),
llamaServerBin = asString(modelsSection["llama_server_bin"], "llama-server"),
host = asString(modelsSection["host"], "127.0.0.1"),
port = asInt(modelsSection["port"], 10000),
)
return CorrexConfig( return CorrexConfig(
server = server, server = server,
tui = tui, tui = tui,
@@ -354,6 +393,8 @@ object ConfigLoader {
tools = tools, tools = tools,
providers = providers, providers = providers,
router = router, router = router,
models = models,
modelsSettings = modelsSettings,
) )
} }
@@ -10,6 +10,8 @@ data class CorrexConfig(
val tools: ToolsConfig = ToolsConfig(), val tools: ToolsConfig = ToolsConfig(),
val providers: List<ProviderConfig> = emptyList(), val providers: List<ProviderConfig> = emptyList(),
val router: RouterConfig = RouterConfig(), val router: RouterConfig = RouterConfig(),
val models: List<ModelConfig> = emptyList(),
val modelsSettings: ModelsSettings = ModelsSettings(),
) )
@Serializable @Serializable
@@ -90,3 +92,20 @@ data class L3Config(
val dim: Int = 1536, val dim: Int = 1536,
val bitWidth: Int = 4, val bitWidth: Int = 4,
) )
@Serializable
data class ModelConfig(
val id: String,
val modelPath: String,
val contextSize: Int = 8192,
val params: Map<String, String> = emptyMap(),
val capabilities: Map<String, Double> = emptyMap(),
)
@Serializable
data class ModelsSettings(
val defaultModel: String? = null,
val llamaServerBin: String = "llama-server",
val host: String = "127.0.0.1",
val port: Int = 10000,
)
@@ -270,4 +270,98 @@ class ConfigLoaderTest {
assertEquals(1536, result.router.l3.dim) assertEquals(1536, result.router.l3.dim)
assertEquals(4, result.router.l3.bitWidth) assertEquals(4, result.router.l3.bitWidth)
} }
@Test
fun `parseToml parses models array-of-tables`() {
val toml = """
[models]
default_model = "mistral-7b"
llama_server_bin = "/usr/local/bin/llama-server"
host = "127.0.0.1"
port = 10001
[[models]]
id = "mistral-7b"
model_path = "/models/mistral-7b.gguf"
context_size = 4096
capabilities = { General = 1.0, Coding = 0.7, Reasoning = 0.6 }
[[models]]
id = "codellama-7b"
model_path = "/models/codellama-7b.gguf"
context_size = 8192
""".trimIndent()
val loader = ConfigLoader::class.java
val parseTomlMethod = loader.getDeclaredMethod("parseToml", String::class.java)
parseTomlMethod.isAccessible = true
val result = parseTomlMethod.invoke(ConfigLoader, toml) as CorrexConfig
assertEquals(2, result.models.size)
assertEquals("mistral-7b", result.models[0].id)
assertEquals("/models/mistral-7b.gguf", result.models[0].modelPath)
assertEquals(4096, result.models[0].contextSize)
assertEquals(3, result.models[0].capabilities.size)
assertEquals(1.0, result.models[0].capabilities["General"])
assertEquals(0.7, result.models[0].capabilities["Coding"])
assertEquals(0.6, result.models[0].capabilities["Reasoning"])
assertEquals("codellama-7b", result.models[1].id)
assertEquals("/models/codellama-7b.gguf", result.models[1].modelPath)
assertEquals(8192, result.models[1].contextSize)
assertEquals(0, result.models[1].capabilities.size)
assertEquals("mistral-7b", result.modelsSettings.defaultModel)
assertEquals("/usr/local/bin/llama-server", result.modelsSettings.llamaServerBin)
assertEquals("127.0.0.1", result.modelsSettings.host)
assertEquals(10001, result.modelsSettings.port)
}
@Test
fun `parseToml returns empty models list and default modelsSettings when sections absent`() {
val toml = """
[server]
host = "localhost"
""".trimIndent()
val loader = ConfigLoader::class.java
val parseTomlMethod = loader.getDeclaredMethod("parseToml", String::class.java)
parseTomlMethod.isAccessible = true
val result = parseTomlMethod.invoke(ConfigLoader, toml) as CorrexConfig
assertEquals(0, result.models.size)
assertEquals(null, result.modelsSettings.defaultModel)
assertEquals("llama-server", result.modelsSettings.llamaServerBin)
assertEquals("127.0.0.1", result.modelsSettings.host)
assertEquals(10000, result.modelsSettings.port)
}
@Test
fun `parseToml parses models alongside providers`() {
val toml = """
[[providers]]
id = "remote-provider"
type = "llamacpp"
model_id = "remote-model"
url = "http://192.168.1.10:10000"
[[models]]
id = "local-model"
model_path = "/models/local.gguf"
""".trimIndent()
val loader = ConfigLoader::class.java
val parseTomlMethod = loader.getDeclaredMethod("parseToml", String::class.java)
parseTomlMethod.isAccessible = true
val result = parseTomlMethod.invoke(ConfigLoader, toml) as CorrexConfig
assertEquals(1, result.providers.size)
assertEquals("remote-provider", result.providers[0].id)
assertEquals(1, result.models.size)
assertEquals("local-model", result.models[0].id)
assertEquals("/models/local.gguf", result.models[0].modelPath)
assertEquals(8192, result.models[0].contextSize)
}
} }
+37 -12
View File
@@ -30,23 +30,48 @@ enabled = true
[tools.file_edit] [tools.file_edit]
enabled = true enabled = true
# Provider configuration (array of tables) # ─────────────────────────────────────────────────────────────
[[providers]] # Managed model configuration (Slice 1: correx spawns + owns llama-server)
id = "local-llama" #
type = "llamacpp" # When [[models]] is present, correx launches llama-server at boot and kills it on
model_id = "mistral-7b" # shutdown. Use this instead of (or alongside) [[providers]].
#
# [models] — global settings for the managed llama-server process
# [[models]] — one entry per model file; correx will load the defaultModel at startup.
# ─────────────────────────────────────────────────────────────
[models]
default_model = "mistral-7b" # which [[models]] entry to load at boot
llama_server_bin = "llama-server" # path to the llama-server binary (default: llama-server)
host = "127.0.0.1" # host for llama-server to bind / correx to connect
port = 10000 # port for llama-server
[[models]]
id = "mistral-7b"
model_path = "~/models/mistral-7b-gguf/model.gguf" model_path = "~/models/mistral-7b-gguf/model.gguf"
url = "http://127.0.0.1:10000" context_size = 8192
capabilities = { General = 1.0, Coding = 0.7, Reasoning = 0.6, Summarization = 0.8, ToolCalling = 0.5 } capabilities = { General = 1.0, Coding = 0.7, Reasoning = 0.6, Summarization = 0.8, ToolCalling = 0.5 }
# Example: second provider (if you have multiple) # Example: additional model (swap via TUI in a later slice)
# [[models]]
# id = "codellama-7b"
# model_path = "~/models/codellama-7b-gguf/model.gguf"
# context_size = 4096
# capabilities = { General = 0.8, Coding = 1.0, Reasoning = 0.7, Summarization = 0.6, ToolCalling = 0.7 }
# ─────────────────────────────────────────────────────────────
# Legacy static provider configuration (array of tables)
# Use [[providers]] when correx should connect to an already-running llama-server
# (i.e. you launch the server yourself externally).
# If [[models]] is configured, [[providers]] entries are registered as additional
# static providers alongside the managed one.
# ─────────────────────────────────────────────────────────────
# [[providers]] # [[providers]]
# id = "alternative-llama" # id = "local-llama"
# type = "llamacpp" # type = "llamacpp"
# model_id = "neural-chat-7b" # model_id = "mistral-7b"
# model_path = "~/models/neural-chat-7b-gguf/model.gguf" # model_path = "~/models/mistral-7b-gguf/model.gguf"
# url = "http://127.0.0.1:10001" # url = "http://127.0.0.1:10000"
# capabilities = { General = 0.9, Coding = 0.8, Reasoning = 0.7, Summarization = 0.75, ToolCalling = 0.6 } # capabilities = { General = 1.0, Coding = 0.7, Reasoning = 0.6, Summarization = 0.8, ToolCalling = 0.5 }
# Router configuration (optional, defaults shown below) # Router configuration (optional, defaults shown below)
[router.embedder] [router.embedder]
+2
View File
@@ -28,6 +28,8 @@ dependencies {
implementation "io.ktor:ktor-client-content-negotiation:$ktor_version" implementation "io.ktor:ktor-client-content-negotiation:$ktor_version"
implementation "io.ktor:ktor-serialization-kotlinx-json:$ktor_version" implementation "io.ktor:ktor-serialization-kotlinx-json:$ktor_version"
implementation "org.slf4j:slf4j-api:2.0.16" implementation "org.slf4j:slf4j-api:2.0.16"
testImplementation "io.ktor:ktor-client-mock:$ktor_version"
testImplementation "org.jetbrains.kotlin:kotlin-test"
} }
tasks.named("koverVerify").configure { enabled = false } tasks.named("koverVerify").configure { enabled = false }
@@ -23,6 +23,8 @@ dependencies {
implementation "io.ktor:ktor-serialization-kotlinx-json:$ktor_version" implementation "io.ktor:ktor-serialization-kotlinx-json:$ktor_version"
testImplementation "org.junit.jupiter:junit-jupiter" testImplementation "org.junit.jupiter:junit-jupiter"
testImplementation "io.ktor:ktor-client-mock:$ktor_version" testImplementation "io.ktor:ktor-client-mock:$ktor_version"
testImplementation project(':infrastructure:persistence')
testImplementation "org.jetbrains.kotlin:kotlin-test"
} }
tasks.named("koverVerify").configure { enabled = false } tasks.named("koverVerify").configure { enabled = false }
@@ -0,0 +1,188 @@
package com.correx.infrastructure.inference.llama.cpp
import com.correx.core.events.EventDispatcher
import com.correx.core.events.events.ModelLoadedEvent
import com.correx.core.events.events.ModelUnloadedEvent
import com.correx.infrastructure.inference.commons.ModelDescriptor
import com.correx.infrastructure.inference.commons.ModelLoadException
import com.correx.infrastructure.inference.commons.ResidencyMode
import com.correx.infrastructure.persistence.InMemoryEventStore
import io.ktor.client.HttpClient
import io.ktor.client.engine.mock.MockEngine
import io.ktor.client.engine.mock.respond
import io.ktor.http.ContentType
import io.ktor.http.HttpHeaders
import io.ktor.http.headersOf
import kotlinx.coroutines.runBlocking
import org.junit.jupiter.api.Test
import org.junit.jupiter.api.assertThrows
import kotlin.test.assertEquals
import kotlin.test.assertNotNull
import kotlin.test.assertNull
/**
* Tests for [DefaultModelManager] lifecycle using a stub process and a mock HTTP health endpoint.
*
* Stub process: uses the system "sleep" binary called with llama-server-style args.
* The process may exit quickly (args are not understood by sleep), but the mock HTTP
* client returns "healthy" before the manager ever checks process liveness, so load()
* succeeds. This approach intentionally avoids any dependency on a real llama-server
* binary.
*/
class DefaultModelManagerTest {
private fun healthyHttpClient(): HttpClient {
val engine = MockEngine {
respond(
content = """{"status":"healthy"}""",
headers = headersOf(HttpHeaders.ContentType, ContentType.Application.Json.toString()),
)
}
return HttpClient(engine)
}
private fun unhealthyHttpClient(): HttpClient {
val engine = MockEngine {
respond(
content = """{"status":"loading"}""",
headers = headersOf(HttpHeaders.ContentType, ContentType.Application.Json.toString()),
)
}
return HttpClient(engine)
}
private fun descriptor(id: String = "test-model") = ModelDescriptor(
modelId = id,
modelPath = "/dev/null",
residencyMode = ResidencyMode.PERSISTENT,
contextSize = 512,
)
@Test
fun `load spawns process, polls health, emits ModelLoadedEvent, returns provider`(): Unit = runBlocking {
val eventStore = InMemoryEventStore()
val dispatcher = EventDispatcher(eventStore)
val manager = DefaultModelManager(
llamaServerBin = "sleep",
host = "127.0.0.1",
port = 29999,
healthTimeoutMs = 5_000L,
eventStore = eventStore,
httpClient = healthyHttpClient(),
eventDispatcher = dispatcher,
)
val provider = manager.load(descriptor())
assertNotNull(provider)
assertEquals("test-model", manager.currentModel()?.modelId)
val events = eventStore.allEvents().toList()
val loadedEvent = events.firstOrNull { it.payload is ModelLoadedEvent }
assertNotNull(loadedEvent, "ModelLoadedEvent must have been emitted")
assertEquals("test-model", (loadedEvent!!.payload as ModelLoadedEvent).modelId)
assertEquals("correx:model-manager", loadedEvent.metadata.sessionId.value)
manager.unload("test-model")
}
@Test
fun `currentModel returns null before any load`() {
val eventStore = InMemoryEventStore()
val dispatcher = EventDispatcher(eventStore)
val manager = DefaultModelManager(
llamaServerBin = "sleep",
host = "127.0.0.1",
port = 29998,
healthTimeoutMs = 5_000L,
eventStore = eventStore,
httpClient = healthyHttpClient(),
eventDispatcher = dispatcher,
)
assertNull(manager.currentModel())
}
@Test
fun `loading same model twice returns provider without re-spawning`(): Unit = runBlocking {
val eventStore = InMemoryEventStore()
val dispatcher = EventDispatcher(eventStore)
val manager = DefaultModelManager(
llamaServerBin = "sleep",
host = "127.0.0.1",
port = 29997,
healthTimeoutMs = 5_000L,
eventStore = eventStore,
httpClient = healthyHttpClient(),
eventDispatcher = dispatcher,
)
manager.load(descriptor("model-a"))
manager.load(descriptor("model-a"))
// Only one ModelLoadedEvent because the second load is a no-op
val loadedEvents = eventStore.allEvents().filter { it.payload is ModelLoadedEvent }.toList()
assertEquals(1, loadedEvents.size)
manager.unload("model-a")
}
@Test
fun `unload emits ModelUnloadedEvent and clears currentModel`(): Unit = runBlocking {
val eventStore = InMemoryEventStore()
val dispatcher = EventDispatcher(eventStore)
val manager = DefaultModelManager(
llamaServerBin = "sleep",
host = "127.0.0.1",
port = 29996,
healthTimeoutMs = 5_000L,
eventStore = eventStore,
httpClient = healthyHttpClient(),
eventDispatcher = dispatcher,
)
manager.load(descriptor("model-b"))
manager.unload("model-b")
assertNull(manager.currentModel())
val unloadedEvents = eventStore.allEvents().filter { it.payload is ModelUnloadedEvent }.toList()
assertEquals(1, unloadedEvents.size)
assertEquals("model-b", (unloadedEvents[0].payload as ModelUnloadedEvent).modelId)
assertEquals("correx:model-manager", unloadedEvents[0].metadata.sessionId.value)
}
@Test
fun `unload without prior load throws ModelLoadException`(): Unit = runBlocking {
val eventStore = InMemoryEventStore()
val dispatcher = EventDispatcher(eventStore)
val manager = DefaultModelManager(
llamaServerBin = "sleep",
host = "127.0.0.1",
port = 29995,
healthTimeoutMs = 5_000L,
eventStore = eventStore,
httpClient = healthyHttpClient(),
eventDispatcher = dispatcher,
)
assertThrows<ModelLoadException> { runBlocking { manager.unload("nonexistent") } }
}
@Test
fun `load throws ModelLoadException on health timeout`(): Unit = runBlocking {
val eventStore = InMemoryEventStore()
val dispatcher = EventDispatcher(eventStore)
val manager = DefaultModelManager(
llamaServerBin = "sleep",
host = "127.0.0.1",
port = 29994,
healthTimeoutMs = 100L,
eventStore = eventStore,
httpClient = unhealthyHttpClient(),
eventDispatcher = dispatcher,
)
assertThrows<ModelLoadException> { runBlocking { manager.load(descriptor("timeout-model")) } }
}
}
@@ -8,6 +8,8 @@ import com.correx.core.artifacts.repository.ArtifactRepository
import com.correx.core.artifactstore.ArtifactStore import com.correx.core.artifactstore.ArtifactStore
import com.correx.core.config.EmbedderConfig import com.correx.core.config.EmbedderConfig
import com.correx.core.config.L3Config import com.correx.core.config.L3Config
import com.correx.core.config.ModelConfig
import com.correx.core.config.ModelsSettings
import com.correx.core.events.EventDispatcher import com.correx.core.events.EventDispatcher
import com.correx.core.events.stores.EventStore import com.correx.core.events.stores.EventStore
import com.correx.core.inference.CapabilityScore import com.correx.core.inference.CapabilityScore
@@ -35,8 +37,11 @@ import com.correx.infrastructure.artifactscas.index.SqliteArtifactIndex
import com.correx.infrastructure.inference.DefaultProviderRegistry import com.correx.infrastructure.inference.DefaultProviderRegistry
import com.correx.infrastructure.inference.commons.ModelDescriptor import com.correx.infrastructure.inference.commons.ModelDescriptor
import com.correx.infrastructure.inference.commons.ResidencyMode import com.correx.infrastructure.inference.commons.ResidencyMode
import com.correx.infrastructure.inference.llama.cpp.DefaultModelManager
import com.correx.infrastructure.inference.llama.cpp.LlamaCppEmbedder import com.correx.infrastructure.inference.llama.cpp.LlamaCppEmbedder
import com.correx.infrastructure.inference.llama.cpp.LlamaCppInferenceProvider import com.correx.infrastructure.inference.llama.cpp.LlamaCppInferenceProvider
import io.ktor.client.HttpClient
import io.ktor.client.engine.cio.CIO
import com.correx.infrastructure.router.turbovec.TurboVecL3MemoryStore import com.correx.infrastructure.router.turbovec.TurboVecL3MemoryStore
import com.correx.infrastructure.router.turbovec.TurboVecSidecarConfig import com.correx.infrastructure.router.turbovec.TurboVecSidecarConfig
import com.correx.infrastructure.persistence.SqliteEventStore import com.correx.infrastructure.persistence.SqliteEventStore
@@ -110,6 +115,49 @@ object InfrastructureModule {
baseUrl = baseUrl, baseUrl = baseUrl,
) )
fun createModelManager(
settings: ModelsSettings,
eventStore: EventStore,
eventDispatcher: EventDispatcher,
httpClient: HttpClient = HttpClient(CIO),
healthTimeoutMs: Long = 30_000L,
): DefaultModelManager = DefaultModelManager(
llamaServerBin = settings.llamaServerBin,
host = settings.host,
port = settings.port,
healthTimeoutMs = healthTimeoutMs,
eventStore = eventStore,
httpClient = httpClient,
eventDispatcher = eventDispatcher,
)
fun modelConfigToDescriptor(config: ModelConfig): ModelDescriptor = ModelDescriptor(
modelId = config.id,
modelPath = config.modelPath,
residencyMode = ResidencyMode.PERSISTENT,
contextSize = config.contextSize,
capabilities = parseCapabilitiesFromConfig(config.capabilities),
)
private fun parseCapabilitiesFromConfig(capsMap: Map<String, Double>): Set<CapabilityScore> {
val result = mutableSetOf<CapabilityScore>()
capsMap.forEach { (capName, score) ->
val capability = when (capName.lowercase()) {
"general" -> ModelCapability.General
"coding" -> ModelCapability.Coding
"reasoning" -> ModelCapability.Reasoning
"summarization" -> ModelCapability.Summarization
"toolcalling" -> ModelCapability.ToolCalling
else -> {
log.warn("Unknown capability name '{}' in [[models]] config", capName)
return@forEach
}
}
result.add(CapabilityScore(capability, score))
}
return result
}
fun createProviderRegistry( fun createProviderRegistry(
providers: List<InferenceProvider> = emptyList(), providers: List<InferenceProvider> = emptyList(),
): DefaultProviderRegistry = DefaultProviderRegistry(providers) ): DefaultProviderRegistry = DefaultProviderRegistry(providers)
@@ -0,0 +1,88 @@
package com.correx.infrastructure
import com.correx.core.config.ModelConfig
import com.correx.core.config.ModelsSettings
import com.correx.core.events.EventDispatcher
import com.correx.core.inference.ModelCapability
import com.correx.infrastructure.inference.commons.ResidencyMode
import com.correx.infrastructure.persistence.InMemoryEventStore
import io.ktor.client.HttpClient
import io.ktor.client.engine.mock.MockEngine
import io.ktor.client.engine.mock.respond
import io.ktor.http.ContentType
import io.ktor.http.HttpHeaders
import io.ktor.http.headersOf
import org.junit.jupiter.api.Test
import kotlin.test.assertEquals
import kotlin.test.assertNotNull
class InfrastructureModuleModelTest {
@Test
fun `modelConfigToDescriptor maps fields correctly`() {
val config = ModelConfig(
id = "my-model",
modelPath = "/models/my-model.gguf",
contextSize = 4096,
capabilities = mapOf("General" to 1.0, "Coding" to 0.8),
)
val descriptor = InfrastructureModule.modelConfigToDescriptor(config)
assertEquals("my-model", descriptor.modelId)
assertEquals("/models/my-model.gguf", descriptor.modelPath)
assertEquals(4096, descriptor.contextSize)
assertEquals(ResidencyMode.PERSISTENT, descriptor.residencyMode)
assertEquals(2, descriptor.capabilities.size)
val generalScore = descriptor.capabilities.firstOrNull { it.capability == ModelCapability.General }
assertNotNull(generalScore)
assertEquals(1.0, generalScore!!.score)
val codingScore = descriptor.capabilities.firstOrNull { it.capability == ModelCapability.Coding }
assertNotNull(codingScore)
assertEquals(0.8, codingScore!!.score)
}
@Test
fun `modelConfigToDescriptor handles empty capabilities`() {
val config = ModelConfig(
id = "bare-model",
modelPath = "/models/bare.gguf",
)
val descriptor = InfrastructureModule.modelConfigToDescriptor(config)
assertEquals("bare-model", descriptor.modelId)
assertEquals(8192, descriptor.contextSize)
assertEquals(0, descriptor.capabilities.size)
}
@Test
fun `createModelManager returns DefaultModelManager with correct settings`() {
val settings = ModelsSettings(
defaultModel = "test-model",
llamaServerBin = "/custom/llama-server",
host = "10.0.0.1",
port = 12345,
)
val eventStore = InMemoryEventStore()
val dispatcher = EventDispatcher(eventStore)
val mockEngine = MockEngine {
respond(
content = """{"status":"healthy"}""",
headers = headersOf(HttpHeaders.ContentType, ContentType.Application.Json.toString()),
)
}
val httpClient = HttpClient(mockEngine)
val manager = InfrastructureModule.createModelManager(
settings = settings,
eventStore = eventStore,
eventDispatcher = dispatcher,
httpClient = httpClient,
)
assertNotNull(manager)
// currentModel is null before any load
assertEquals(null, manager.currentModel())
}
}