Repository navigation
feat(plugins): add SaveFilesAsArtifactsPlugin, porting adk-python's file-upload offload #1412
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,108 @@ | ||
| /* | ||
| * Copyright 2026 Google LLC | ||
| * | ||
| * Licensed under the Apache License, Version 2.0 (the "License"); | ||
| * you may not use this file except in compliance with the License. | ||
| * You may obtain a copy of the License at | ||
| * | ||
| * http://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, software | ||
| * distributed under the License is distributed on an "AS IS" BASIS, | ||
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| * See the License for the specific language governing permissions and | ||
| * limitations under the License. | ||
| */ | ||
| package com.google.adk.plugins; | ||
|
|
||
| import static com.google.common.collect.ImmutableMap.toImmutableMap; | ||
|
|
||
| import com.google.adk.agents.CallbackContext; | ||
| import com.google.adk.agents.InvocationContext; | ||
| import com.google.common.collect.ImmutableMap; | ||
| import java.util.Map; | ||
|
|
||
| /** | ||
| * Carries artifact versions from {@code onUserMessageCallback}, which has no {@link | ||
| * com.google.adk.events.EventActions} to write to, across to {@code beforeAgentCallback}, which | ||
| * does. | ||
| * | ||
| * <p>The hand-off rides on the session state under a {@link | ||
| * com.google.adk.sessions.State#TEMP_PREFIX} key, so {@code BaseSessionService.appendEvent} skips | ||
| * it when applying the state delta and the bookkeeping never reaches persisted session state. The | ||
| * key also carries the invocation id, so concurrent invocations sharing one session cannot read | ||
| * each other's pending versions. | ||
| * | ||
| * <p>Draining overwrites the entry with an empty map rather than removing it. {@code State.remove} | ||
| * would write the {@code State.REMOVED} sentinel into the event's state delta, which does not | ||
| * survive a JSON round trip. | ||
| */ | ||
| final class PendingArtifactDelta { | ||
|
|
||
| private static final String KEY = "temp:%s:pending_delta:%s"; | ||
|
|
||
| private PendingArtifactDelta() {} | ||
|
|
||
| /** Stores the versions produced during this invocation. */ | ||
| static void stash( | ||
| InvocationContext invocationContext, String pluginName, ImmutableMap<String, Integer> delta) { | ||
| invocationContext | ||
| .session() | ||
| .state() | ||
| .put(key(pluginName, invocationContext.invocationId()), delta); | ||
| } | ||
|
|
||
| /** | ||
| * Returns the stored versions and clears them, so only the first agent callback reports them. | ||
| * | ||
| * <p>The write is guarded on there being something to clear, and the guard is load-bearing: a | ||
| * write through {@link CallbackContext#state()} sets the state delta, and {@code BaseAgent} emits | ||
| * an event for any before-agent callback that left one. Clearing unconditionally would therefore | ||
| * emit an event carrying an empty artifact delta for every agent after the first. | ||
| */ | ||
| static ImmutableMap<String, Integer> drain(CallbackContext callbackContext, String pluginName) { | ||
| String key = key(pluginName, callbackContext.invocationId()); | ||
| ImmutableMap<String, Integer> pending = read(callbackContext.state().get(key)); | ||
| if (!pending.isEmpty()) { | ||
| callbackContext.state().put(key, ImmutableMap.of()); | ||
| } | ||
| return pending; | ||
| } | ||
|
|
||
| /** | ||
| * Discards versions left behind by an invocation that ended before any agent ran, which happens | ||
| * when a {@code beforeRunCallback} on another plugin halts the run. There is no {@link | ||
| * com.google.adk.events.EventActions} to report them on at that point, so they are dropped rather | ||
| * than returned. | ||
| * | ||
| * <p>Only writes when something is actually stashed: an unconditional write would add an entry | ||
| * for every invocation that uploaded nothing, which is more state noise than the leak it | ||
| * prevents. | ||
| */ | ||
| static void clear(InvocationContext invocationContext, String pluginName) { | ||
| Map<String, Object> state = invocationContext.session().state(); | ||
| String key = key(pluginName, invocationContext.invocationId()); | ||
| if (!read(state.get(key)).isEmpty()) { | ||
| state.put(key, ImmutableMap.of()); | ||
| } | ||
| } | ||
|
|
||
| private static String key(String pluginName, String invocationId) { | ||
| return KEY.formatted(pluginName, invocationId); | ||
| } | ||
|
|
||
| private static ImmutableMap<String, Integer> read(Object stashed) { | ||
| if (!(stashed instanceof Map<?, ?> entries)) { | ||
| return ImmutableMap.of(); | ||
| } | ||
| return entries.entrySet().stream() | ||
| .filter(PendingArtifactDelta::isVersionEntry) | ||
| .collect( | ||
| toImmutableMap(entry -> (String) entry.getKey(), entry -> (Integer) entry.getValue())); | ||
| } | ||
|
|
||
| /** State values survive a JSON round trip untyped, so each entry is checked before it is kept. */ | ||
| private static boolean isVersionEntry(Map.Entry<?, ?> entry) { | ||
| return entry.getKey() instanceof String && entry.getValue() instanceof Integer; | ||
| } | ||
| } | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,227 @@ | ||
| /* | ||
| * Copyright 2026 Google LLC | ||
| * | ||
| * Licensed under the Apache License, Version 2.0 (the "License"); | ||
| * you may not use this file except in compliance with the License. | ||
| * You may obtain a copy of the License at | ||
| * | ||
| * http://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, software | ||
| * distributed under the License is distributed on an "AS IS" BASIS, | ||
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| * See the License for the specific language governing permissions and | ||
| * limitations under the License. | ||
| */ | ||
| package com.google.adk.plugins; | ||
|
|
||
| import static com.google.common.collect.ImmutableList.toImmutableList; | ||
| import static com.google.common.collect.ImmutableMap.toImmutableMap; | ||
|
|
||
| import com.google.adk.agents.BaseAgent; | ||
| import com.google.adk.agents.CallbackContext; | ||
| import com.google.adk.agents.InvocationContext; | ||
| import com.google.common.collect.ImmutableList; | ||
| import com.google.common.collect.ImmutableMap; | ||
| import com.google.genai.types.Blob; | ||
| import com.google.genai.types.Content; | ||
| import com.google.genai.types.Part; | ||
| import io.reactivex.rxjava3.core.Completable; | ||
| import io.reactivex.rxjava3.core.Flowable; | ||
| import io.reactivex.rxjava3.core.Maybe; | ||
| import io.reactivex.rxjava3.core.Single; | ||
| import java.util.List; | ||
| import java.util.Optional; | ||
| import org.slf4j.Logger; | ||
| import org.slf4j.LoggerFactory; | ||
|
|
||
| /** | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. LOW: This Javadoc runs to about ten sentences, while Kotlin's covers the same points in three (L34-L45). Please trim it and use |
||
| * Plugin that saves files embedded in user messages as artifacts. | ||
| * | ||
| * <p>This allows users to upload files in the chat experience and have those files available to the | ||
| * agent within the current session. Each {@code inlineData} part of the incoming user message is | ||
| * written to the configured {@link com.google.adk.artifacts.BaseArtifactService} and replaced, in | ||
| * the message that reaches the model, by a short text placeholder naming the artifact. The bytes | ||
| * themselves are therefore stored once and not resent on every turn. | ||
| * | ||
| * <p>The artifact name is taken from {@link Blob#displayName()} when present, so an uploaded {@code | ||
| * report.pdf} is stored under that name. When the blob carries no display name, a name is generated | ||
| * from the invocation id and the part index. Uploading the same name again saves a new version of | ||
| * it, and the version saved is what this plugin reports for that name. | ||
| * | ||
| * <p>Add the {@code load_artifacts} tool to the agent, or load the artifacts from your own tool, to | ||
| * let the model read the stored bytes back. | ||
| * | ||
| * <p>Register it on the runner: | ||
| * | ||
| * <pre>{@code | ||
| * Runner runner = | ||
| * Runner.builder() | ||
| * .agent(agent) | ||
| * .appName("my-app") | ||
| * .artifactService(new InMemoryArtifactService()) | ||
| * .sessionService(new InMemorySessionService()) | ||
| * .plugins(new SaveFilesAsArtifactsPlugin()) | ||
| * .build(); | ||
| * }</pre> | ||
| * | ||
| * <p>The plugin is a no-op when no artifact service is configured on the runner. | ||
| */ | ||
| public class SaveFilesAsArtifactsPlugin extends BasePlugin { | ||
|
|
||
| /** Name used when the plugin is constructed without an explicit one. Matches adk-python's. */ | ||
| public static final String DEFAULT_NAME = "save_files_as_artifacts_plugin"; | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. LOW: |
||
|
|
||
| private static final Logger logger = LoggerFactory.getLogger(SaveFilesAsArtifactsPlugin.class); | ||
|
|
||
| private static final String GENERATED_FILE_NAME = "artifact_%s_%d"; | ||
| private static final String PLACEHOLDER_TEXT = "[Uploaded Artifact: \"%s\"]"; | ||
|
|
||
| public SaveFilesAsArtifactsPlugin() { | ||
| this(DEFAULT_NAME); | ||
| } | ||
|
|
||
| public SaveFilesAsArtifactsPlugin(String name) { | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. HIGH: Please port |
||
| super(name); | ||
| } | ||
|
|
||
| @Override | ||
| public Maybe<Content> onUserMessageCallback( | ||
| InvocationContext invocationContext, Content userMessage) { | ||
| if (invocationContext.artifactService() == null) { | ||
| logger.warn("No artifact service is configured; plugin '{}' is disabled.", getName()); | ||
| return Maybe.empty(); | ||
| } | ||
| ImmutableList<Part> parts = | ||
| ImmutableList.copyOf(userMessage.parts().orElse(ImmutableList.of())); | ||
| if (parts.stream().noneMatch(SaveFilesAsArtifactsPlugin::hasInlineData)) { | ||
| return Maybe.empty(); | ||
| } | ||
| return Flowable.range(0, parts.size()) | ||
| .concatMapSingle(index -> savePart(invocationContext, parts.get(index), index)) | ||
| .collect(toImmutableList()) | ||
| .map(results -> rebuildMessage(invocationContext, userMessage, results)) | ||
| .filter(Optional::isPresent) | ||
| .map(Optional::get); | ||
| } | ||
|
|
||
| /** | ||
| * Records the artifact versions stashed by {@link #onUserMessageCallback} on the first event | ||
| * actions of the invocation. {@code onUserMessageCallback} runs before any {@link | ||
| * com.google.adk.events.EventActions} exists, so the versions cannot be reported from there. | ||
| * | ||
| * <p>The reporting is a side effect and the return is always empty, deliberately: {@code | ||
| * PluginManager} stops at the first plugin that returns a value, so returning content here would | ||
| * both skip every later plugin's callback and halt the agent. | ||
| */ | ||
| @Override | ||
| public Maybe<Content> beforeAgentCallback(BaseAgent agent, CallbackContext callbackContext) { | ||
| PendingArtifactDelta.drain(callbackContext, getName()) | ||
| .forEach(callbackContext.eventActions().artifactDelta()::put); | ||
| return Maybe.empty(); | ||
| } | ||
|
|
||
| /** | ||
| * Discards a stash that {@link #beforeAgentCallback} never got to report, which happens when a | ||
| * {@code beforeRunCallback} on another plugin halts the invocation before any agent runs. | ||
| * | ||
| * <p>Only a run that completes reaches this hook; {@link #onRunErrorCallback} clears the same | ||
| * stash when one fails. | ||
| */ | ||
| @Override | ||
| public Completable afterRunCallback(InvocationContext invocationContext) { | ||
| PendingArtifactDelta.clear(invocationContext, getName()); | ||
| return Completable.complete(); | ||
| } | ||
|
|
||
| /** | ||
| * Discards the stash when the invocation fails, the window {@link #afterRunCallback} never sees. | ||
| * | ||
| * <p>The versions cannot be reported from here: this hook is handed an {@link InvocationContext}, | ||
| * which carries no {@link com.google.adk.events.EventActions} to write an artifact delta to. The | ||
| * clear is idempotent, so a run reaching both hooks is no different from one reaching either. | ||
| */ | ||
| @Override | ||
| public Completable onRunErrorCallback(InvocationContext invocationContext, Throwable error) { | ||
| PendingArtifactDelta.clear(invocationContext, getName()); | ||
| return Completable.complete(); | ||
| } | ||
|
|
||
| /** Saves one part if it carries inline data, leaving every other part untouched. */ | ||
| private Single<SavedPart> savePart(InvocationContext invocationContext, Part part, int index) { | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. MEDIUM: Please port the 20 MB cap from Kotlin (L78-L89; Python L116-L128). An oversized file is not saved and is replaced with an |
||
| if (!hasInlineData(part)) { | ||
| return Single.just(SavedPart.unchanged(part)); | ||
| } | ||
| String fileName = resolveFileName(invocationContext, part, index); | ||
| return invocationContext | ||
| .artifactService() | ||
| .saveArtifact( | ||
| invocationContext.appName(), | ||
| invocationContext.userId(), | ||
| invocationContext.session().id(), | ||
| fileName, | ||
| part) | ||
| .map(version -> SavedPart.saved(placeholderFor(fileName), fileName, version)) | ||
| .onErrorReturn(error -> keepOriginal(part, fileName, error)); | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. HIGH: |
||
| } | ||
|
|
||
| /** A failed save must not fail the invocation: the original part is passed through unchanged. */ | ||
| private SavedPart keepOriginal(Part part, String fileName, Throwable error) { | ||
| logger.error("Failed to save artifact '{}'; keeping the original part.", fileName, error); | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. HIGH: This logs the user-chosen file name, and the logged exception can repeat it. To avoid that, Kotlin logs only the part index and the exception type (L96-L97), for example |
||
| return SavedPart.unchanged(part); | ||
| } | ||
|
|
||
| /** Returns the rewritten message, or empty when no part was actually offloaded. */ | ||
| private Optional<Content> rebuildMessage( | ||
| InvocationContext invocationContext, Content userMessage, List<SavedPart> results) { | ||
| ImmutableMap<String, Integer> delta = toArtifactDelta(results); | ||
| if (delta.isEmpty()) { | ||
| return Optional.empty(); | ||
| } | ||
| PendingArtifactDelta.stash(invocationContext, getName(), delta); | ||
| ImmutableList<Part> parts = results.stream().map(SavedPart::part).collect(toImmutableList()); | ||
| return Optional.of(userMessage.toBuilder().parts(parts).build()); | ||
| } | ||
|
|
||
| private static ImmutableMap<String, Integer> toArtifactDelta(List<SavedPart> results) { | ||
| return results.stream() | ||
| .filter(SavedPart::isSaved) | ||
| .collect( | ||
| toImmutableMap(SavedPart::savedFileName, SavedPart::version, (older, newer) -> newer)); | ||
| } | ||
|
|
||
| private static String resolveFileName(InvocationContext invocationContext, Part part, int index) { | ||
| return part.inlineData() | ||
| .flatMap(Blob::displayName) | ||
| .filter(displayName -> !displayName.isEmpty()) | ||
| .orElseGet(() -> GENERATED_FILE_NAME.formatted(invocationContext.invocationId(), index)); | ||
| } | ||
|
|
||
| private static Part placeholderFor(String fileName) { | ||
| return Part.fromText(PLACEHOLDER_TEXT.formatted(fileName)); | ||
| } | ||
|
|
||
| private static boolean hasInlineData(Part part) { | ||
| return part.inlineData().isPresent(); | ||
| } | ||
|
|
||
| /** One input part after the offload attempt: either untouched, or replaced by a placeholder. */ | ||
| private record SavedPart(Part part, Optional<String> fileName, int version) { | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. LOW: Kotlin handles this in one loop that builds the new parts and a map of versions (L66-L113). Using that structure would remove |
||
|
|
||
| static SavedPart unchanged(Part part) { | ||
| return new SavedPart(part, Optional.empty(), 0); | ||
| } | ||
|
|
||
| static SavedPart saved(Part placeholder, String fileName, int version) { | ||
| return new SavedPart(placeholder, Optional.of(fileName), version); | ||
| } | ||
|
|
||
| boolean isSaved() { | ||
| return fileName.isPresent(); | ||
| } | ||
|
|
||
| String savedFileName() { | ||
| return fileName.orElseThrow(); | ||
| } | ||
| } | ||
| } | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
HIGH: The stash leaves one key in stored session state for every upload turn. Running two turns through
InMemoryRunnerleft twotemp:save_files_as_artifacts_plugin:pending_delta:<invocationId>={}entries. The direct write mutates theSessionobject stored byInMemorySessionService(L340-L358), andFirestoreSessionServicepersists it as well (L657-L696). Please follow Kotlin and keep the versions in per-invocation data (L107-L113, L120-L128), which in Java isInvocationContext.callbackContextData().