From e45a626cc4dd181f5007d12bf1be0c9a581a0ced Mon Sep 17 00:00:00 2001 From: kami Date: Mon, 1 Jun 2026 11:03:05 +0400 Subject: [PATCH] =?UTF-8?q?feat:=20correx-managed=20model=20lifecycle=20sl?= =?UTF-8?q?ice=201=20=E2=80=94=20config=20+=20manager=20factory=20+=20boot?= =?UTF-8?q?/shutdown?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit [[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). --- .../kotlin/com/correx/apps/server/Main.kt | 60 ++++-- .../com/correx/core/config/ConfigLoader.kt | 65 ++++-- .../com/correx/core/config/CorrexConfig.kt | 19 ++ .../correx/core/config/ConfigLoaderTest.kt | 94 +++++++++ docs/sample-config.toml | 49 +++-- infrastructure/build.gradle | 2 + .../inference/llama_cpp/build.gradle | 2 + .../llama/cpp/DefaultModelManagerTest.kt | 188 ++++++++++++++++++ .../infrastructure/InfrastructureModule.kt | 48 +++++ .../InfrastructureModuleModelTest.kt | 88 ++++++++ 10 files changed, 576 insertions(+), 39 deletions(-) create mode 100644 infrastructure/inference/llama_cpp/src/test/kotlin/com/correx/infrastructure/inference/llama/cpp/DefaultModelManagerTest.kt create mode 100644 infrastructure/src/test/kotlin/com/correx/infrastructure/InfrastructureModuleModelTest.kt diff --git a/apps/server/src/main/kotlin/com/correx/apps/server/Main.kt b/apps/server/src/main/kotlin/com/correx/apps/server/Main.kt index b93a0434..32bc9344 100644 --- a/apps/server/src/main/kotlin/com/correx/apps/server/Main.kt +++ b/apps/server/src/main/kotlin/com/correx/apps/server/Main.kt @@ -47,6 +47,7 @@ import com.correx.core.toolintent.rules.PathContainmentRule import com.correx.infrastructure.InfrastructureModule import com.correx.infrastructure.inference.DefaultProviderRegistry import com.correx.infrastructure.inference.FirstAvailableRoutingStrategy +import com.correx.core.inference.InferenceProvider import com.correx.infrastructure.inference.llama.cpp.LlamaCppInferenceProvider import com.correx.infrastructure.tools.FileEditConfig import com.correx.infrastructure.tools.FileReadConfig @@ -71,15 +72,48 @@ fun main() { logStoresInfo(eventStore, artifactStore) val correxConfig = ConfigLoader.load() - val providers = buildProviders(correxConfig) - 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)" + val eventDispatcher = EventDispatcher(eventStore) - 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 approvalEngine = DefaultApprovalEngine() @@ -107,7 +141,7 @@ fun main() { ) val toolExecutor = InfrastructureModule.createToolExecutor( registry = toolRegistry, - eventDispatcher = EventDispatcher(eventStore), + eventDispatcher = eventDispatcher, workDir = sandboxRoot, artifactStore = artifactStore, ) @@ -156,7 +190,7 @@ fun main() { engines = engines, retryCoordinator = DefaultRetryCoordinator(eventStore), artifactStore = artifactStore, - tokenizer = llamaProvider.tokenizer, + tokenizer = firstProvider.tokenizer, ) val defaultOrchestrationConfig = OrchestrationConfig( sandboxRoot = sandboxRoot, @@ -173,7 +207,7 @@ fun main() { eventStore = eventStore, inferenceRouter = inferenceRouter, config = RouterConfig(), - tokenizer = llamaProvider.tokenizer, + tokenizer = firstProvider.tokenizer, embedder = embedder, l3MemoryStore = l3MemoryStore, ) @@ -206,14 +240,10 @@ fun main() { private fun logModelInfo( modelId: String, - modelPath: String, - llamaUrl: String, infraRegistry: DefaultProviderRegistry, ) { log.info("==========MODEL INFO==========") log.info(" model id : {}", modelId) - log.info(" model path : {}", modelPath) - log.info(" llama url : {}", llamaUrl) log.info(" providers : {} registered", infraRegistry.listAll().size) } diff --git a/core/config/src/main/kotlin/com/correx/core/config/ConfigLoader.kt b/core/config/src/main/kotlin/com/correx/core/config/ConfigLoader.kt index 129024b4..97a8748f 100644 --- a/core/config/src/main/kotlin/com/correx/core/config/ConfigLoader.kt +++ b/core/config/src/main/kotlin/com/correx/core/config/ConfigLoader.kt @@ -35,7 +35,9 @@ object ConfigLoader { var currentSection = "" val sections = mutableMapOf>() val providers = mutableListOf>() + val models = mutableListOf>() var currentProvider: MutableMap? = null + var currentModel: MutableMap? = null for ((lineNum, line) in lines.withIndex()) { val trimmed = line.trim() @@ -46,18 +48,28 @@ object ConfigLoader { } trimmed == "[[providers]]" -> { // Start a new provider entry - if (currentProvider != null) { - providers.add(currentProvider) - } + if (currentProvider != null) providers.add(currentProvider) + if (currentModel != null) { models.add(currentModel); currentModel = null } currentProvider = mutableMapOf() 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("]") -> { // Parse section headers like [server] or [tools.shell] if (currentProvider != null) { providers.add(currentProvider) currentProvider = null } + if (currentModel != null) { + models.add(currentModel) + currentModel = null + } currentSection = trimmed.substring(1, trimmed.length - 1).trim() sections.putIfAbsent(currentSection, mutableMapOf()) } @@ -69,22 +81,21 @@ object ConfigLoader { val valueStr = trimmed.substring(eqIndex + 1).trim() val parsedValue = parseValue(valueStr, lineNum + 1) - if (currentProvider != null) { - currentProvider[key] = parsedValue - } else if (currentSection.isNotEmpty()) { - sections[currentSection]?.put(key, parsedValue) + when { + currentProvider != null -> currentProvider[key] = parsedValue + currentModel != null -> currentModel[key] = parsedValue + currentSection.isNotEmpty() -> sections[currentSection]?.put(key, parsedValue) } } } } } - // Don't forget the last provider if file ends with [[providers]] block - if (currentProvider != null) { - providers.add(currentProvider) - } + // Don't forget the last array-of-table entry if file ends with one + if (currentProvider != null) 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 { @@ -200,6 +211,7 @@ object ConfigLoader { private fun buildConfig( sections: Map>, providersList: List> = emptyList(), + modelsList: List> = emptyList(), ): CorrexConfig { val serverSection = sections["server"] ?: emptyMap() val tuiSection = sections["tui"] ?: emptyMap() @@ -211,6 +223,7 @@ object ConfigLoader { val toolsFileEditSection = sections["tools.file_edit"] ?: emptyMap() val routerEmbedderSection = sections["router.embedder"] ?: emptyMap() val routerL3Section = sections["router.l3"] ?: emptyMap() + val modelsSection = sections["models"] ?: emptyMap() val server = ServerConfig( host = asString(serverSection["host"], "localhost"), @@ -347,6 +360,32 @@ object ConfigLoader { 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) + ?.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( server = server, tui = tui, @@ -354,6 +393,8 @@ object ConfigLoader { tools = tools, providers = providers, router = router, + models = models, + modelsSettings = modelsSettings, ) } diff --git a/core/config/src/main/kotlin/com/correx/core/config/CorrexConfig.kt b/core/config/src/main/kotlin/com/correx/core/config/CorrexConfig.kt index 045eb0ad..3d006ce0 100644 --- a/core/config/src/main/kotlin/com/correx/core/config/CorrexConfig.kt +++ b/core/config/src/main/kotlin/com/correx/core/config/CorrexConfig.kt @@ -10,6 +10,8 @@ data class CorrexConfig( val tools: ToolsConfig = ToolsConfig(), val providers: List = emptyList(), val router: RouterConfig = RouterConfig(), + val models: List = emptyList(), + val modelsSettings: ModelsSettings = ModelsSettings(), ) @Serializable @@ -90,3 +92,20 @@ data class L3Config( val dim: Int = 1536, val bitWidth: Int = 4, ) + +@Serializable +data class ModelConfig( + val id: String, + val modelPath: String, + val contextSize: Int = 8192, + val params: Map = emptyMap(), + val capabilities: Map = 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, +) diff --git a/core/config/src/test/kotlin/com/correx/core/config/ConfigLoaderTest.kt b/core/config/src/test/kotlin/com/correx/core/config/ConfigLoaderTest.kt index 7f360c22..6ea9b22e 100644 --- a/core/config/src/test/kotlin/com/correx/core/config/ConfigLoaderTest.kt +++ b/core/config/src/test/kotlin/com/correx/core/config/ConfigLoaderTest.kt @@ -270,4 +270,98 @@ class ConfigLoaderTest { assertEquals(1536, result.router.l3.dim) 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) + } } diff --git a/docs/sample-config.toml b/docs/sample-config.toml index de9b093c..eedc25f7 100644 --- a/docs/sample-config.toml +++ b/docs/sample-config.toml @@ -30,23 +30,48 @@ enabled = true [tools.file_edit] enabled = true -# Provider configuration (array of tables) -[[providers]] -id = "local-llama" -type = "llamacpp" -model_id = "mistral-7b" +# ───────────────────────────────────────────────────────────── +# Managed model configuration (Slice 1: correx spawns + owns llama-server) +# +# When [[models]] is present, correx launches llama-server at boot and kills it on +# 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" -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 } -# 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]] -# id = "alternative-llama" +# id = "local-llama" # type = "llamacpp" -# model_id = "neural-chat-7b" -# model_path = "~/models/neural-chat-7b-gguf/model.gguf" -# url = "http://127.0.0.1:10001" -# capabilities = { General = 0.9, Coding = 0.8, Reasoning = 0.7, Summarization = 0.75, ToolCalling = 0.6 } +# model_id = "mistral-7b" +# model_path = "~/models/mistral-7b-gguf/model.gguf" +# url = "http://127.0.0.1:10000" +# capabilities = { General = 1.0, Coding = 0.7, Reasoning = 0.6, Summarization = 0.8, ToolCalling = 0.5 } # Router configuration (optional, defaults shown below) [router.embedder] diff --git a/infrastructure/build.gradle b/infrastructure/build.gradle index ea60381d..28e68ac6 100644 --- a/infrastructure/build.gradle +++ b/infrastructure/build.gradle @@ -28,6 +28,8 @@ dependencies { implementation "io.ktor:ktor-client-content-negotiation:$ktor_version" implementation "io.ktor:ktor-serialization-kotlinx-json:$ktor_version" 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 } diff --git a/infrastructure/inference/llama_cpp/build.gradle b/infrastructure/inference/llama_cpp/build.gradle index 1ee5603c..fac65f8d 100644 --- a/infrastructure/inference/llama_cpp/build.gradle +++ b/infrastructure/inference/llama_cpp/build.gradle @@ -23,6 +23,8 @@ dependencies { implementation "io.ktor:ktor-serialization-kotlinx-json:$ktor_version" testImplementation "org.junit.jupiter:junit-jupiter" testImplementation "io.ktor:ktor-client-mock:$ktor_version" + testImplementation project(':infrastructure:persistence') + testImplementation "org.jetbrains.kotlin:kotlin-test" } tasks.named("koverVerify").configure { enabled = false } diff --git a/infrastructure/inference/llama_cpp/src/test/kotlin/com/correx/infrastructure/inference/llama/cpp/DefaultModelManagerTest.kt b/infrastructure/inference/llama_cpp/src/test/kotlin/com/correx/infrastructure/inference/llama/cpp/DefaultModelManagerTest.kt new file mode 100644 index 00000000..51389573 --- /dev/null +++ b/infrastructure/inference/llama_cpp/src/test/kotlin/com/correx/infrastructure/inference/llama/cpp/DefaultModelManagerTest.kt @@ -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 { 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 { runBlocking { manager.load(descriptor("timeout-model")) } } + } +} diff --git a/infrastructure/src/main/kotlin/com/correx/infrastructure/InfrastructureModule.kt b/infrastructure/src/main/kotlin/com/correx/infrastructure/InfrastructureModule.kt index 6b40c433..33f05508 100644 --- a/infrastructure/src/main/kotlin/com/correx/infrastructure/InfrastructureModule.kt +++ b/infrastructure/src/main/kotlin/com/correx/infrastructure/InfrastructureModule.kt @@ -8,6 +8,8 @@ import com.correx.core.artifacts.repository.ArtifactRepository import com.correx.core.artifactstore.ArtifactStore import com.correx.core.config.EmbedderConfig 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.stores.EventStore 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.commons.ModelDescriptor 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.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.TurboVecSidecarConfig import com.correx.infrastructure.persistence.SqliteEventStore @@ -110,6 +115,49 @@ object InfrastructureModule { 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): Set { + val result = mutableSetOf() + 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( providers: List = emptyList(), ): DefaultProviderRegistry = DefaultProviderRegistry(providers) diff --git a/infrastructure/src/test/kotlin/com/correx/infrastructure/InfrastructureModuleModelTest.kt b/infrastructure/src/test/kotlin/com/correx/infrastructure/InfrastructureModuleModelTest.kt new file mode 100644 index 00000000..4556de84 --- /dev/null +++ b/infrastructure/src/test/kotlin/com/correx/infrastructure/InfrastructureModuleModelTest.kt @@ -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()) + } +}