Resolve LLM roles through an SPI so a role is not tied to one provider

embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/Asyncer.kt

@@ -19,6 +19,18 @@ import java.util.concurrent.CompletableFuture
 
 /**
  * Simple Java-friendly async interface.
+ *
+ * An implementation is a supported extension point, and one obligation comes with it: **capture
+ * the calling thread's context and re-establish it on the worker**, restoring whatever the worker
+ * held before. Three things travel that way today - the [com.embabel.agent.core.AgentProcess], the
+ * current Micrometer observation, and the
+ * [com.embabel.common.ai.model.ModelSelectionContext] - and
+ * [com.embabel.agent.spi.support.ExecutorAsyncer] is the reference for how.
+ *
+ * Dropping them does not fail loudly, which is what makes this worth stating. A lost model
+ * selection context means role resolution quietly falls back to deployment configuration and
+ * serves a model the deployment is billed for, on a call the user brought their own key for. A
+ * lost `AgentProcess` breaks the platform's own bookkeeping just as silently.
  */
 interface Asyncer {
 

embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/config/spring/AgentPlatformConfiguration.kt

@@ -36,8 +36,10 @@ import com.embabel.common.util.EmbabelObjectMapperHolder
 import com.embabel.common.ai.autoconfig.ProviderInitialization
 import com.embabel.common.ai.model.ConfigurableModelProvider
 import com.embabel.common.ai.model.ConfigurableModelProviderProperties
+import com.embabel.common.ai.model.CredentialLlmServiceFactory
 import com.embabel.common.ai.model.EmbeddingService
 import com.embabel.common.ai.model.ModelProvider
+import com.embabel.common.ai.model.RoleResolver
 import com.embabel.common.core.MobyNameGenerator
 import com.embabel.common.core.NameGenerator
 import com.embabel.common.textio.template.JinjavaTemplateRenderer
@@ -194,6 +196,12 @@ class AgentPlatformConfiguration(
             llms = applicationContext.getBeansOfType(LlmService::class.java).values.toList(),
             embeddingServices = applicationContext.getBeansOfType(EmbeddingService::class.java).values.toList(),
             properties = properties,
+            // Ordered, so that one application resolver can take precedence over another
+            roleResolvers = applicationContext.getBeanProvider(RoleResolver::class.java)
+                .orderedStream().toList(),
+            credentialLlmServiceFactories = applicationContext
+                .getBeanProvider(CredentialLlmServiceFactory::class.java)
+                .orderedStream().toList(),
         )
     }
 

embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/AbstractLlmOperations.kt

@@ -34,6 +34,7 @@ import com.embabel.agent.spi.validation.ValidationPromptGenerator
 import com.embabel.chat.Message
 import com.embabel.chat.UserMessage
 import com.embabel.common.ai.model.AutoModelSelectionCriteria
+import com.embabel.common.ai.model.ByRoleModelSelectionCriteria
 import com.embabel.common.ai.model.LlmOptions
 import com.embabel.common.ai.model.ModelProvider
 import com.embabel.common.ai.model.ModelSelectionCriteria
@@ -144,6 +145,14 @@ abstract class AbstractLlmOperations(
         agentProcess: AgentProcess,
         action: Action?,
     ): O {
+        // Shadowed deliberately. After this line the resolved interaction IS the interaction for
+        // the rest of the method, and shadowing makes the unresolved one unreachable. A distinct
+        // name would leave both in scope, differing only in whether a role has become a concrete
+        // model plus its hyperparameters - and picking the wrong one is not a compile error, it is
+        // a call that silently skips role resolution and runs on the default model.
+        @Suppress("NAME_SHADOWING")
+        val interaction = withRoleResolved(interaction)
+
         val (allTools, llmRequestEvent) = getToolsAndEvent(
             agentProcess = agentProcess,
             interaction = interaction,
@@ -260,6 +269,14 @@ abstract class AbstractLlmOperations(
         agentProcess: AgentProcess,
         action: Action?,
     ): Result<O> {
+        // Shadowed deliberately. After this line the resolved interaction IS the interaction for
+        // the rest of the method, and shadowing makes the unresolved one unreachable. A distinct
+        // name would leave both in scope, differing only in whether a role has become a concrete
+        // model plus its hyperparameters - and picking the wrong one is not a compile error, it is
+        // a call that silently skips role resolution and runs on the default model.
+        @Suppress("NAME_SHADOWING")
+        val interaction = withRoleResolved(interaction)
+
         val (allTools, llmRequestEvent) = getToolsAndEvent(
             agentProcess = agentProcess,
             interaction = interaction,
@@ -312,6 +329,14 @@ abstract class AbstractLlmOperations(
         agentProcess: AgentProcess,
         action: Action?,
     ): ThinkingResponse<O> {
+        // Shadowed deliberately. After this line the resolved interaction IS the interaction for
+        // the rest of the method, and shadowing makes the unresolved one unreachable. A distinct
+        // name would leave both in scope, differing only in whether a role has become a concrete
+        // model plus its hyperparameters - and picking the wrong one is not a compile error, it is
+        // a call that silently skips role resolution and runs on the default model.
+        @Suppress("NAME_SHADOWING")
+        val interaction = withRoleResolved(interaction)
+
         val (allTools, llmRequestEvent) = getToolsAndEvent(
             agentProcess = agentProcess,
             interaction = interaction,
@@ -364,6 +389,14 @@ abstract class AbstractLlmOperations(
         agentProcess: AgentProcess,
         action: Action?,
     ): Result<ThinkingResponse<O>> {
+        // Shadowed deliberately. After this line the resolved interaction IS the interaction for
+        // the rest of the method, and shadowing makes the unresolved one unreachable. A distinct
+        // name would leave both in scope, differing only in whether a role has become a concrete
+        // model plus its hyperparameters - and picking the wrong one is not a compile error, it is
+        // a call that silently skips role resolution and runs on the default model.
+        @Suppress("NAME_SHADOWING")
+        val interaction = withRoleResolved(interaction)
+
         val (allTools, llmRequestEvent) = getToolsAndEvent(
             agentProcess = agentProcess,
             interaction = interaction,
@@ -409,6 +442,33 @@ abstract class AbstractLlmOperations(
         return response
     }
 
+    /**
+     * Resolve any role named by this interaction before anything reads its options: a role can
+     * carry hyperparameters, and which model it means depends on the provider active for this call.
+     *
+     * Interactions naming no role are returned untouched, so the common path costs nothing.
+     *
+     * Idempotent, so a subclass may call it on a path this class has already resolved: resolution
+     * replaces the role criteria with a pre-resolved one, and a second call sees no role and does
+     * nothing. That is what lets the low-level `doTransform` entry points resolve for themselves
+     * without double-resolving the `createObject` path that reaches them.
+     */
+    protected fun withRoleResolved(interaction: LlmInteraction): LlmInteraction {
+        val resolved = withRoleResolved(interaction.llm)
+        return if (resolved === interaction.llm) interaction else interaction.copy(llm = resolved)
+    }
+
+    /**
+     * As above, for the paths that carry options rather than a whole interaction - streaming,
+     * and the capability queries that pick a model without running a prompt.
+     */
+    private fun withRoleResolved(options: LlmOptions): LlmOptions =
+        if (options.criteria is ByRoleModelSelectionCriteria) {
+            modelProvider.resolveLlmOptions(options)
+        } else {
+            options
+        }
+
     protected fun chooseLlm(
         llmOptions: LlmOptions,
     ): LlmService<*> {
@@ -426,18 +486,22 @@ abstract class AbstractLlmOperations(
     }
 
     override fun supportsStreaming(options: LlmOptions): Boolean {
-        val llmService = chooseLlm(options)
+        val llmService = chooseLlm(withRoleResolved(options))
         return llmService.supportsStreaming()
     }
 
     override fun supportsThinking(options: LlmOptions): Boolean {
-        val llmService = chooseLlm(options)
+        val llmService = chooseLlm(withRoleResolved(options))
         return llmService.supportsThinking()
     }
 
     override fun createStreamingOperations(options: LlmOptions): StreamingLlmOperations {
-        val llmService = chooseLlm(options)
-        val messageStreamer = llmService.createMessageStreamer(options)
+        // Resolve once and stream with the SAME options. The streamer reads hyperparameters, so
+        // resolving only far enough to pick a model would silently drop the tuning a role carries -
+        // and streaming is the chat path, where that tuning matters most.
+        val resolved = withRoleResolved(options)
+        val llmService = chooseLlm(resolved)
+        val messageStreamer = llmService.createMessageStreamer(resolved)
         return StreamingLlmOperationsImpl(
             messageStreamer = messageStreamer,
             objectMapper = objectMapper,

embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/ExecutorAsyncer.kt

@@ -16,6 +16,7 @@
 package com.embabel.agent.spi.support
 
 import com.embabel.agent.api.common.Asyncer
+import com.embabel.common.ai.model.ModelSelectionContextHolder
 import io.micrometer.context.ContextSnapshotFactory
 import javax.annotation.concurrent.ThreadSafe
 import java.util.concurrent.CompletableFuture
@@ -24,9 +25,16 @@ import java.util.concurrent.Semaphore
 
 /**
  * Asyncer implementation that uses an Executor for async operations, propagating to worker
- * threads the [AgentProcess] (a domain concern, via [AgentProcessAccessor]) and the current
- * Micrometer Observation (via the official [ContextSnapshotFactory], so spans nest across
- * threads; a no-op when no observation is current, e.g. a NOOP registry).
+ * threads the [AgentProcess] (a domain concern, via [AgentProcessAccessor]), the
+ * [com.embabel.common.ai.model.ModelSelectionContext] (so a user's own provider key still
+ * decides model selection off the request thread), and the current Micrometer Observation
+ * (via the official [ContextSnapshotFactory], so spans nest across threads; a no-op when no
+ * observation is current, e.g. a NOOP registry).
+ *
+ * The model selection context matters here because the platform itself moves work off the
+ * calling thread - `AgentPlatform.start`, parallel actions, `OperationContext.parallelMap`.
+ * Losing it does not fail: role resolution quietly falls back to deployment configuration and
+ * serves a model the deployment pays for, on a call the user brought their own key for.
  */
 @ThreadSafe
 class ExecutorAsyncer(
@@ -36,21 +44,25 @@ class ExecutorAsyncer(
     private val contextSnapshotFactory = ContextSnapshotFactory.builder().clearMissing(true).build()
 
     override fun <T> async(block: () -> T): CompletableFuture<T> {
-        // Capture AgentProcess and the current observation from the calling thread
+        // Capture AgentProcess, model selection context and the current observation from the calling thread
         val agentProcess = AgentProcessAccessor.getValue()
+        val modelSelectionContext = ModelSelectionContextHolder.get()
         val contextSnapshot = contextSnapshotFactory.captureAll()
 
         return CompletableFuture.supplyAsync({
             contextSnapshot.setThreadLocals().use {
-                if (agentProcess != null) {
-                    AgentProcessAccessor.setValue(agentProcess)
-                    try {
+                // with() restores whatever the pooled thread held before, so nothing leaks between tasks
+                ModelSelectionContextHolder.with(modelSelectionContext) {
+                    if (agentProcess != null) {
+                        AgentProcessAccessor.setValue(agentProcess)
+                        try {
+                            block()
+                        } finally {
+                            AgentProcessAccessor.reset() // cleanup
+                        }
+                    } else {
                         block()
-                    } finally {
-                        AgentProcessAccessor.reset() // cleanup
                     }
-                } else {
-                    block()
                 }
             }
         }, executor)

embabel-agent-api/src/main/kotlin/com/embabel/agent/spi/support/ToolLoopLlmOperations.kt

@@ -151,6 +151,14 @@ open class ToolLoopLlmOperations(
         outputClass: Class<O>,
         llmRequestEvent: LlmRequestEvent<O>?,
     ): O {
+        // Shadowed deliberately, on the same terms as AbstractLlmOperations: after this line the
+        // resolved interaction IS the interaction. These are the low-level entry points on
+        // LlmOperations - reachable directly, not only through createObject - so a role named here
+        // has to resolve here too, or it silently runs on the default model. Idempotent, so the
+        // createObject path that already resolved pays nothing.
+        @Suppress("NAME_SHADOWING")
+        val interaction = withRoleResolved(interaction)
+
         val llm = chooseLlm(interaction.llm)
         val promptContributions = buildPromptContributions(interaction, llm)
 
@@ -239,6 +247,14 @@ open class ToolLoopLlmOperations(
         outputClass: Class<O>,
         llmRequestEvent: LlmRequestEvent<O>,
     ): Result<O> {
+        // Shadowed deliberately, on the same terms as AbstractLlmOperations: after this line the
+        // resolved interaction IS the interaction. These are the low-level entry points on
+        // LlmOperations - reachable directly, not only through createObject - so a role named here
+        // has to resolve here too, or it silently runs on the default model. Idempotent, so the
+        // createObject path that already resolved pays nothing.
+        @Suppress("NAME_SHADOWING")
+        val interaction = withRoleResolved(interaction)
+
         val llm = chooseLlm(interaction.llm)
         val promptContributions = buildPromptContributions(interaction, llm)
 
@@ -344,6 +360,14 @@ open class ToolLoopLlmOperations(
         outputClass: Class<O>,
         llmRequestEvent: LlmRequestEvent<O>?,
     ): ThinkingResponse<O> {
+        // Shadowed deliberately, on the same terms as AbstractLlmOperations: after this line the
+        // resolved interaction IS the interaction. These are the low-level entry points on
+        // LlmOperations - reachable directly, not only through createObject - so a role named here
+        // has to resolve here too, or it silently runs on the default model. Idempotent, so the
+        // createObject path that already resolved pays nothing.
+        @Suppress("NAME_SHADOWING")
+        val interaction = withRoleResolved(interaction)
+
         val llm = chooseLlm(interaction.llm)
         val promptContributions = buildPromptContributions(interaction, llm)
 
@@ -446,6 +470,14 @@ open class ToolLoopLlmOperations(
         outputClass: Class<O>,
         llmRequestEvent: LlmRequestEvent<O>?,
     ): Result<ThinkingResponse<O>> {
+        // Shadowed deliberately, on the same terms as AbstractLlmOperations: after this line the
+        // resolved interaction IS the interaction. These are the low-level entry points on
+        // LlmOperations - reachable directly, not only through createObject - so a role named here
+        // has to resolve here too, or it silently runs on the default model. Idempotent, so the
+        // createObject path that already resolved pays nothing.
+        @Suppress("NAME_SHADOWING")
+        val interaction = withRoleResolved(interaction)
+
         return try {
             val llm = chooseLlm(interaction.llm)
             val promptContributions = buildPromptContributions(interaction, llm)