feat(inference): capability-aware routing for tool-heavy stages
Add CapabilityAwareRoutingStrategy — it hard-filters to providers that declare every required capability (like FirstAvailable) but ranks the matches by summed required-capability score; ties keep list order, so it is a strict superset of first-available. The server now wires it as the routing policy. The orchestrator augments a stage's required capabilities with ModelCapability.ToolCalling when the stage grants tools, so tool-heavy stages route to the best tool-calling model rather than whichever healthy provider comes first. Scores come from each provider's declared capabilities(); observed per-model reliability (GET /metrics/tool-reliability) can feed in later as an additional weight. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
@@ -11,7 +11,7 @@ Adapter for LLM inference backends. Implements `core:inference` interfaces. No d
|
||||
## Local Contracts
|
||||
|
||||
- `DefaultProviderRegistry` implements `ProviderRegistry` from `core:inference`.
|
||||
- `FirstAvailableRoutingStrategy` is the default routing policy; extend only in `core:inference`, not here.
|
||||
- Two `RoutingStrategy` implementations (both hard-filter to providers declaring every required capability): `FirstAvailableRoutingStrategy` picks the first match; `CapabilityAwareRoutingStrategy` ranks the matches by summed required-capability score (ties keep list order, so it's a superset of first-available). The server wires `CapabilityAwareRoutingStrategy` so tool-heavy stages (which request `ModelCapability.ToolCalling`) route to the best tool-caller. Extend the `RoutingStrategy` interface only in `core:inference`, not here.
|
||||
- All LLM responses are proposals — they must be validated by the core before affecting state (invariant #7). Adapters return raw responses; they do not validate.
|
||||
- Network calls to inference backends are environment observations; results must be recorded as events by callers if replay must reproduce them (invariant #9).
|
||||
- Model lifecycle (spawn/own the llama-server process) lives in `llama_cpp/`; the autonomous scheduler is intentionally unbuilt (see root CLAUDE.md).
|
||||
|
||||
+38
@@ -0,0 +1,38 @@
|
||||
package com.correx.infrastructure.inference
|
||||
|
||||
import com.correx.core.events.types.StageId
|
||||
import com.correx.core.inference.InferenceProvider
|
||||
import com.correx.core.inference.ModelCapability
|
||||
import com.correx.core.inference.NoEligibleProviderException
|
||||
import com.correx.core.inference.RoutingStrategy
|
||||
|
||||
/**
|
||||
* Like [FirstAvailableRoutingStrategy] it hard-filters to providers that declare every
|
||||
* [requiredCapabilities] — but among those it picks the one with the highest summed capability
|
||||
* **score** for the required set (falling back to total declared score when nothing is required).
|
||||
* Ties keep list order, so behaviour is identical to first-available when scores are equal.
|
||||
*
|
||||
* This is the seam the orchestrator relies on to steer tool-heavy stages (which request
|
||||
* `ModelCapability.ToolCalling`) to the model that calls tools best, rather than whichever healthy
|
||||
* provider happens to come first. The scores come from each provider's declared
|
||||
* `capabilities()` (config-supplied); observed per-model reliability (the `/metrics/tool-reliability`
|
||||
* groundwork) can later feed in as an additional weight.
|
||||
*/
|
||||
class CapabilityAwareRoutingStrategy : RoutingStrategy {
|
||||
override fun select(
|
||||
candidates: List<InferenceProvider>,
|
||||
requiredCapabilities: Set<ModelCapability>,
|
||||
): InferenceProvider {
|
||||
val eligible = candidates.filter { provider ->
|
||||
provider.capabilities().map { it.capability }.toSet().containsAll(requiredCapabilities)
|
||||
}
|
||||
// maxByOrNull keeps the first element on ties, preserving first-available order.
|
||||
return eligible.maxByOrNull { score(it, requiredCapabilities) }
|
||||
?: throw NoEligibleProviderException(StageId("routing"), requiredCapabilities)
|
||||
}
|
||||
|
||||
private fun score(provider: InferenceProvider, required: Set<ModelCapability>): Double {
|
||||
val scores = provider.capabilities().associate { it.capability to it.score }
|
||||
return if (required.isEmpty()) scores.values.sum() else required.sumOf { scores[it] ?: 0.0 }
|
||||
}
|
||||
}
|
||||
+73
@@ -0,0 +1,73 @@
|
||||
package com.correx.infrastructure.inference
|
||||
|
||||
import com.correx.core.events.types.ProviderId
|
||||
import com.correx.core.inference.CapabilityScore
|
||||
import com.correx.core.inference.InferenceProvider
|
||||
import com.correx.core.inference.InferenceRequest
|
||||
import com.correx.core.inference.InferenceResponse
|
||||
import com.correx.core.inference.ModelCapability
|
||||
import com.correx.core.inference.NoEligibleProviderException
|
||||
import com.correx.core.inference.ProviderHealth
|
||||
import com.correx.core.inference.Token
|
||||
import com.correx.core.inference.Tokenizer
|
||||
import com.correx.core.utils.TypeId
|
||||
import org.junit.jupiter.api.Assertions.assertEquals
|
||||
import org.junit.jupiter.api.Test
|
||||
import org.junit.jupiter.api.assertThrows
|
||||
|
||||
class CapabilityAwareRoutingStrategyTest {
|
||||
|
||||
private val strategy = CapabilityAwareRoutingStrategy()
|
||||
|
||||
private fun fakeProvider(id: String, vararg caps: Pair<ModelCapability, Double>): InferenceProvider =
|
||||
object : InferenceProvider {
|
||||
override val id: ProviderId = TypeId(id)
|
||||
override val name: String = id
|
||||
override val tokenizer: Tokenizer = object : Tokenizer {
|
||||
override suspend fun tokenize(text: String): List<Token> = emptyList()
|
||||
override suspend fun countTokens(text: String): Int = 0
|
||||
}
|
||||
override suspend fun infer(request: InferenceRequest): InferenceResponse = error("unused")
|
||||
override suspend fun healthCheck(): ProviderHealth = ProviderHealth.Healthy
|
||||
override fun capabilities(): Set<CapabilityScore> =
|
||||
caps.map { CapabilityScore(it.first, it.second) }.toSet()
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `picks the highest-scoring provider for the required capability`() {
|
||||
val weak = fakeProvider("weak", ModelCapability.ToolCalling to 0.4, ModelCapability.General to 1.0)
|
||||
val strong = fakeProvider("strong", ModelCapability.ToolCalling to 0.95, ModelCapability.General to 1.0)
|
||||
val selected = strategy.select(listOf(weak, strong), setOf(ModelCapability.ToolCalling))
|
||||
assertEquals("strong", selected.id.value)
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `still hard-filters to providers declaring the required capability`() {
|
||||
val noTools = fakeProvider("noTools", ModelCapability.General to 1.0)
|
||||
val tools = fakeProvider("tools", ModelCapability.ToolCalling to 0.3)
|
||||
val selected = strategy.select(listOf(noTools, tools), setOf(ModelCapability.ToolCalling))
|
||||
assertEquals("tools", selected.id.value)
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `throws when no candidate declares the required capability`() {
|
||||
val general = fakeProvider("general", ModelCapability.General to 1.0)
|
||||
assertThrows<NoEligibleProviderException> {
|
||||
strategy.select(listOf(general), setOf(ModelCapability.ToolCalling))
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `keeps list order on a score tie (first-available parity)`() {
|
||||
val a = fakeProvider("a", ModelCapability.ToolCalling to 0.5)
|
||||
val b = fakeProvider("b", ModelCapability.ToolCalling to 0.5)
|
||||
assertEquals("a", strategy.select(listOf(a, b), setOf(ModelCapability.ToolCalling)).id.value)
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `with no required capabilities ranks by total declared score`() {
|
||||
val lean = fakeProvider("lean", ModelCapability.General to 0.5)
|
||||
val rich = fakeProvider("rich", ModelCapability.General to 1.0, ModelCapability.Coding to 0.9)
|
||||
assertEquals("rich", strategy.select(listOf(lean, rich), emptySet()).id.value)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user