diff --git a/core/src/main/java/com/google/adk/plugins/agentanalytics/BigQueryAgentAnalyticsPlugin.java b/core/src/main/java/com/google/adk/plugins/agentanalytics/BigQueryAgentAnalyticsPlugin.java index a71785c49..6982c7df0 100644 --- a/core/src/main/java/com/google/adk/plugins/agentanalytics/BigQueryAgentAnalyticsPlugin.java +++ b/core/src/main/java/com/google/adk/plugins/agentanalytics/BigQueryAgentAnalyticsPlugin.java @@ -327,6 +327,7 @@ private Completable logEvent( // Log common fields Map row = new HashMap<>(); row.put("timestamp", Instant.now()); + row.put("event_id", UUID.randomUUID().toString().replace("-", "")); row.put("event_type", eventType); row.put("agent", resolveAgentName(invocationContext, eventData)); row.put("session_id", invocationContext.session().id()); diff --git a/core/src/main/java/com/google/adk/plugins/agentanalytics/BigQuerySchema.java b/core/src/main/java/com/google/adk/plugins/agentanalytics/BigQuerySchema.java index 9a7e76f88..bf50a26f1 100644 --- a/core/src/main/java/com/google/adk/plugins/agentanalytics/BigQuerySchema.java +++ b/core/src/main/java/com/google/adk/plugins/agentanalytics/BigQuerySchema.java @@ -49,7 +49,7 @@ private BigQuerySchema() {} * The version of the BigQuery schema. Each time the schema is changed(new fields are added), this * should be incremented. */ - static final String SCHEMA_VERSION = "1"; + static final String SCHEMA_VERSION = "2"; static final String SCHEMA_VERSION_LABEL_KEY = "adk_schema_version"; @@ -75,6 +75,12 @@ public static Schema getEventsSchema() { .setMode(Field.Mode.REQUIRED) .setDescription("The UTC timestamp when the event occurred.") .build(), + Field.newBuilder("event_id", StandardSQLTypeName.STRING) + .setMode(Field.Mode.NULLABLE) + .setDescription( + "A unique identifier assigned before enqueue. Storage Write API retries preserve" + + " this value so duplicate rows can be identified reliably.") + .build(), Field.newBuilder("event_type", StandardSQLTypeName.STRING) .setMode(Field.Mode.NULLABLE) .setDescription("The category of the event.") diff --git a/core/src/main/java/com/google/adk/plugins/agentanalytics/BigQueryUtils.java b/core/src/main/java/com/google/adk/plugins/agentanalytics/BigQueryUtils.java index 58aa66bcc..3d79a888c 100644 --- a/core/src/main/java/com/google/adk/plugins/agentanalytics/BigQueryUtils.java +++ b/core/src/main/java/com/google/adk/plugins/agentanalytics/BigQueryUtils.java @@ -53,9 +53,10 @@ final class BigQueryUtils { static final String A2A_TASK_ID_KEY = "a2a:task_id"; static final String A2A_CONTEXT_ID_KEY = "a2a:context_id"; - private static final ImmutableList VIEW_COMMON_COLUMNS = + static final ImmutableList VIEW_COMMON_COLUMNS = ImmutableList.of( "timestamp", + "event_id", "event_type", "agent", "session_id", diff --git a/core/src/test/java/com/google/adk/plugins/agentanalytics/BigQueryAgentAnalyticsPluginTest.java b/core/src/test/java/com/google/adk/plugins/agentanalytics/BigQueryAgentAnalyticsPluginTest.java index 9506a7298..90ea29866 100644 --- a/core/src/test/java/com/google/adk/plugins/agentanalytics/BigQueryAgentAnalyticsPluginTest.java +++ b/core/src/test/java/com/google/adk/plugins/agentanalytics/BigQueryAgentAnalyticsPluginTest.java @@ -21,6 +21,7 @@ import static java.util.concurrent.TimeUnit.MILLISECONDS; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotEquals; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertNull; import static org.junit.Assert.assertTrue; @@ -86,8 +87,10 @@ import io.opentelemetry.sdk.testing.junit4.OpenTelemetryRule; import io.reactivex.rxjava3.core.Flowable; import java.time.Duration; +import java.util.ArrayList; import java.util.HashMap; import java.util.List; +import java.util.Locale; import java.util.Map; import java.util.Objects; import java.util.Optional; @@ -476,6 +479,11 @@ public void arrowSchema_handlesFieldNullability() { assertNotNull(timestampField); assertFalse(timestampField.isNullable()); + // event_id is NULLABLE in BigQuerySchema.getEventsSchema() + Field eventIdField = schema.findField("event_id"); + assertNotNull(eventIdField); + assertTrue(eventIdField.isNullable()); + // event_type is NULLABLE in BigQuerySchema.getEventsSchema() Field eventTypeField = schema.findField("event_type"); assertNotNull(eventTypeField); @@ -505,6 +513,11 @@ public void logEvent_populatesCommonFields() throws Exception { "USER_MESSAGE_RECEIVED")) { failureMessage[0] = "Wrong event_type: " + root.getVector("event_type").getObject(0); + } else if (!root.getVector("event_id") + .getObject(0) + .toString() + .matches("^[0-9a-f]{32}$")) { + failureMessage[0] = "Wrong event_id: " + root.getVector("event_id").getObject(0); } else if (!root.getVector("agent").getObject(0).toString().equals("agent_name")) { failureMessage[0] = "Wrong agent: " + root.getVector("agent").getObject(0); } else if (!root.getVector("session_id") @@ -560,6 +573,51 @@ public void logEvent_populatesCommonFields() throws Exception { assertTrue(failureMessage[0], checksPassed[0]); } + @Test + public void schemaAndViews_exposeEventId() { + com.google.cloud.bigquery.Schema schema = BigQuerySchema.getEventsSchema(); + com.google.cloud.bigquery.Field eventIdField = schema.getFields().get("event_id"); + assertNotNull(eventIdField); + assertEquals(StandardSQLTypeName.STRING, eventIdField.getType().getStandardType()); + assertEquals(Mode.NULLABLE, eventIdField.getMode()); + assertTrue(BigQueryUtils.VIEW_COMMON_COLUMNS.contains("event_id")); + } + + @Test + public void eachEmittedRow_hasDistinctHexEventId() throws Exception { + List eventIds = new ArrayList<>(); + when(mockWriter.append(any(ArrowRecordBatch.class))) + .thenAnswer( + invocation -> { + ArrowRecordBatch recordedBatch = invocation.getArgument(0); + Schema schema = BigQuerySchema.getArrowSchema(); + try (VectorSchemaRoot root = + VectorSchemaRoot.create( + schema, state.getBatchProcessor("invocation_id").allocator)) { + VectorLoader loader = new VectorLoader(root); + loader.load(recordedBatch); + for (int i = 0; i < root.getRowCount(); i++) { + eventIds.add(root.getVector("event_id").getObject(i).toString()); + } + } + return ApiFutures.immediateFuture(AppendRowsResponse.getDefaultInstance()); + }); + + Content content1 = Content.fromParts(Part.fromText("hello")); + Content content2 = Content.fromParts(Part.fromText("world")); + plugin.onUserMessageCallback(mockInvocationContext, content1).blockingSubscribe(); + plugin.onUserMessageCallback(mockInvocationContext, content2).blockingSubscribe(); + state.getBatchProcessor("invocation_id").flush(); + + assertEquals(2, eventIds.size()); + assertNotEquals(eventIds.get(0), eventIds.get(1)); + for (String eventId : eventIds) { + assertEquals(32, eventId.length()); + assertEquals(eventId.toLowerCase(Locale.ROOT), eventId); + assertTrue(eventId.matches("^[0-9a-f]{32}$")); + } + } + @Test public void logEvent_populatesTraceDetails() throws Exception { String traceId = "4bf92f3577b34da6a3ce929d0e0e4736"; @@ -1165,7 +1223,7 @@ public void beforeToolCallback_concurrentTool_stampsEnclosingParentNotSibling() Map row; while ((row = state.getBatchProcessor("invocation_id").queue.poll()) != null) { if (Objects.equals(row.get("event_type"), "TOOL_STARTING") - && "tool_b".equals(((ObjectNode) row.get("content")).get("tool").asText())) { + && Objects.equals(((ObjectNode) row.get("content")).get("tool").asText(), "tool_b")) { startingB = row; } }