Ability to register hand-crafted agents to the registry

sentinel-ai-configured-agents/src/main/java/com/phonepe/sentinelai/configuredagents/AgentConfiguration.java

@@ -1,10 +1,18 @@
 package com.phonepe.sentinelai.configuredagents;
 
 import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
 import com.phonepe.sentinelai.configuredagents.capabilities.AgentCapability;
-import lombok.*;
+import com.phonepe.sentinelai.core.utils.JsonUtils;
+import lombok.AllArgsConstructor;
+import lombok.Builder;
+import lombok.NonNull;
+import lombok.Singular;
+import lombok.Value;
+import lombok.With;
 
 import java.util.List;
+import java.util.Objects;
 
 /**
  * Configuration for dynamically spun up {@link ConfiguredAgent}.
@@ -47,4 +55,16 @@ public class AgentConfiguration {
     @Singular
     List<AgentCapability> capabilities;
 
+    public static AgentConfiguration fixConfiguration(@NonNull AgentConfiguration configuration,
+                                                      final ObjectMapper mapper) {
+        return new AgentConfiguration(
+                configuration.getAgentName(),
+                configuration.getDescription(),
+                configuration.getPrompt(),
+                Objects.requireNonNullElseGet(configuration.getInputSchema(),
+                                              () -> JsonUtils.schemaForPrimitive(String.class, "data", mapper)),
+                Objects.requireNonNullElseGet(configuration.getOutputSchema(),
+                                              () -> JsonUtils.schema(String.class)),
+                Objects.requireNonNullElseGet(configuration.getCapabilities(), List::of));
+    }
 }

sentinel-ai-configured-agents/src/main/java/com/phonepe/sentinelai/configuredagents/AgentRegistry.java

@@ -19,6 +19,7 @@
 import java.nio.file.Files;
 import java.nio.file.Paths;
 import java.util.*;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.BiFunction;
 import java.util.function.Predicate;
 
@@ -48,9 +49,11 @@ public class AgentRegistry<R, T, A extends Agent<R, T, A>> implements AgentExten
 
     private final AgentMetadataAccessMode agentMetadataAccessMode;
 
+    private final Map<String, ConfiguredAgent> externallyRegisteredAgents = new ConcurrentHashMap<>();
+
     private final SimpleCache<ConfiguredAgent> agentCache;
 
-    private A agent;
+    private A parent;
 
     @Builder
     public AgentRegistry(
@@ -65,10 +68,15 @@ public AgentRegistry(
         this.agentMetadataAccessMode = Objects.requireNonNullElse(agentMetadataAccessMode,
                                                                   DEFAULT_METADATA_ACCESS_MODE);
         this.agentCache = new SimpleCache<>(agentId -> {
+            final var externallyRegisteredAgent = externallyRegisteredAgents.get(agentId);
+            if (null != externallyRegisteredAgent) {
+                log.info("Using externally registered agent for: {}", agentId);
+                return externallyRegisteredAgent;
+            }
             log.info("Building new agent for: {}", agentId);
             return agentFactory.apply(
                     agentSource.read(agentId)
-                            .orElse(null), this.agent);
+                            .orElse(null), this.parent);
         });
     }
 
@@ -92,20 +100,39 @@ public List<AgentMetadata> loadAgentsFromContent(byte[] content) {
         return agents;
     }
 
+    /**
+     * Register the configuration of a new agent. Will be created when needed using the
+     * {@link ConfiguredAgentFactory::createAgent(AgentMetadata, Agent)} method.
+     *
+     * @param configuration Agent configuration
+     * @return The metadata of the configured agent if successful, empty otherwise
+     */
     public Optional<AgentMetadata> configureAgent(@NonNull final AgentConfiguration configuration) {
-        final var fixedConfig = new AgentConfiguration(
-                configuration.getAgentName(),
-                configuration.getDescription(),
-                configuration.getPrompt(),
-                Objects.requireNonNullElseGet(configuration.getInputSchema(),
-                                              () -> JsonUtils.schemaForPrimitive(String.class, "data", mapper)),
-                Objects.requireNonNullElseGet(configuration.getOutputSchema(),
-                                              () -> JsonUtils.schema(String.class)),
-                Objects.requireNonNullElseGet(configuration.getCapabilities(), List::of));
+        final var fixedConfig = AgentConfiguration.fixConfiguration(configuration, mapper);
         final var agentId = AgentUtils.id(fixedConfig.getAgentName());
         return agentSource.save(agentId, fixedConfig);
     }
 
+    /**
+     * Register a new agent. This will be treated as a singleton and reused when needed.
+     *
+     * @param agent The agent to be registered
+     * @return The metadata of the configured agent if successful, empty otherwise
+     */
+    public Optional<AgentMetadata> configureAgent(@NonNull final RegisterableAgent<? extends RegisterableAgent<?>> agent) {
+        final var fixedConfig = AgentConfiguration.fixConfiguration(agent.agentConfiguration(), mapper);
+        final var agentId = AgentUtils.id(fixedConfig.getAgentName());
+        final var metadata = agentSource.save(agentId, fixedConfig);
+        if (metadata.isPresent()) {
+            externallyRegisteredAgents.put(agentId, new ConfiguredAgent(agent));
+            log.info("Registered external agent: {} ({})", agent.agentConfiguration().getAgentName(), agentId);
+        }
+        else {
+            log.warn("Agent with id {} could not be configured", agentId);
+        }
+        return metadata;
+    }
+
     @Tool("Get agent metadata. Use this to get agent id, name, description, input and output schema etc")
     public ExposedAgentMetadata getAgentMetadata(
             @JsonPropertyDescription("ID of the agent to get metadata for") String agentId) {
@@ -240,7 +267,7 @@ public void consume(JsonNode output, A agent) {
 
     @Override
     public void onExtensionRegistrationCompleted(A agent) {
-        this.agent = agent;
+        this.parent = agent;
     }
 
     private static IllegalArgumentException agentNotFoundError(String agentId) {

sentinel-ai-configured-agents/src/main/java/com/phonepe/sentinelai/configuredagents/ConfiguredAgent.java

@@ -1,8 +1,10 @@
 package com.phonepe.sentinelai.configuredagents;
 
-import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.phonepe.sentinelai.core.agent.*;
+import com.phonepe.sentinelai.core.earlytermination.NeverTerminateEarlyStrategy;
+import com.phonepe.sentinelai.core.errorhandling.DefaultErrorHandler;
+import com.phonepe.sentinelai.core.outputvalidation.DefaultOutputValidator;
 import com.phonepe.sentinelai.core.tools.ToolBox;
 import com.phonepe.sentinelai.core.utils.JsonUtils;
 import lombok.SneakyThrows;
@@ -18,63 +20,39 @@
 @SuppressWarnings({"unused", "FieldCanBeLocal"})
 public class ConfiguredAgent {
 
-    private final String name;
-    private final String description;
-    private final RootAgent rootAgent;
-    private final JsonNode inputSchema;
-    private final JsonNode outputSchema;
+    public static final class RootAgent extends RegisterableAgent<RootAgent> {
 
-    public static final class RootAgent extends Agent<String, String, RootAgent> {
-
-        private final String name;
-        private final JsonNode outputSchema;
+        private final AgentConfiguration agentConfiguration;
 
         public RootAgent(
-                final String name,
-                final String prompt,
-                final JsonNode outputSchema,
-                final List<AgentExtension<String, String, RootAgent>> extensions,
+                final AgentConfiguration agentConfiguration,
+                final List<AgentExtension<String, String, RootAgent>> agentExtensions,
                 final ToolBox toolBox) {
-            super(String.class,
-                  prompt,
+            super(agentConfiguration,
                   AgentSetup.builder().build(),
-                  extensions,
-                  Map.of());
-            this.name = name;
-            this.outputSchema = outputSchema;
+                  agentExtensions,
+                  Map.of(),
+                  new ApproveAllToolRuns<>(),
+                  new DefaultOutputValidator<>(),
+                  new DefaultErrorHandler<>(),
+                  new NeverTerminateEarlyStrategy());
+            this.agentConfiguration = agentConfiguration;
             this.registerToolbox(toolBox);
         }
 
-        @Override
-        public String name() {
-            return name;
-        }
+    }
 
-        @Override
-        protected JsonNode outputSchema() {
-            return outputSchema;
-        }
+    private final Agent<String, String, ? extends RegisterableAgent<?>> rootAgent;
 
-        @Override
-        protected String translateData(JsonNode output, AgentSetup agentSetup) throws JsonProcessingException {
-            return agentSetup.getMapper()
-                    .writeValueAsString(output);
-        }
+    public ConfiguredAgent(Agent<String, String, ? extends RegisterableAgent<?>> agent) {
+        this.rootAgent = agent;
     }
 
     public ConfiguredAgent(
-            final String name,
-            final String description,
-            final String prompt,
+            final AgentConfiguration agentConfiguration,
             final List<AgentExtension<String, String, RootAgent>> rootAgentExtensions,
-            final ToolBox availableTools,
-            final JsonNode inputSchema,
-            final JsonNode outputSchema) {
-        this.name = name;
-        this.description = description;
-        this.inputSchema = inputSchema;
-        this.outputSchema = outputSchema;
-        this.rootAgent = new RootAgent(name, prompt, outputSchema, rootAgentExtensions, availableTools);
+            final ToolBox availableTools) {
+        this.rootAgent = new RootAgent(agentConfiguration, rootAgentExtensions, availableTools);
     }
 
     @SneakyThrows
@@ -86,8 +64,7 @@ public final CompletableFuture<AgentOutput<JsonNode>> executeAsync(AgentInput<Js
                                               input.getRequestMetadata(),
                                               input.getOldMessages(),
                                               input.getAgentSetup()
-                                      )
-                                     )
+                                      ))
                 .thenApply(output -> {
                     try {
                         final var json = Objects.requireNonNullElseGet(mapper, JsonUtils::createMapper)

sentinel-ai-configured-agents/src/main/java/com/phonepe/sentinelai/configuredagents/ConfiguredAgentFactory.java

@@ -1,5 +1,7 @@
 package com.phonepe.sentinelai.configuredagents;
 
+import com.phonepe.sentinelai.configuredagents.capabilities.AgentCapability;
+import com.phonepe.sentinelai.configuredagents.capabilities.AgentCapabilityVisitor;
 import com.phonepe.sentinelai.configuredagents.capabilities.impl.AgentCustomToolCapability;
 import com.phonepe.sentinelai.configuredagents.capabilities.impl.AgentMCPCapability;
 import com.phonepe.sentinelai.configuredagents.capabilities.impl.AgentMemoryCapability;
@@ -12,8 +14,6 @@
 import com.phonepe.sentinelai.core.tools.ToolBox;
 import com.phonepe.sentinelai.toolbox.mcp.MCPToolBox;
 import com.phonepe.sentinelai.toolbox.remotehttp.HttpToolBox;
-import com.phonepe.sentinelai.configuredagents.capabilities.AgentCapability;
-import com.phonepe.sentinelai.configuredagents.capabilities.AgentCapabilityVisitor;
 import lombok.Builder;
 import lombok.NonNull;
 import lombok.extern.slf4j.Slf4j;
@@ -138,13 +138,9 @@ public Void visit(ParentToolInheritanceCapability parentToolInheritanceCapabilit
         toolBoxes.addAll(extensions); //Because all extensions are also toolboxes
 
         return new ConfiguredAgent(
-                agentConfiguration.getAgentName(),
-                agentConfiguration.getDescription(),
-                agentConfiguration.getPrompt(),
+                agentConfiguration,
                 extensions,
-                new ComposingToolBox(toolBoxes, Set.of()),
-                agentConfiguration.getInputSchema(),
-                agentConfiguration.getOutputSchema());
+                new ComposingToolBox(toolBoxes, Set.of()));
     }
 
 }

sentinel-ai-configured-agents/src/main/java/com/phonepe/sentinelai/configuredagents/RegisterableAgent.java

@@ -0,0 +1,108 @@
+package com.phonepe.sentinelai.configuredagents;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.JsonNode;
+import com.phonepe.sentinelai.core.agent.Agent;
+import com.phonepe.sentinelai.core.agent.AgentExtension;
+import com.phonepe.sentinelai.core.agent.AgentSetup;
+import com.phonepe.sentinelai.core.agent.ApproveAllToolRuns;
+import com.phonepe.sentinelai.core.earlytermination.EarlyTerminationStrategy;
+import com.phonepe.sentinelai.core.earlytermination.NeverTerminateEarlyStrategy;
+import com.phonepe.sentinelai.core.errorhandling.DefaultErrorHandler;
+import com.phonepe.sentinelai.core.errorhandling.ErrorResponseHandler;
+import com.phonepe.sentinelai.core.outputvalidation.DefaultOutputValidator;
+import com.phonepe.sentinelai.core.outputvalidation.OutputValidator;
+import com.phonepe.sentinelai.core.tools.ExecutableTool;
+import com.phonepe.sentinelai.core.tools.ToolRunApprovalSeeker;
+import com.phonepe.sentinelai.core.utils.JsonUtils;
+import lombok.NonNull;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+
+/**
+ * An agent that can be registered in the {@link AgentRegistry}.
+ * @param <T> The type of the agent extending this class.
+ */
+@SuppressWarnings("java:S107")
+public abstract class RegisterableAgent<T  extends RegisterableAgent<T>> extends Agent<String, String, T> {
+
+    private final AgentConfiguration agentConfiguration;
+
+    protected RegisterableAgent(
+            AgentConfiguration agentConfiguration,
+            @NonNull AgentSetup setup,
+            List<AgentExtension<String, String, T>> agentExtensions,
+            Map<String, ExecutableTool> knownTools) {
+        this(agentConfiguration,
+             setup,
+             agentExtensions,
+             knownTools,
+             new ApproveAllToolRuns<>(),
+             new DefaultOutputValidator<>(),
+             new DefaultErrorHandler<>(),
+             new NeverTerminateEarlyStrategy());
+    }
+
+    protected RegisterableAgent(
+            AgentConfiguration agentConfiguration,
+            @NonNull AgentSetup setup,
+            List<AgentExtension<String, String, T>> agentExtensions,
+            Map<String, ExecutableTool> knownTools,
+            ToolRunApprovalSeeker<String, String, T> toolRunApprovalSeeker,
+            OutputValidator<String, String> outputValidator,
+            ErrorResponseHandler<String> errorHandler,
+            EarlyTerminationStrategy earlyTerminationStrategy) {
+        super(String.class,
+              agentConfiguration.getPrompt(),
+              setup,
+              agentExtensions,
+              knownTools,
+              toolRunApprovalSeeker,
+              outputValidator,
+              errorHandler,
+              earlyTerminationStrategy);
+        this.agentConfiguration = AgentConfiguration.fixConfiguration(
+                agentConfiguration, Objects.requireNonNullElseGet(setup.getMapper(), JsonUtils::createMapper));
+    }
+
+    public final AgentConfiguration agentConfiguration() {
+        return agentConfiguration;
+    }
+
+
+    /**
+     * You cannot override this here. Set the correct schema in the agent configuration.
+     * @return The output schema for this agent
+     */
+    @Override
+    protected final JsonNode outputSchema() {
+        return agentConfiguration.getOutputSchema();
+    }
+
+    /**
+     * The name of the agent as specified in the configuration. Feel free to override if needed for more exotic
+     * implementations
+     * @return The name of the agent
+     */
+    @Override
+    public String name() {
+        return agentConfiguration.getAgentName();
+    }
+
+    /**
+     * We don't do much here. Basically if the model sends out put in the required format, we just pass it through.
+     * You can override this if something more exotic needs to be done
+     * @param output     The model output
+     * @param agentSetup The agent setup
+     * @return A String serialized version of the output. It is a little wasteful to serialize and deserialize again, but
+     *         this keeps things simple.
+     * @throws JsonProcessingException If serialization fails
+     */
+    @Override
+    protected String translateData(JsonNode output, AgentSetup agentSetup) throws JsonProcessingException {
+        return agentSetup.getMapper()
+                .writeValueAsString(output);
+    }
+}