Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 15 additions & 11 deletions core/src/main/java/com/google/adk/runner/Runner.java
Original file line number Diff line number Diff line change
Expand Up @@ -347,6 +347,9 @@ public Completable close() {
/**
* Appends a new user message to the session history with optional state delta.
*
* <p>{@code newMessage} is never modified; when inline blobs are saved as artifacts, the appended
* event carries a copy in which the blob data is replaced by placeholders.
*
* @throws IllegalArgumentException if message has no parts.
*/
private Single<Event> appendNewMessageToSession(
Expand All @@ -357,12 +360,16 @@ private Single<Event> appendNewMessageToSession(
@Nullable Map<String, Object> stateDelta) {
checkArgument(newMessage.parts().isPresent(), "No parts in the new_message.");

Content messageToAppend = newMessage;
Completable saveArtifactsFlow = Completable.complete();
if (this.artifactService != null && saveInputBlobsAsArtifacts) {
// The runner directly saves the artifacts (if applicable) in the user message and replaces
// the artifact data with a file name placeholder.
for (int i = 0; i < newMessage.parts().get().size(); i++) {
Part part = newMessage.parts().get().get(i);
// the artifact data with a file name placeholder. The rewrite happens on a copy of the parts
// list: the caller's list may be immutable, and the caller does not expect the message it
// passed to runAsync to be modified.
List<Part> parts = new ArrayList<>(newMessage.parts().get());
for (int i = 0; i < parts.size(); i++) {
Part part = parts.get(i);
if (part.inlineData().isEmpty()) {
continue;
}
Expand All @@ -373,22 +380,19 @@ private Single<Event> appendNewMessageToSession(
.saveArtifact(this.appName, session.userId(), session.id(), fileName, part)
.ignoreElement());

newMessage
.parts()
.get()
.set(
i,
Part.fromText(
"Uploaded file: " + fileName + ". It has been saved to the artifacts"));
parts.set(
i,
Part.fromText("Uploaded file: " + fileName + ". It has been saved to the artifacts"));
}
messageToAppend = newMessage.toBuilder().parts(ImmutableList.copyOf(parts)).build();
}
// Appends only. We do not yield the event because it's not from the model.
Event.Builder eventBuilder =
Event.builder()
.id(Event.generateEventId())
.invocationId(invocationContext.invocationId())
.author("user")
.content(newMessage);
.content(messageToAppend);

// Add state delta if provided
if (stateDelta != null && !stateDelta.isEmpty()) {
Expand Down
271 changes: 271 additions & 0 deletions core/src/test/java/com/google/adk/runner/RunnerTest.java
Comment thread
MiloszSobczyk marked this conversation as resolved.
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
import static com.google.adk.testing.TestUtils.createTestLlm;
import static com.google.adk.testing.TestUtils.createTextLlmResponse;
import static com.google.adk.testing.TestUtils.simplifyEvents;
import static com.google.common.collect.ImmutableList.toImmutableList;
import static com.google.common.truth.Truth.assertThat;
import static java.nio.charset.StandardCharsets.UTF_8;
import static java.util.Arrays.stream;
Expand Down Expand Up @@ -50,6 +51,7 @@
import com.google.adk.apps.App;
import com.google.adk.apps.ResumabilityConfig;
import com.google.adk.artifacts.BaseArtifactService;
import com.google.adk.artifacts.InMemoryArtifactService;
import com.google.adk.events.Event;
import com.google.adk.flows.llmflows.Functions;
import com.google.adk.models.LlmRequest;
Expand Down Expand Up @@ -3042,6 +3044,275 @@ public void runner_executesSaveArtifactFlow() {
assertThat(simplifyEvents(events.values())).containsExactly("test agent: from llm");
}

private static final String BLOB_MIME_TYPE = "example/octet-stream";
private static final String BLOB_PAYLOAD = "blob payload";
private static final String PLACEHOLDER_FORMAT =
"Uploaded file: %s. It has been saved to the artifacts";

private static Part blobPart() {
return Part.fromBytes(BLOB_PAYLOAD.getBytes(UTF_8), BLOB_MIME_TYPE);
}

/** The text the runner substitutes for the blob it offloaded to {@code fileName}. */
private static String placeholderFor(String fileName) {
return PLACEHOLDER_FORMAT.formatted(fileName);
}

/**
* A message whose parts list is immutable: {@code Content.Builder.parts(List)} stores the
* caller's list without copying it.
*/
private static Content immutablePartsMessage() {
return Content.builder()
.role("user")
.parts(ImmutableList.of(Part.fromText("hello"), blobPart()))
.build();
}

/** A message whose parts list genai itself collected into an {@code ImmutableList}. */
private static Content partBuilderPartsMessage() {
return Content.builder()
.role("user")
.parts(Part.fromText("hello").toBuilder(), blobPart().toBuilder())
.build();
}

/**
* A message whose parts list accepts {@code set}. Used where the assertion is that the runner
* leaves the caller's message alone: with an immutable list the runner could not have modified it
* either way, so only a mutable one distinguishes copying from rewriting in place.
*/
private static Content mutablePartsMessage() {
return Content.builder()
.role("user")
.parts(new ArrayList<>(ImmutableList.of(Part.fromText("hello"), blobPart())))
.build();
}

/**
* A message carrying two blobs, at part indices 1 and 2. The runner names each artifact after the
* index of the part it came from, so only a message with more than one blob distinguishes that
* from a running counter.
*/
private static Content twoBlobsMessage() {
return Content.builder()
.role("user")
.parts(ImmutableList.of(Part.fromText("hello"), blobPart(), blobPart()))
.build();
}

private static RunConfig saveInputBlobs(boolean enabled) {
return RunConfig.builder().saveInputBlobsAsArtifacts(enabled).build();
}

/**
* Points {@link #runner} at a runner backed by a fresh {@link InMemoryArtifactService}, with a
* fresh {@link #session} on it. What the service stored is read back with {@link #artifactNames}
* and {@link Runner#artifactService()}.
*/
private void useRunnerWithArtifactService() {
this.runner =
Runner.builder()
.app(App.builder().name("test").rootAgent(agent).build())
.artifactService(new InMemoryArtifactService())
.build();
this.session = this.runner.sessionService().createSession("test", "user").blockingGet();
}

/** The names of the artifacts saved for {@link #session}. */
private ImmutableList<String> artifactNames() {
return ImmutableList.copyOf(
runner
.artifactService()
.listArtifactKeys("test", "user", session.id())
.blockingGet()
.filenames());
}

/** The name of the single saved artifact whose file name ends in {@code suffix}. */
private String artifactNameEndingIn(String suffix) {
ImmutableList<String> matches =
artifactNames().stream().filter(name -> name.endsWith(suffix)).collect(toImmutableList());
assertThat(matches).hasSize(1);
return matches.get(0);
}

/** The user message that was actually appended to the session. */
private Content appendedUserMessage() {
Session stored =
runner
.sessionService()
.getSession("test", "user", session.id(), Optional.empty())
.blockingGet();
return stored.events().stream()
.filter(event -> event.author().equals("user"))
.findFirst()
.flatMap(Event::content)
.orElseThrow(() -> new AssertionError("No user message was appended to the session."));
}

/** The parts of the user message that was actually appended to the session. */
private List<Part> appendedUserParts() {
return appendedUserMessage()
.parts()
.orElseThrow(() -> new AssertionError("The appended user message has no parts."));
}

/** Asserts the run reached the model and emitted the agent's reply. */
private static void assertAgentReplied(TestSubscriber<Event> events) {
events.assertComplete();
assertThat(simplifyEvents(events.values())).containsExactly("test agent: from llm");
}

@Test
public void saveInputBlobsAsArtifacts_immutablePartsList_savesArtifactAndCompletes() {
useRunnerWithArtifactService();

var events =
runner.runAsync("user", session.id(), immutablePartsMessage(), saveInputBlobs(true)).test();

assertAgentReplied(events);
assertThat(artifactNames()).hasSize(1);
}

@Test
public void saveInputBlobsAsArtifacts_partBuilderPartsList_savesArtifactAndCompletes() {
useRunnerWithArtifactService();

var events =
runner
.runAsync("user", session.id(), partBuilderPartsMessage(), saveInputBlobs(true))
.test();

assertAgentReplied(events);
assertThat(artifactNames()).hasSize(1);
}

@Test
public void saveInputBlobsAsArtifacts_doesNotModifyCallerMessage() {
useRunnerWithArtifactService();
Content callerMessage = mutablePartsMessage();

var events = runner.runAsync("user", session.id(), callerMessage, saveInputBlobs(true)).test();

assertAgentReplied(events);
assertThat(artifactNames()).hasSize(1);
assertThat(callerMessage.parts().get().get(1).inlineData()).isPresent();
assertThat(callerMessage.parts().get().get(1).text()).isEmpty();
}

@Test
public void saveInputBlobsAsArtifacts_appendedEventReplacesBlobWithPlaceholder() {
useRunnerWithArtifactService();

var events =
runner.runAsync("user", session.id(), immutablePartsMessage(), saveInputBlobs(true)).test();

assertAgentReplied(events);
// The appended message is a copy of the caller's, so the role has to survive the copy.
assertThat(appendedUserMessage().role()).hasValue("user");
List<Part> appended = appendedUserParts();
assertThat(appended).hasSize(2);
assertThat(appended.get(0).text()).hasValue("hello");
assertThat(appended.get(1).inlineData()).isEmpty();
assertThat(appended.get(1).text()).hasValue(placeholderFor(artifactNames().get(0)));
}

@Test
public void saveInputBlobsAsArtifacts_twoBlobs_namesEachArtifactAfterItsPartIndex() {
useRunnerWithArtifactService();

var events =
runner.runAsync("user", session.id(), twoBlobsMessage(), saveInputBlobs(true)).test();

assertAgentReplied(events);
assertThat(artifactNames()).hasSize(2);
List<Part> appended = appendedUserParts();
assertThat(appended).hasSize(3);
assertThat(appended.get(1).text()).hasValue(placeholderFor(artifactNameEndingIn("_1")));
assertThat(appended.get(2).text()).hasValue(placeholderFor(artifactNameEndingIn("_2")));
}

@Test
public void saveInputBlobsAsArtifacts_storesBlobVerbatim() {
useRunnerWithArtifactService();

var events =
runner.runAsync("user", session.id(), immutablePartsMessage(), saveInputBlobs(true)).test();

assertAgentReplied(events);
assertThat(artifactNames()).hasSize(1);
Part stored =
runner
.artifactService()
.loadArtifact("test", "user", session.id(), artifactNames().get(0))
.blockingGet();
assertThat(new String(stored.inlineData().get().data().get(), UTF_8)).isEqualTo(BLOB_PAYLOAD);
assertThat(stored.inlineData().get().mimeType()).hasValue(BLOB_MIME_TYPE);
}

@Test
public void saveInputBlobsAsArtifacts_textOnlyMessage_passesThroughUnchanged() {
useRunnerWithArtifactService();
Content callerMessage = Content.fromParts(Part.fromText("hello"));

var events = runner.runAsync("user", session.id(), callerMessage, saveInputBlobs(true)).test();

assertAgentReplied(events);
assertThat(artifactNames()).isEmpty();
List<Part> appended = appendedUserParts();
assertThat(appended).hasSize(1);
assertThat(appended.get(0).text()).hasValue("hello");
assertThat(callerMessage.parts().get().get(0).text()).hasValue("hello");
}

@Test
public void saveInputBlobsAsArtifacts_disabledWithTextOnlyMessage_passesThroughUnchanged() {
// The default path for every ordinary agent call: no blob, and the option at its default false.
// The runner must not touch the message at all.
useRunnerWithArtifactService();
Content callerMessage = Content.fromParts(Part.fromText("hello"));

var events = runner.runAsync("user", session.id(), callerMessage, saveInputBlobs(false)).test();

assertAgentReplied(events);
assertThat(artifactNames()).isEmpty();
List<Part> appended = appendedUserParts();
assertThat(appended).hasSize(1);
assertThat(appended.get(0).text()).hasValue("hello");
assertThat(callerMessage.parts().get().get(0).text()).hasValue("hello");
}

@Test
public void saveInputBlobsAsArtifacts_disabled_keepsBlobAndSavesNothing() {
useRunnerWithArtifactService();

var events =
runner
.runAsync("user", session.id(), immutablePartsMessage(), saveInputBlobs(false))
.test();

assertAgentReplied(events);
assertThat(artifactNames()).isEmpty();
assertThat(appendedUserParts().get(1).inlineData()).isPresent();
}

@Test
public void saveInputBlobsAsArtifacts_fromPartsConstruction_savesArtifactAndCompletes() {
useRunnerWithArtifactService();
Content fromPartsMessage = Content.fromParts(Part.fromText("hello"), blobPart());

var events =
runner.runAsync("user", session.id(), fromPartsMessage, saveInputBlobs(true)).test();

assertAgentReplied(events);
assertThat(artifactNames()).hasSize(1);
List<Part> appended = appendedUserParts();
assertThat(appended).hasSize(2);
assertThat(appended.get(1).inlineData()).isEmpty();
assertThat(appended.get(1).text()).hasValue(placeholderFor(artifactNames().get(0)));
}

@Test
public void runAsync_partialEvent_streamedButNotPassedToSessionService() {
// The model streams a partial event followed by the final aggregated event in one turn.
Expand Down