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
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,8 @@

package com.google.adk.plugins.agentanalytics;

import static com.google.adk.plugins.agentanalytics.BigQueryUtils.createAnalyticsViews;
import static com.google.adk.plugins.agentanalytics.BigQueryUtils.maybeUpgradeSchema;
import static com.google.adk.plugins.agentanalytics.JsonFormatter.convertToJsonNode;
import static com.google.adk.plugins.agentanalytics.JsonFormatter.smartTruncate;
import static com.google.adk.plugins.agentanalytics.JsonFormatter.toJavaObject;
Expand Down Expand Up @@ -143,6 +145,8 @@ private static BigQuery createBigQuery(BigQueryLoggerConfig config) throws IOExc
builder.setCredentials(
GoogleCredentials.getApplicationDefault().createScoped(DEFAULT_AUTH_SCOPES));
}
builder = builder.setLocation(config.location());
builder.setProjectId(config.projectId());
return builder.build().getService();
}

Expand Down Expand Up @@ -172,28 +176,37 @@ private void ensureTableExists(BigQuery bigQuery, BigQueryLoggerConfig config) {
tableDefinitionBuilder.setClustering(
Clustering.newBuilder().setFields(config.clusteringFields()).build());
}
TableInfo tableInfo = TableInfo.newBuilder(tableId, tableDefinitionBuilder.build()).build();
TableInfo tableInfo =
TableInfo.newBuilder(tableId, tableDefinitionBuilder.build())
.setLabels(
ImmutableMap.of(
BigQuerySchema.SCHEMA_VERSION_LABEL_KEY, BigQuerySchema.SCHEMA_VERSION))
.build();
bigQuery.create(tableInfo);
} else if (config.autoSchemaUpgrade()) {
// TODO(b/491851868): Implement auto-schema upgrade.
logger.info("BigQuery table already exists and auto-schema upgrade is enabled: " + tableId);
logger.info("Auto-schema upgrade is not implemented yet.");
maybeUpgradeSchema(bigQuery, table);
}
} catch (BigQueryException e) {
if (e.getMessage().contains("invalid_grant")) {
logger.log(
Level.SEVERE,
"Failed to authenticate with BigQuery. Please run 'gcloud auth application-default"
+ " login' to refresh your credentials or provide valid credentials in"
+ " BigQueryLoggerConfig.",
e);
} else {
logger.log(
Level.WARNING, "Failed to check or create/upgrade BigQuery table: " + tableId, e);
}
processBigQueryException(e, "Failed to check or create/upgrade BigQuery table: " + tableId);
} catch (RuntimeException e) {
logger.log(Level.WARNING, "Failed to check or create/upgrade BigQuery table: " + tableId, e);
}

try {
if (config.createViews()) {
var unused = executor.submit(() -> createAnalyticsViews(bigQuery, config));
}
} catch (RuntimeException e) {
logger.log(Level.WARNING, "Failed to create/update BigQuery views for table: " + tableId, e);
}
}

private void processBigQueryException(BigQueryException e, String logMessage) {
if (e.getMessage().contains("invalid_grant")) {
logger.log(Level.SEVERE, "Failed to authenticate with BigQuery.", e);
} else {
logger.log(Level.WARNING, logMessage, e);
}
}

protected BigQueryWriteClient createWriteClient(BigQueryLoggerConfig config) throws IOException {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,9 @@ public abstract class BigQueryLoggerConfig {
// Max length for text content before truncation.
public abstract int maxContentLength();

// BigQuery location.
public abstract String location();

// Project ID for the BigQuery table.
public abstract String projectId();

Expand Down Expand Up @@ -93,9 +96,16 @@ public abstract class BigQueryLoggerConfig {
// Automatically add new columns to existing tables when the plugin
// schema evolves. Only additive changes are made (columns are never
// dropped or altered).
// TODO(b/491852782): Implement auto-schema upgrade.
public abstract boolean autoSchemaUpgrade();

// Automatically create per-event-type BigQuery views that unnest
// JSON columns into typed, queryable columns.
public abstract boolean createViews();

// Prefix for auto-created per-event-type view names.
// Default "v" produces views like ``v_llm_request``.
public abstract String viewPrefix();

@Nullable
public abstract Credentials credentials();

Expand All @@ -105,6 +115,7 @@ public static Builder builder() {
return new AutoValue_BigQueryLoggerConfig.Builder()
.enabled(true)
.maxContentLength(500 * 1024)
.location("us") // Default location.
.datasetId("agent_analytics")
.tableName("events")
.clusteringFields(ImmutableList.of("event_type", "agent", "user_id"))
Expand All @@ -118,8 +129,9 @@ public static Builder builder() {
.customTags(ImmutableMap.of())
.eventAllowlist(ImmutableList.of())
.eventDenylist(ImmutableList.of())
// TODO(b/491851868): Enable auto-schema upgrade once implemented.
.autoSchemaUpgrade(false);
.autoSchemaUpgrade(true)
.createViews(false)
.viewPrefix("v");
}

/** Builder for {@link BigQueryLoggerConfig}. */
Expand All @@ -138,6 +150,9 @@ public abstract static class Builder {
@CanIgnoreReturnValue
public abstract Builder maxContentLength(int maxContentLength);

@CanIgnoreReturnValue
public abstract Builder location(String location);

@CanIgnoreReturnValue
public abstract Builder projectId(String projectId);

Expand Down Expand Up @@ -184,6 +199,12 @@ public abstract Builder contentFormatter(
@CanIgnoreReturnValue
public abstract Builder autoSchemaUpgrade(boolean autoSchemaUpgrade);

@CanIgnoreReturnValue
public abstract Builder createViews(boolean createViews);

@CanIgnoreReturnValue
public abstract Builder viewPrefix(String viewPrefix);

@CanIgnoreReturnValue
public abstract Builder credentials(Credentials credentials);

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,14 @@ public final class BigQuerySchema {

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_LABEL_KEY = "adk_schema_version";

private static final ImmutableMap<StandardSQLTypeName, ImmutableMap<String, String>>
FIELD_TYPE_TO_ARROW_FIELD_METADATA =
ImmutableMap.of(
Expand Down
Loading