diff --git a/core/src/main/java/com/google/adk/plugins/agentanalytics/BigQueryLoggerConfig.java b/core/src/main/java/com/google/adk/plugins/agentanalytics/BigQueryLoggerConfig.java index 8c9219406..a56183357 100644 --- a/core/src/main/java/com/google/adk/plugins/agentanalytics/BigQueryLoggerConfig.java +++ b/core/src/main/java/com/google/adk/plugins/agentanalytics/BigQueryLoggerConfig.java @@ -58,7 +58,12 @@ public abstract class BigQueryLoggerConfig { // Fields to cluster the table by. public abstract ImmutableList clusteringFields(); - // Whether to log multi-modal content. + /** + * Whether to log detailed content parts in the {@code content_parts} column. When {@code false}, + * {@code content_parts} is written as {@code []} and GCS offloading ({@code gcsBucketName}) is + * disabled so large text remains inline (truncated at {@code maxContentLength}) and binary parts + * are recorded as placeholders. + */ public abstract boolean logMultiModalContent(); // Retry configuration for BigQuery writes. @@ -92,7 +97,12 @@ public abstract class BigQueryLoggerConfig { */ public abstract @Nullable BiFunction contentFormatter(); - // GCS bucket name to store multi-modal content. + /** + * GCS bucket name for offloading large text and binary content. Offloading only runs when {@code + * logMultiModalContent} is {@code true}; otherwise {@code GcsOffloader} is not created, a warning + * is logged, large text stays inline (truncated at {@code maxContentLength}), and binary parts + * are recorded as placeholders. + */ public abstract String gcsBucketName(); // Optional BigQuery connection ID for ObjectRef columns diff --git a/core/src/main/java/com/google/adk/plugins/agentanalytics/Parser.java b/core/src/main/java/com/google/adk/plugins/agentanalytics/Parser.java index c489273e1..ddc446329 100644 --- a/core/src/main/java/com/google/adk/plugins/agentanalytics/Parser.java +++ b/core/src/main/java/com/google/adk/plugins/agentanalytics/Parser.java @@ -63,17 +63,15 @@ final class Parser { private final @Nullable GcsOffloader offloader; private final int maxLength; private final @Nullable String connectionId; - private final boolean logMultiModalContent; Parser( @Nullable GcsOffloader offloader, int maxLength, @Nullable String connectionId, boolean logMultiModalContent) { - this.offloader = offloader; + this.offloader = logMultiModalContent ? offloader : null; this.maxLength = maxLength; this.connectionId = connectionId; - this.logMultiModalContent = logMultiModalContent; } @AutoValue @@ -352,8 +350,10 @@ private CompletableFuture processPart( // CASE B: It is Binary/Inline Data (Image/Blob) if (part.inlineData().isPresent()) { Blob blob = part.inlineData().get(); - String mimeType = blob.mimeType().orElse("application/octet-stream"); - if (logMultiModalContent && offloader != null) { + String mimeType = + blob.mimeType().filter(m -> !m.isEmpty()).orElse("application/octet-stream"); + partBuilder.setMimeType(mimeType); + if (offloader != null) { String ext = MimeTypeMapper.getExtension(mimeType); if (ext.isEmpty()) { ext = DEFAULT_EXTENSION; @@ -378,7 +378,6 @@ private CompletableFuture processPart( partBuilder .setStorageMode("GCS_REFERENCE") .setUri(uri) - .setMimeType(mimeType) .setText(MEDIA_OFFLOADED_MESSAGE) .setObjectRef( mapper.valueToTree(ObjectRef.create(uri, null, connectionId, details))); @@ -386,7 +385,7 @@ private CompletableFuture processPart( return TruncationResult.create(mapper.valueToTree(partBuilder.build()), false); }); } else { - partBuilder.setText(BINARY_DATA_MESSAGE).setMimeType(mimeType); + partBuilder.setText(BINARY_DATA_MESSAGE); return CompletableFuture.completedFuture( TruncationResult.create(mapper.valueToTree(partBuilder.build()), false)); } diff --git a/core/src/main/java/com/google/adk/plugins/agentanalytics/PluginState.java b/core/src/main/java/com/google/adk/plugins/agentanalytics/PluginState.java index 1f712047d..2318dd6f0 100644 --- a/core/src/main/java/com/google/adk/plugins/agentanalytics/PluginState.java +++ b/core/src/main/java/com/google/adk/plugins/agentanalytics/PluginState.java @@ -571,6 +571,16 @@ void appendRow(InvocationLifecycle lifecycle, String invocationId, Map builtOffloaders = new ArrayList<>(); + PluginState gcsState = + new PluginState(gcsConfig) { + @Override + protected BigQueryWriteClient createWriteClient(BigQueryLoggerConfig config) { + return mockWriteClient; + } + + @Override + protected StreamWriter createWriter() { + return mockWriter; + } + + @Override + protected GcsOffloader getGcsOffloader(BigQueryLoggerConfig config) { + // Runs the real factory; a built offloader is swapped for a recording fake so a + // regression fails the assertions below without reaching GCS. + GcsOffloader offloader = super.getGcsOffloader(config); + if (offloader == null) { + return null; + } + builtOffloaders.add(offloader); + return recordingOffloader; + } + }; + BigQueryAgentAnalyticsPlugin gcsPlugin = + new BigQueryAgentAnalyticsPlugin(gcsConfig, mockBigQuery, gcsState); + String largeText = "x".repeat(40 * 1024); + + gcsPlugin + .onUserMessageCallback(mockInvocationContext, Content.fromParts(Part.fromText(largeText))) + .test() + .await() + .assertResult(); + CompletableFuture.allOf( + gcsState + .getPendingTasksForInvocation("invocation_id") + .toArray(new CompletableFuture[0])) + .join(); + Map row = gcsState.getBatchProcessor("invocation_id").queue.poll(); + + assertTrue("No GcsOffloader may be built", builtOffloaders.isEmpty()); + verify(recordingOffloader, never()).uploadContent(anyString(), anyString(), anyString()); + assertNotNull(row); + assertEquals("USER_MESSAGE_RECEIVED", row.get("event_type")); + assertEquals(ImmutableList.of(), row.get("content_parts")); + assertEquals(false, row.get("is_truncated")); + String textSummary = ((JsonNode) row.get("content")).get("text_summary").asText(); + assertEquals(40 * 1024, textSummary.length()); + } + private static class FakeAgent extends BaseAgent { FakeAgent(String name) { super(name, "description", null, null, null); diff --git a/core/src/test/java/com/google/adk/plugins/agentanalytics/ParserTest.java b/core/src/test/java/com/google/adk/plugins/agentanalytics/ParserTest.java index 385e81082..d642e408a 100644 --- a/core/src/test/java/com/google/adk/plugins/agentanalytics/ParserTest.java +++ b/core/src/test/java/com/google/adk/plugins/agentanalytics/ParserTest.java @@ -17,8 +17,18 @@ package com.google.adk.plugins.agentanalytics; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; +import static org.mockito.Mockito.when; +import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.node.ArrayNode; import com.fasterxml.jackson.databind.node.ObjectNode; import com.google.genai.types.Blob; @@ -109,4 +119,147 @@ public void parse_multipartContent_coversLine310() throws Exception { assertTrue(result.isTruncated()); } + + @Test + public void parse_inlineDataUploadFailure_preservesMimeTypeAndFallback() throws Exception { + GcsOffloader failingOffloader = mock(GcsOffloader.class); + when(failingOffloader.uploadContent(any(byte[].class), anyString(), anyString())) + .thenReturn(CompletableFuture.failedFuture(new RuntimeException("GCS boom"))); + Parser offloadingParser = new Parser(failingOffloader, 100, "connectionId", true); + + JsonNode png = parseOnlyPart(offloadingParser, inlinePart("image/png")); + JsonNode pdf = parseOnlyPart(offloadingParser, inlinePart("application/pdf")); + JsonNode empty = parseOnlyPart(offloadingParser, inlinePart("")); + JsonNode absent = parseOnlyPart(offloadingParser, inlinePartWithoutMimeType()); + + verify(failingOffloader).uploadContent(any(byte[].class), eq("image/png"), anyString()); + verify(failingOffloader).uploadContent(any(byte[].class), eq("application/pdf"), anyString()); + verify(failingOffloader, times(2)) + .uploadContent(any(byte[].class), eq("application/octet-stream"), anyString()); + assertInlinePart(png, "image/png", "[UPLOAD FAILED]"); + assertInlinePart(pdf, "application/pdf", "[UPLOAD FAILED]"); + assertInlinePart(empty, "application/octet-stream", "[UPLOAD FAILED]"); + assertInlinePart(absent, "application/octet-stream", "[UPLOAD FAILED]"); + } + + @Test + public void parse_inlineDataUploadSuccess_preservesMimeTypeAndSetsGcsReference() + throws Exception { + GcsOffloader succeedingOffloader = mock(GcsOffloader.class); + when(succeedingOffloader.uploadContent(any(byte[].class), anyString(), anyString())) + .thenReturn(CompletableFuture.completedFuture("gs://test-bucket/blob.bin")); + Parser offloadingParser = new Parser(succeedingOffloader, 100, "connectionId", true); + + JsonNode png = parseOnlyPart(offloadingParser, inlinePart("image/png")); + JsonNode empty = parseOnlyPart(offloadingParser, inlinePart("")); + JsonNode absent = parseOnlyPart(offloadingParser, inlinePartWithoutMimeType()); + + verify(succeedingOffloader).uploadContent(any(byte[].class), eq("image/png"), anyString()); + verify(succeedingOffloader, times(2)) + .uploadContent(any(byte[].class), eq("application/octet-stream"), anyString()); + assertOffloadedBinaryPart(png, "image/png", "gs://test-bucket/blob.bin"); + assertOffloadedBinaryPart(empty, "application/octet-stream", "gs://test-bucket/blob.bin"); + assertOffloadedBinaryPart(absent, "application/octet-stream", "gs://test-bucket/blob.bin"); + } + + @Test + public void parse_inlineDataWithoutOffload_preservesMimeTypeAndEmptyFallback() throws Exception { + GcsOffloader unusedOffloader = mock(GcsOffloader.class); + Parser noOffloaderParser = new Parser(null, 100, "connectionId", true); + Parser multiModalOffParser = new Parser(unusedOffloader, 100, "connectionId", false); + + JsonNode wavNoOffloader = parseOnlyPart(noOffloaderParser, inlinePart("audio/wav")); + JsonNode emptyNoOffloader = parseOnlyPart(noOffloaderParser, inlinePart("")); + JsonNode absentNoOffloader = parseOnlyPart(noOffloaderParser, inlinePartWithoutMimeType()); + JsonNode wavMultiModalOff = parseOnlyPart(multiModalOffParser, inlinePart("audio/wav")); + JsonNode emptyMultiModalOff = parseOnlyPart(multiModalOffParser, inlinePart("")); + JsonNode absentMultiModalOff = parseOnlyPart(multiModalOffParser, inlinePartWithoutMimeType()); + + verifyNoInteractions(unusedOffloader); + assertInlinePart(wavNoOffloader, "audio/wav", "[BINARY DATA]"); + assertInlinePart(emptyNoOffloader, "application/octet-stream", "[BINARY DATA]"); + assertInlinePart(absentNoOffloader, "application/octet-stream", "[BINARY DATA]"); + assertInlinePart(wavMultiModalOff, "audio/wav", "[BINARY DATA]"); + assertInlinePart(emptyMultiModalOff, "application/octet-stream", "[BINARY DATA]"); + assertInlinePart(absentMultiModalOff, "application/octet-stream", "[BINARY DATA]"); + } + + @Test + public void parse_largeTextWithLogMultiModalFalse_skipsOffloadAndKeepsFullText() + throws Exception { + String largeText = "a".repeat(40 * 1024); + GcsOffloader unusedOffloader = succeedingTextOffloader(); + + Parser.ParsedContent kept = + new Parser(unusedOffloader, 100 * 1024, "connectionId", false) + .parse(Content.fromParts(Part.fromText(largeText)), "traceId", "spanId") + .get(); + + verifyNoInteractions(unusedOffloader); + assertFalse(kept.isTruncated()); + String keptSummary = kept.content().get("text_summary").asText(); + assertEquals(largeText.length(), keptSummary.length()); + assertTrue("text_summary must keep the full text", keptSummary.equals(largeText)); + assertEquals("INLINE", kept.parts().get(0).get("storage_mode").asText()); + } + + @Test + public void parse_largeTextWithLogMultiModalTrue_offloadsToGcsAndTruncatesInlinePreview() + throws Exception { + String largeText = "a".repeat(40 * 1024); + GcsOffloader usedOffloader = succeedingTextOffloader(); + + Parser.ParsedContent offloaded = + new Parser(usedOffloader, 100 * 1024, "connectionId", true) + .parse(Content.fromParts(Part.fromText(largeText)), "traceId", "spanId") + .get(); + + verify(usedOffloader).uploadContent(eq(largeText), eq("text/plain"), anyString()); + assertTrue(offloaded.isTruncated()); + assertEquals("GCS_REFERENCE", offloaded.parts().get(0).get("storage_mode").asText()); + assertTrue(offloaded.parts().get(0).get("text").asText().endsWith("... [OFFLOADED]")); + } + + private static Part inlinePart(String mimeType) { + return Part.builder() + .inlineData(Blob.builder().mimeType(mimeType).data(new byte[] {1, 2, 3}).build()) + .build(); + } + + private static Part inlinePartWithoutMimeType() { + return Part.builder().inlineData(Blob.builder().data(new byte[] {1, 2, 3}).build()).build(); + } + + private static JsonNode parseOnlyPart(Parser parser, Part part) throws Exception { + Parser.ParsedContent result = parser.parse(part, "traceId", "spanId").get(); + assertEquals(1, result.parts().size()); + return result.parts().get(0); + } + + private static void assertInlinePart(JsonNode part, String mimeType, String text) { + assertEquals(mimeType, part.get("mime_type").asText()); + assertEquals(text, part.get("text").asText()); + assertEquals("INLINE", part.get("storage_mode").asText()); + assertTrue(part.get("uri").isNull()); + assertTrue(part.get("object_ref").isNull()); + } + + private static void assertOffloadedBinaryPart(JsonNode part, String mimeType, String uri) { + assertEquals(mimeType, part.get("mime_type").asText()); + assertEquals("[MEDIA OFFLOADED]", part.get("text").asText()); + assertEquals("GCS_REFERENCE", part.get("storage_mode").asText()); + assertEquals(uri, part.get("uri").asText()); + JsonNode objectRef = part.get("object_ref"); + assertEquals(uri, objectRef.get("uri").asText()); + assertEquals("connectionId", objectRef.get("authorizer").asText()); + assertEquals( + mimeType, objectRef.get("details").get("gcs_metadata").get("content_type").asText()); + } + + private static GcsOffloader succeedingTextOffloader() { + GcsOffloader offloader = mock(GcsOffloader.class); + when(offloader.uploadContent(anyString(), anyString(), anyString())) + .thenReturn(CompletableFuture.completedFuture("gs://test-bucket/large.txt")); + return offloader; + } } diff --git a/core/src/test/java/com/google/adk/plugins/agentanalytics/PluginStateTest.java b/core/src/test/java/com/google/adk/plugins/agentanalytics/PluginStateTest.java index 59f6087ed..3f70b4340 100644 --- a/core/src/test/java/com/google/adk/plugins/agentanalytics/PluginStateTest.java +++ b/core/src/test/java/com/google/adk/plugins/agentanalytics/PluginStateTest.java @@ -16,6 +16,7 @@ package com.google.adk.plugins.agentanalytics; +import static com.google.common.collect.ImmutableList.toImmutableList; import static java.util.concurrent.TimeUnit.SECONDS; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; @@ -34,6 +35,7 @@ import com.google.cloud.bigquery.storage.v1.BigQueryWriteClient; import com.google.cloud.bigquery.storage.v1.StreamWriter; import com.google.common.cache.Cache; +import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.util.concurrent.Uninterruptibles; import java.io.IOException; @@ -121,6 +123,31 @@ public void getGcsOffloader_emptyBucketName_returnsNull() { assertNull(pluginState.getGcsOffloader(config)); } + @Test + public void getGcsOffloader_logMultiModalContentFalse_returnsNullAndLogsWarning() { + BigQueryLoggerConfig gcsConfig = + config.toBuilder().gcsBucketName("test-bucket").logMultiModalContent(false).build(); + + GcsOffloader offloader = pluginState.getGcsOffloader(gcsConfig); + + assertNull(offloader); + ArgumentCaptor captor = ArgumentCaptor.forClass(LogRecord.class); + verify(mockHandler, atLeastOnce()).publish(captor.capture()); + ImmutableList offloadWarnings = + captor.getAllValues().stream() + .filter(record -> record.getLevel().equals(Level.WARNING)) + .map(LogRecord::getMessage) + .filter(message -> message.startsWith("GCS offload disabled")) + .collect(toImmutableList()); + assertEquals( + ImmutableList.of( + "GCS offload disabled: logMultiModalContent is false, so content_parts (which holds the" + + " offloaded object reference) is not persisted; large text is kept inline" + + " (truncated) and binary parts are recorded as placeholders instead of being" + + " uploaded to test-bucket."), + offloadWarnings); + } + @Test public void addPendingTask_removedTaskOnCompletion() { String invocationId = "testInvocation";