001/*
002 * SPDX-License-Identifier: Apache-2.0
003 *
004 * Copyright 2025-2026 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 static dev.enola.ai.adk.core.Contents.replaceText;
021import static dev.enola.ai.adk.core.InvocationContexts.replaceUserContent;
022
023import com.google.adk.agents.BaseAgent;
024import com.google.adk.agents.InvocationContext;
025import com.google.adk.events.Event;
026
027import dev.enola.common.function.CheckedFunction;
028
029import io.reactivex.rxjava3.core.Flowable;
030
031import java.io.IOException;
032import java.util.List;
033
034public class UserContentReplacingAgent extends BaseAgent {
035
036    // TODO This doesn't actually really work as-is just yet; see the (failing)
037    //   UserContentReplacingAgentTest... :-( See https://github.com/google/adk-java/issues/288.
038
039    // TODO Raise ADK issue to simplify this; w.o. need for Contents, InvocationContexts
040
041    private final CheckedFunction<String, String, IOException> userContentTextReplacer;
042    private final BaseAgent delegateAgent;
043
044    public UserContentReplacingAgent(
045            String name,
046            String description,
047            CheckedFunction<String, String, IOException> userContentTextReplacer,
048            BaseAgent delegateAgent
049            // TODO List<Callbacks.BeforeAgentCallback> beforeAgentCallback,
050            // List<Callbacks.AfterAgentCallback> afterAgentCallback
051            ) {
052        super(name, description, List.of(), List.of(), List.of());
053        this.userContentTextReplacer = userContentTextReplacer;
054        this.delegateAgent = delegateAgent;
055    }
056
057    // TODO Support BeforeAgentCallback & AfterAgentCallback? Share code with LlmAgent...
058
059    // TODO Telemetry Tracer Scope... (see LLmAgent)
060
061    @Override
062    protected Flowable<Event> runAsyncImpl(InvocationContext parentContext) {
063        try {
064            var newContext = run(parentContext);
065            return delegateAgent.runAsync(newContext);
066        } catch (Throwable t) {
067            return Flowable.error(t);
068        }
069    }
070
071    @Override
072    protected Flowable<Event> runLiveImpl(InvocationContext parentContext) {
073        try {
074            var newContext = run(parentContext);
075            return delegateAgent.runLive(newContext);
076        } catch (Throwable t) {
077            return Flowable.error(t);
078        }
079    }
080
081    private InvocationContext run(InvocationContext parentContext) throws IOException {
082        var uc = parentContext.userContent();
083        var replacedUserContent = replaceText(uc, userContentTextReplacer);
084        var childContext = createInvocationContext(parentContext);
085        return replaceUserContent(childContext, replacedUserContent);
086    }
087
088    // TODO Make this protected instead of private in BaseAgent, instead of copy/pasta it here:
089    private InvocationContext createInvocationContext(InvocationContext parentContext) {
090        InvocationContext invocationContext = InvocationContext.copyOf(parentContext);
091        invocationContext.agent(this);
092        // Check for branch to be truthy (not None, not empty string),
093        if (parentContext.branch().filter(s -> !s.isEmpty()).isPresent()) {
094            invocationContext.branch(parentContext.branch().get() + "." + name());
095        }
096        return invocationContext;
097    }
098}