Skip to content

feat: Refactor BigQueryAgentAnalyticsPlugin for async in preparation for GCS offloading #1177

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.

Sign up for GitHub

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

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: Refactor BigQueryAgentAnalyticsPlugin for async in preparation for GCS offloading #1177

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
880 changes: 434 additions & 446 deletions core/src/main/java/com/google/adk/plugins/agentanalytics/BigQueryAgentAnalyticsPlugin.java
View file
Open in desktop

Large diffs are not rendered by default.

6 changes: 2 additions & 4 deletions core/src/main/java/com/google/adk/plugins/agentanalytics/BigQueryLoggerConfig.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,6 @@ public abstract class BigQueryLoggerConfig {
public abstract ImmutableList<String> clusteringFields();

// Whether to log multi-modal content.
// TODO(b/491852782): Implement logging of multi-modal content.
public abstract boolean logMultiModalContent();

// Retry configuration for BigQuery writes.
Expand Down Expand Up @@ -96,7 +95,7 @@ public abstract class BigQueryLoggerConfig {
// GCS bucket name to store multi-modal content.
public abstract String gcsBucketName();

// TODO(b/491852782): Implement connection id.
// Optional BigQuery connection ID for ObjectRef columns
public abstract Optional<String> connectionId();

// Toggle for session metadata (e.g. gchat thread-id).
Expand All @@ -118,8 +117,7 @@ public abstract class BigQueryLoggerConfig {
// Default "v" produces views like ``v_llm_request``.
public abstract String viewPrefix();

@Nullable
public abstract Credentials credentials();
public abstract @Nullable Credentials credentials();

public abstract Builder toBuilder();

Expand Down
313 changes: 47 additions & 266 deletions core/src/main/java/com/google/adk/plugins/agentanalytics/JsonFormatter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,29 +16,20 @@

package com.google.adk.plugins.agentanalytics;

import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.google.adk.models.LlmRequest;
import com.google.auto.value.AutoValue;
import com.google.common.collect.ImmutableList;
import com.google.genai.types.Blob;
import com.google.genai.types.Content;
import com.google.genai.types.FileData;
import com.google.genai.types.FunctionCall;
import com.google.genai.types.Part;
import java.util.ArrayList;
import java.util.List;
import com.google.common.base.Utf8;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import org.jspecify.annotations.Nullable;

/** Utility for parsing, formatting and truncating content for BigQuery logging. */
final class JsonFormatter {
private static final ObjectMapper mapper = new ObjectMapper().findAndRegisterModules();
static final ObjectMapper mapper = new ObjectMapper().findAndRegisterModules();
static final String TRUNCATION_SUFFIX = "...[truncated]";

@AutoValue
abstract static class TruncationResult {
Expand All @@ -51,254 +42,6 @@ static TruncationResult create(JsonNode node, boolean isTruncated) {
}
}

@AutoValue
abstract static class ParsedContent {
abstract ImmutableList<JsonNode> parts();

abstract JsonNode content();

abstract boolean isTruncated();

static ParsedContent create(
ImmutableList<JsonNode> parts, JsonNode content, boolean isTruncated) {
return new AutoValue_JsonFormatter_ParsedContent(parts, content, isTruncated);
}
}

@AutoValue
abstract static class ParsedContentObject {
abstract ArrayNode parts();

abstract String summary();

abstract boolean isTruncated();

static ParsedContentObject create(ArrayNode parts, String summary, boolean isTruncated) {
return new AutoValue_JsonFormatter_ParsedContentObject(parts, summary, isTruncated);
}
}

@AutoValue
abstract static class ContentPart {
@JsonProperty("part_index")
abstract int partIndex();

@JsonProperty("mime_type")
abstract @Nullable String mimeType();

@JsonProperty("uri")
abstract @Nullable String uri();

@JsonProperty("text")
abstract @Nullable String text();

@JsonProperty("part_attributes")
abstract String partAttributes();

@JsonProperty("storage_mode")
abstract String storageMode();

@JsonProperty("object_ref")
abstract @Nullable String objectRef();

static Builder builder() {
return new AutoValue_JsonFormatter_ContentPart.Builder();
}

@AutoValue.Builder
abstract static class Builder {
abstract Builder setPartIndex(int value);

abstract Builder setMimeType(@Nullable String value);

abstract Builder setUri(@Nullable String value);

abstract Builder setText(@Nullable String value);

abstract Builder setPartAttributes(String value);

abstract Builder setStorageMode(String value);

abstract Builder setObjectRef(@Nullable String value);

abstract ContentPart build();
}
}

/**
* Parses content into JSON payload and content parts, matching Python implementation.
*
* @param content the content to parse
* @param maxLength the maximum length for text fields
* @return a ParsedContent object
*/
static ParsedContent parse(Object content, int maxLength) {
JsonNode contentNode = mapper.nullNode();
ArrayNode contentParts = mapper.createArrayNode();
boolean isTruncated = false;

if (content instanceof LlmRequest llmRequest) {
ObjectNode jsonPayload = mapper.createObjectNode();
// Handle prompt
ArrayNode messages = mapper.createArrayNode();
List<Content> contents = llmRequest.contents();
for (Content c : contents) {
String role = c.role().orElse("unknown");
ParsedContentObject parsedContentObject = parseContentObject(c, maxLength);
isTruncated = isTruncated || parsedContentObject.isTruncated();
contentParts.addAll(parsedContentObject.parts());

ObjectNode message = mapper.createObjectNode();
message.put("role", role);
message.put("content", parsedContentObject.summary());
messages.add(message);
}
if (!messages.isEmpty()) {
jsonPayload.set("prompt", messages);
}
// Handle system instruction
if (llmRequest.config().isPresent()
&& llmRequest.config().get().systemInstruction().isPresent()) {
Content systemInstruction = llmRequest.config().get().systemInstruction().get();
ParsedContentObject parsedSystemInstruction =
parseContentObject(systemInstruction, maxLength);
isTruncated = isTruncated || parsedSystemInstruction.isTruncated();
contentParts.addAll(parsedSystemInstruction.parts());
jsonPayload.put("system_prompt", parsedSystemInstruction.summary());
}
contentNode = jsonPayload;
} else if (content instanceof Content || content instanceof Part) {
ParsedContentObject parsedContentObject = parseContentObject(content, maxLength);
ObjectNode summaryNode = mapper.createObjectNode();
summaryNode.put("text_summary", parsedContentObject.summary());
return ParsedContent.create(
ImmutableList.copyOf(parsedContentObject.parts()),
summaryNode,
parsedContentObject.isTruncated());
} else if (content instanceof String s) {
TruncationResult result = truncateWithStatus(s, maxLength);
contentNode = result.node();
isTruncated = result.isTruncated();
} else {
TruncationResult result = smartTruncate(content, maxLength);
contentNode = result.node();
isTruncated = result.isTruncated();
}
return ParsedContent.create(ImmutableList.copyOf(contentParts), contentNode, isTruncated);
}

/**
* Parses a Content or Part object into summary text and content parts.
*
* @param content the Content or Part object to parse
* @param maxLength the maximum length of text fields before truncation
* @return a ParsedContentObject containing parts, summary, and truncation flag
*/
private static ParsedContentObject parseContentObject(Object content, int maxLength) {
ArrayNode contentParts = mapper.createArrayNode();
boolean isTruncated = false;
List<String> summaryText = new ArrayList<>();

List<Part> parts;
if (content instanceof Content c) {
parts = c.parts().orElse(ImmutableList.of());
} else if (content instanceof Part p) {
parts = ImmutableList.of(p);
} else {
return ParsedContentObject.create(contentParts, "", false);
}

for (int i = 0; i < parts.size(); i++) {
Part part = parts.get(i);
ContentPart.Builder partBuilder =
ContentPart.builder()
.setPartIndex(i)
.setMimeType("text/plain")
.setUri(null)
.setText(null)
.setPartAttributes("{}")
.setStorageMode("INLINE")
.setObjectRef(null);

// CASE A: It is already a URI (e.g. from user input)
if (part.fileData().isPresent()) {
FileData fileData = part.fileData().get();
partBuilder
.setStorageMode("EXTERNAL_URI")
.setUri(fileData.fileUri().orElse(null))
.setMimeType(fileData.mimeType().orElse(null));
}
// CASE B: It is Binary/Inline Data (Image/Blob)
else if (part.inlineData().isPresent()) {
// TODO: (b/485571635) Implement GCS offloading here.
partBuilder
.setText("[BINARY DATA]")
.setMimeType(part.inlineData().get().mimeType().orElse(""));
}
// CASE C: Text
else if (part.text().isPresent()) {
String text = part.text().get();
// TODO: (b/485571635) Implement GCS offloading if text length exceeds maxLength.
if (text.length() > maxLength) {
text = truncate(text, maxLength);
isTruncated = true;
}
partBuilder.setText(text);
summaryText.add(text);
} else if (part.functionCall().isPresent()) {
FunctionCall fc = part.functionCall().get();
ObjectNode partAttributes = mapper.createObjectNode();
partAttributes.put("function_name", fc.name().orElse("unknown"));
partBuilder
.setMimeType("application/json")
.setText("Function: " + fc.name().orElse("unknown"))
.setPartAttributes(partAttributes.toString());
}
contentParts.add(mapper.valueToTree(partBuilder.build()));
}

String summaryResult = String.join(" | ", summaryText);
if (summaryResult.length() > maxLength) {
summaryResult = truncate(summaryResult, maxLength);
isTruncated = true;
}

return ParsedContentObject.create(contentParts, summaryResult, isTruncated);
}

/** Formats Content parts into an ArrayNode for BigQuery logging. */
static ArrayNode formatContentParts(Optional<Content> content, int maxLength) {
ArrayNode partsArray = mapper.createArrayNode();
if (content.isEmpty() || content.get().parts() == null) {
return partsArray;
}

List<Part> parts = content.get().parts().orElse(ImmutableList.of());

for (int i = 0; i < parts.size(); i++) {
Part part = parts.get(i);
ObjectNode partObj = mapper.createObjectNode();
partObj.put("part_index", i);
partObj.put("storage_mode", "INLINE");

if (part.text().isPresent()) {
partObj.put("mime_type", "text/plain");
partObj.put("text", truncate(part.text().get(), maxLength));
} else if (part.inlineData().isPresent()) {
Blob blob = part.inlineData().get();
partObj.put("mime_type", blob.mimeType().orElse(""));
partObj.put("text", "[BINARY DATA]");
} else if (part.fileData().isPresent()) {
FileData fileData = part.fileData().get();
partObj.put("mime_type", fileData.mimeType().orElse(""));
partObj.put("uri", fileData.fileUri().orElse(""));
partObj.put("storage_mode", "EXTERNAL_URI");
}
partsArray.add(partObj);
}
return partsArray;
}

/** Recursively truncates long strings inside an object and returns a TruncationResult. */
static TruncationResult smartTruncate(Object obj, int maxLength) {
if (obj == null) {
Expand Down Expand Up @@ -328,7 +71,7 @@ private static TruncationResult recursiveSmartTruncate(JsonNode node, int maxLen
boolean isTruncated = false;
if (node.isTextual()) {
String text = node.asText();
if (text.length() > maxLength) {
if (Utf8.encodedLength(text) > maxLength) {
return TruncationResult.create(mapper.valueToTree(truncate(text, maxLength)), true);
}
return TruncationResult.create(node, false);
Expand All @@ -353,21 +96,59 @@ private static TruncationResult recursiveSmartTruncate(JsonNode node, int maxLen
return TruncationResult.create(node, false);
}

private static TruncationResult truncateWithStatus(String s, int maxLength) {
static TruncationResult truncateWithStatus(String s, int maxLength) {
if (s == null) {
return TruncationResult.create(mapper.nullNode(), false);
}
if (s.length() <= maxLength) {
if (Utf8.encodedLength(s) <= maxLength) {
return TruncationResult.create(mapper.valueToTree(s), false);
}
return TruncationResult.create(mapper.valueToTree(truncate(s, maxLength)), true);
}

private static String truncate(String s, int maxLength) {
if (s == null || s.length() <= maxLength) {
static @Nullable String truncate(String s, int budget) {
return truncateAndAddSuffix(s, budget, TRUNCATION_SUFFIX);
}

static @Nullable String truncateAndAddSuffix(String s, int budget, String suffix) {
if (s == null) {
return null;
}
if (Utf8.encodedLength(s) <= budget) {
return s;
}
return s.substring(0, maxLength) + "...[truncated]";
int suffixBytes = Utf8.encodedLength(suffix);
int effectiveBudget = Math.max(0, budget - suffixBytes);
// Fallback in case the budget is too small
if (effectiveBudget == 0) {
return suffix.substring(0, budget);
}

int byteCount = 0;
int charIndex = 0;
for (int i = 0; i < s.length(); ) {
int codePoint = s.codePointAt(i);
int codePointLen = Character.charCount(codePoint);
int codePointBytes;
if (codePoint < 0x80) {
codePointBytes = 1;
} else if (codePoint < 0x800) {
codePointBytes = 2;
} else if (codePoint < 0x10000) {
codePointBytes = 3;
} else {
codePointBytes = 4;
}

if (byteCount + codePointBytes > effectiveBudget) {
break;
}
byteCount += codePointBytes;
charIndex += codePointLen;
i += codePointLen;
}

return s.substring(0, charIndex) + suffix;
}

/** Converts a JsonNode to a standard Java object (Map, List, etc.). */
Expand Down
Loading