001/*
002 * SPDX-License-Identifier: Apache-2.0
003 *
004 * Copyright 2025 The Enola <https://enola.dev> Authors
005 *
006 * Licensed under the Apache License, Version 2.0 (the "License");
007 * you may not use this file except in compliance with the License.
008 * You may obtain a copy of the License at
009 *
010 *     https://www.apache.org/licenses/LICENSE-2.0
011 *
012 * Unless required by applicable law or agreed to in writing, software
013 * distributed under the License is distributed on an "AS IS" BASIS,
014 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
015 * See the License for the specific language governing permissions and
016 * limitations under the License.
017 */
018package dev.enola.ai.adk.core;
019
020import com.google.adk.agents.BaseAgent;
021import com.google.adk.artifacts.BaseArtifactService;
022import com.google.adk.artifacts.InMemoryArtifactService;
023import com.google.adk.runner.Runner;
024import com.google.adk.sessions.BaseSessionService;
025import com.google.adk.sessions.InMemorySessionService;
026import com.google.common.collect.ImmutableList;
027import com.google.common.collect.ImmutableMap;
028
029import org.slf4j.Logger;
030import org.slf4j.LoggerFactory;
031
032import java.util.Map;
033import java.util.Set;
034import java.util.concurrent.ConcurrentHashMap;
035
036// Copy/paste from com.google.adk.web.AdkWebServer#RunnerService, with the following changes:
037//   * De-Spring-ified by removing @Component @Autowired @Qualifier("loadedAgentRegistry")
038//   * getRunner() IllegalArgumentException instead of ResponseStatusException(HttpStatus.NOT_FOUND)
039//   * Added additional convenience constructors
040public class RunnersCache {
041
042    private static final Logger log = LoggerFactory.getLogger(RunnersCache.class);
043
044    private final Map<String, BaseAgent> agentRegistry;
045    private final BaseArtifactService artifactService;
046    private final BaseSessionService sessionService;
047    private final Map<String, Runner> runnerCache = new ConcurrentHashMap<>();
048
049    public RunnersCache(
050            Map<String, BaseAgent> agentRegistry,
051            BaseArtifactService artifactService,
052            BaseSessionService sessionService) {
053        this.agentRegistry = ImmutableMap.copyOf(agentRegistry); // skipcq: JAVA-E1086
054        this.artifactService = artifactService;
055        this.sessionService = sessionService;
056    }
057
058    public RunnersCache(Map<String, BaseAgent> agentRegistry) {
059        this(agentRegistry, new InMemoryArtifactService(), new InMemorySessionService());
060    }
061
062    public RunnersCache(Iterable<BaseAgent> agents) {
063        this(Agents.toMap(agents));
064    }
065
066    public RunnersCache(BaseAgent agents) {
067        this(Agents.toMap(ImmutableList.of(agents)));
068    }
069
070    public Set<String> appNames() {
071        return agentRegistry.keySet();
072    }
073
074    /**
075     * Gets the Runner instance for a given application name. Handles potential agent engine ID
076     * overrides.
077     *
078     * @param appName The application name requested by the user.
079     * @return A configured Runner instance.
080     * @throws IllegalArgumentException If no agent matching appName found in the agent registry.
081     */
082    public Runner getRunner(String appName) throws IllegalArgumentException {
083        return runnerCache.computeIfAbsent(
084                appName,
085                key -> {
086                    BaseAgent agent = agentRegistry.get(key);
087                    if (agent == null) {
088                        log.error(
089                                "Agent/App named '{}' not found in registry. Available apps: {}",
090                                key,
091                                agentRegistry.keySet());
092                        throw new IllegalArgumentException("Agent/App not found: " + key);
093                    }
094                    log.info(
095                            "RunnerService: Creating Runner for appName: {}, using agent"
096                                    + " definition: {}",
097                            appName,
098                            agent.name());
099                    return new Runner(agent, appName, this.artifactService, this.sessionService);
100                });
101    }
102}