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 32bc9344..8498fb31 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 @@ -16,6 +16,7 @@ import com.correx.core.inference.CapabilityScore import com.correx.core.inference.DefaultInferenceRouter import com.correx.core.inference.InferenceProjector import com.correx.core.inference.InferenceRepository +import com.correx.core.inference.InferenceRouter import com.correx.core.inference.ModelCapability import com.correx.core.kernel.orchestration.DefaultOrchestrationReducer import com.correx.core.kernel.orchestration.DefaultSessionOrchestrator @@ -47,6 +48,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.infrastructure.inference.commons.ManagedInferenceRouter import com.correx.core.inference.InferenceProvider import com.correx.infrastructure.inference.llama.cpp.LlamaCppInferenceProvider import com.correx.infrastructure.tools.FileEditConfig @@ -78,10 +80,12 @@ fun main() { // Static path: [[models]] absent → legacy [[providers]] / env-var behavior val firstProvider: InferenceProvider val infraRegistry: DefaultProviderRegistry + val inferenceRouter: InferenceRouter if (correxConfig.models.isNotEmpty()) { val settings = correxConfig.modelsSettings val modelManager = InfrastructureModule.createModelManager(settings, eventStore, eventDispatcher) + val descriptors = correxConfig.models.map { InfrastructureModule.modelConfigToDescriptor(it) } val targetId = settings.defaultModel ?: correxConfig.models.first().id val targetModelConfig = correxConfig.models.firstOrNull { it.id == targetId } ?: correxConfig.models.first() @@ -97,6 +101,9 @@ fun main() { else -> { log.error("Unknown provider type: {}", providerConfig.type); null } } }.forEach { infraRegistry.register(it) } + // Managed router owns per-stage model selection: it ensures the resolved model is resident + // before routing (stage.modelId > capability match > default), evicting the prior model. + inferenceRouter = ManagedInferenceRouter(modelManager, descriptors, targetModelConfig.id) // Shutdown hook to kill the managed llama-server Runtime.getRuntime().addShutdownHook(Thread { log.info("Shutdown: unloading managed model '{}'", targetModelConfig.id) @@ -110,6 +117,7 @@ fun main() { require(staticProviders.isNotEmpty()) { "At least one provider must be configured" } firstProvider = staticProviders.first() infraRegistry = InfrastructureModule.createProviderRegistry(staticProviders) + inferenceRouter = DefaultInferenceRouter(infraRegistry, FirstAvailableRoutingStrategy()) } val modelId = firstProvider.id.value @@ -161,7 +169,6 @@ fun main() { ), ) - val inferenceRouter = DefaultInferenceRouter(infraRegistry, FirstAvailableRoutingStrategy()) val engines = OrchestratorEngines( transitionResolver = DefaultTransitionResolver { condition, ctx -> condition.evaluate(ctx) }, contextPackBuilder = DefaultContextPackBuilder(DefaultContextCompressor()), diff --git a/core/inference/src/main/kotlin/com/correx/core/inference/InferenceRouter.kt b/core/inference/src/main/kotlin/com/correx/core/inference/InferenceRouter.kt index 323971c3..b2436f95 100644 --- a/core/inference/src/main/kotlin/com/correx/core/inference/InferenceRouter.kt +++ b/core/inference/src/main/kotlin/com/correx/core/inference/InferenceRouter.kt @@ -27,6 +27,18 @@ interface InferenceRouter { stageId: StageId, requiredCapabilities: Set, ): InferenceProvider + + /** + * Resolves and selects a provider for a stage, optionally pinned to a specific model id. + * Implementations that own model lifecycle (e.g. a managed router) use [modelId] to select + * and load the target model before routing. The default implementation ignores [modelId] and + * falls back to capability-based selection, so static routers remain unaffected. + */ + suspend fun route( + stageId: StageId, + requiredCapabilities: Set, + modelId: String?, + ): InferenceProvider = route(stageId, requiredCapabilities) } class NoEligibleProviderException( diff --git a/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/SessionOrchestrator.kt b/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/SessionOrchestrator.kt index a5ee4cfd..f5d9fde4 100644 --- a/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/SessionOrchestrator.kt +++ b/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/SessionOrchestrator.kt @@ -753,7 +753,7 @@ abstract class SessionOrchestrator( timeoutMs: Long, responseFormat: ResponseFormat = ResponseFormat.Text, ): InferenceResult { - val provider = inferenceRouter.route(stageId, stageConfig.requiredCapabilities) + val provider = inferenceRouter.route(stageId, stageConfig.requiredCapabilities, stageConfig.modelId) log.debug( "[Orchestrator] inference session={} stage={} provider={} timeoutMs={}", sessionId.value, stageId.value, provider.id.value, timeoutMs, diff --git a/core/transitions/src/main/kotlin/com/correx/core/transitions/graph/StageConfig.kt b/core/transitions/src/main/kotlin/com/correx/core/transitions/graph/StageConfig.kt index 5082c3b2..87f9b84d 100644 --- a/core/transitions/src/main/kotlin/com/correx/core/transitions/graph/StageConfig.kt +++ b/core/transitions/src/main/kotlin/com/correx/core/transitions/graph/StageConfig.kt @@ -7,6 +7,7 @@ import com.correx.core.inference.ModelCapability data class StageConfig( val requiredCapabilities: Set = emptySet(), + val modelId: String? = null, val tokenBudget: Int = 4096, val generationConfig: GenerationConfig = GenerationConfig( temperature = 0.7, diff --git a/infrastructure/inference/commons/src/main/kotlin/com/correx/infrastructure/inference/commons/ManagedInferenceRouter.kt b/infrastructure/inference/commons/src/main/kotlin/com/correx/infrastructure/inference/commons/ManagedInferenceRouter.kt new file mode 100644 index 00000000..fe375147 --- /dev/null +++ b/infrastructure/inference/commons/src/main/kotlin/com/correx/infrastructure/inference/commons/ManagedInferenceRouter.kt @@ -0,0 +1,70 @@ +package com.correx.infrastructure.inference.commons + +import com.correx.core.events.types.StageId +import com.correx.core.inference.InferenceProvider +import com.correx.core.inference.InferenceRouter +import com.correx.core.inference.ModelCapability +import com.correx.core.inference.NoEligibleProviderException + +/** + * An [InferenceRouter] that owns the local model lifecycle: it resolves a single target model + * per stage and ensures it is resident (via [ModelManager.ensureLoaded]) before returning the + * managed provider to run inference against. + * + * Resolution precedence (per the model-lifecycle spec): + * ``` + * stage.modelId > capability match over configured models > defaultModelId + * ``` + * The manual-swap pin (a higher-precedence override) layers in on top of this in a later slice. + * + * Because the returned provider is the live managed provider (its [InferenceProvider.id] changes + * with the resident model), there is no stale router-level health cache to invalidate on swap — + * selection always reflects the model the manager just made resident. + */ +class ManagedInferenceRouter( + private val manager: ModelManager, + descriptors: List, + private val defaultModelId: String, +) : InferenceRouter { + + private val byId: Map = descriptors.associateBy { it.modelId } + private val ordered: List = descriptors + + override suspend fun route( + stageId: StageId, + requiredCapabilities: Set, + ): InferenceProvider = route(stageId, requiredCapabilities, null) + + override suspend fun route( + stageId: StageId, + requiredCapabilities: Set, + modelId: String?, + ): InferenceProvider { + val targetId = modelId + ?: capabilityMatch(requiredCapabilities) + ?: defaultModelId + val descriptor = byId[targetId] + ?: throw NoEligibleProviderException(stageId, requiredCapabilities) + return manager.ensureLoaded(descriptor) + } + + /** + * Selects the configured model whose declared capabilities cover every required capability, + * preferring the highest summed score over those required capabilities. Returns null when no + * capabilities are required or none of the models satisfy them (callers fall back to default). + */ + private fun capabilityMatch(required: Set): String? { + if (required.isEmpty()) return null + return ordered + .filter { descriptor -> + val declared = descriptor.capabilities.map { it.capability }.toSet() + declared.containsAll(required) + } + .maxByOrNull { descriptor -> + descriptor.capabilities + .filter { it.capability in required } + .sumOf { it.score } + } + ?.modelId + } +} diff --git a/infrastructure/inference/commons/src/main/kotlin/com/correx/infrastructure/inference/commons/ModelManager.kt b/infrastructure/inference/commons/src/main/kotlin/com/correx/infrastructure/inference/commons/ModelManager.kt index 03020324..8e2ca129 100644 --- a/infrastructure/inference/commons/src/main/kotlin/com/correx/infrastructure/inference/commons/ModelManager.kt +++ b/infrastructure/inference/commons/src/main/kotlin/com/correx/infrastructure/inference/commons/ModelManager.kt @@ -16,6 +16,13 @@ interface ModelManager { */ suspend fun load(descriptor: ModelDescriptor): InferenceProvider + /** + * Ensures the model described by [descriptor] is the resident model, returning its provider. + * A no-op load if that model is already current; otherwise loads it (evicting the previous + * model). [load] is idempotent for the already-loaded model, so the default delegates to it. + */ + suspend fun ensureLoaded(descriptor: ModelDescriptor): InferenceProvider = load(descriptor) + /** * Unloads the currently loaded model. * diff --git a/testing/contracts/build.gradle b/testing/contracts/build.gradle index 55459ea1..cda5a4a9 100644 --- a/testing/contracts/build.gradle +++ b/testing/contracts/build.gradle @@ -15,6 +15,7 @@ dependencies { testImplementation(project(":testing:fixtures")) testImplementation(project(":core:inference")) testImplementation(project(":core:context")) + testImplementation(project(":infrastructure:inference:commons")) testFixturesImplementation(project(":core:events")) testFixturesImplementation(project(":core:sessions")) testFixturesImplementation(project(":core:artifacts-store")) diff --git a/testing/contracts/src/test/kotlin/com/correx/testing/contracts/model/inference/ManagedInferenceRouterTest.kt b/testing/contracts/src/test/kotlin/com/correx/testing/contracts/model/inference/ManagedInferenceRouterTest.kt new file mode 100644 index 00000000..ba83f9f6 --- /dev/null +++ b/testing/contracts/src/test/kotlin/com/correx/testing/contracts/model/inference/ManagedInferenceRouterTest.kt @@ -0,0 +1,140 @@ +package com.correx.testing.contracts.model.inference + +import com.correx.core.events.types.ProviderId +import com.correx.core.events.types.StageId +import com.correx.core.inference.CapabilityScore +import com.correx.core.inference.InferenceProvider +import com.correx.core.inference.ModelCapability +import com.correx.core.inference.NoEligibleProviderException +import com.correx.core.inference.ProviderHealth +import com.correx.infrastructure.inference.commons.ManagedInferenceRouter +import com.correx.infrastructure.inference.commons.ModelDescriptor +import com.correx.infrastructure.inference.commons.ModelManager +import com.correx.infrastructure.inference.commons.ResidencyMode +import com.correx.testing.fixtures.inference.MockInferenceProvider +import kotlinx.coroutines.runBlocking +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Test +import org.junit.jupiter.api.assertThrows + +class ManagedInferenceRouterTest { + + private val stage = StageId("s1") + + /** Records every (idempotent) load/ensureLoaded so we can assert which model was made resident. */ + private class FakeModelManager : ModelManager { + val loads = mutableListOf() + private var current: ModelDescriptor? = null + + override suspend fun load(descriptor: ModelDescriptor): InferenceProvider { + loads.add(descriptor.modelId) + current = descriptor + return MockInferenceProvider(id = ProviderId("managed:${descriptor.modelId}")) + } + + override suspend fun unload(modelId: String) { current = null } + override fun currentModel(): ModelDescriptor? = current + override fun healthCheck(): ProviderHealth = ProviderHealth.Healthy + } + + private fun desc(id: String, vararg caps: Pair): ModelDescriptor = + ModelDescriptor( + modelId = id, + modelPath = "/models/$id", + residencyMode = ResidencyMode.PERSISTENT, + capabilities = caps.map { CapabilityScore(it.first, it.second) }.toSet(), + ) + + @Test + fun `stage modelId takes precedence over capability and default`(): Unit = runBlocking { + val manager = FakeModelManager() + val router = ManagedInferenceRouter( + manager = manager, + descriptors = listOf(desc("default-model"), desc("coder", ModelCapability.Coding to 1.0)), + defaultModelId = "default-model", + ) + + val provider = router.route(stage, setOf(ModelCapability.Coding), modelId = "coder") + + assertEquals("managed:coder", provider.id.value) + assertEquals(listOf("coder"), manager.loads) + } + + @Test + fun `capability match selects a model that covers the required capabilities`(): Unit = runBlocking { + val manager = FakeModelManager() + val router = ManagedInferenceRouter( + manager = manager, + descriptors = listOf( + desc("generalist", ModelCapability.General to 1.0), + desc("coder", ModelCapability.Coding to 1.0), + ), + defaultModelId = "generalist", + ) + + val provider = router.route(stage, setOf(ModelCapability.Coding), modelId = null) + + assertEquals("managed:coder", provider.id.value) + assertEquals(listOf("coder"), manager.loads) + } + + @Test + fun `capability match prefers the highest summed score`(): Unit = runBlocking { + val manager = FakeModelManager() + val router = ManagedInferenceRouter( + manager = manager, + descriptors = listOf( + desc("weak", ModelCapability.Coding to 0.3), + desc("strong", ModelCapability.Coding to 0.9), + ), + defaultModelId = "weak", + ) + + val provider = router.route(stage, setOf(ModelCapability.Coding), modelId = null) + + assertEquals("managed:strong", provider.id.value) + } + + @Test + fun `falls back to default when no modelId and no required capabilities`(): Unit = runBlocking { + val manager = FakeModelManager() + val router = ManagedInferenceRouter( + manager = manager, + descriptors = listOf(desc("default-model"), desc("coder", ModelCapability.Coding to 1.0)), + defaultModelId = "default-model", + ) + + val provider = router.route(stage, emptySet(), modelId = null) + + assertEquals("managed:default-model", provider.id.value) + assertEquals(listOf("default-model"), manager.loads) + } + + @Test + fun `falls back to default when no model satisfies the required capabilities`(): Unit = runBlocking { + val manager = FakeModelManager() + val router = ManagedInferenceRouter( + manager = manager, + descriptors = listOf(desc("default-model")), + defaultModelId = "default-model", + ) + + val provider = router.route(stage, setOf(ModelCapability.Reasoning), modelId = null) + + assertEquals("managed:default-model", provider.id.value) + } + + @Test + fun `unknown explicit modelId throws NoEligibleProviderException`() { + val manager = FakeModelManager() + val router = ManagedInferenceRouter( + manager = manager, + descriptors = listOf(desc("default-model")), + defaultModelId = "default-model", + ) + + assertThrows { + runBlocking { router.route(stage, emptySet(), modelId = "ghost") } + } + } +}