Skip to content

Commit 2bd480b

Browse files
google-genai-botcopybara-github
authored andcommitted
feat: Add GcsOffloader for asynchronously uploading content to Google Cloud Storage
This change enables user with possibility to upload content to GCP PiperOrigin-RevId: 919953473
1 parent d608909 commit 2bd480b

11 files changed

Lines changed: 786 additions & 40 deletions

File tree

core/src/main/java/com/google/adk/plugins/agentanalytics/BigQueryAgentAnalyticsPlugin.java

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -253,7 +253,10 @@ private Completable logEvent(
253253
parseFuture =
254254
state
255255
.getParser()
256-
.parse(content)
256+
.parse(
257+
content,
258+
traceIds.traceId(),
259+
traceIds.spanId() != null ? traceIds.spanId() : "no_span")
257260
.thenAccept(
258261
parsedContent -> {
259262
row.put(
Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.adk.plugins.agentanalytics;
18+
19+
import static java.nio.charset.StandardCharsets.UTF_8;
20+
21+
import com.google.auth.Credentials;
22+
import com.google.cloud.storage.BlobId;
23+
import com.google.cloud.storage.BlobInfo;
24+
import com.google.cloud.storage.Storage;
25+
import com.google.cloud.storage.StorageOptions;
26+
import java.util.concurrent.CompletableFuture;
27+
import java.util.concurrent.Executor;
28+
import java.util.concurrent.RejectedExecutionException;
29+
import org.jspecify.annotations.Nullable;
30+
31+
/** Offloads content to GCS. */
32+
class GcsOffloader {
33+
private final Storage storage;
34+
private final String bucketName;
35+
private final Executor executor;
36+
private final boolean isStorageOverride;
37+
38+
GcsOffloader(
39+
String projectId,
40+
String bucketName,
41+
Executor executor,
42+
@Nullable Credentials credentials,
43+
@Nullable Storage storageOverride) {
44+
if (storageOverride != null) {
45+
this.isStorageOverride = true;
46+
this.storage = storageOverride;
47+
} else {
48+
this.isStorageOverride = false;
49+
StorageOptions.Builder builder = StorageOptions.newBuilder().setProjectId(projectId);
50+
if (credentials != null) {
51+
builder.setCredentials(credentials);
52+
}
53+
this.storage = builder.build().getService();
54+
}
55+
this.bucketName = bucketName;
56+
this.executor = executor;
57+
}
58+
59+
/** Async wrapper around blocking GCS upload for binary data. */
60+
CompletableFuture<String> uploadContent(byte[] data, String contentType, String path) {
61+
try {
62+
return CompletableFuture.supplyAsync(
63+
() -> {
64+
BlobId blobId = BlobId.of(bucketName, path);
65+
BlobInfo blobInfo = BlobInfo.newBuilder(blobId).setContentType(contentType).build();
66+
storage.create(blobInfo, data);
67+
return String.format("gs://%s/%s", bucketName, path);
68+
},
69+
executor);
70+
} catch (RejectedExecutionException e) {
71+
return CompletableFuture.failedFuture(e);
72+
}
73+
}
74+
75+
/** Async wrapper around blocking GCS upload for text data. */
76+
CompletableFuture<String> uploadContent(String data, String contentType, String path) {
77+
try {
78+
return CompletableFuture.supplyAsync(() -> data.getBytes(UTF_8), executor)
79+
.thenCompose(bytes -> uploadContent(bytes, contentType, path));
80+
} catch (RejectedExecutionException e) {
81+
return CompletableFuture.failedFuture(e);
82+
}
83+
}
84+
85+
String getBucketName() {
86+
return bucketName;
87+
}
88+
89+
void close() throws Exception {
90+
if (storage != null && !isStorageOverride) {
91+
storage.close();
92+
}
93+
}
94+
}
Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.adk.plugins.agentanalytics;
18+
19+
import com.google.common.collect.ImmutableMap;
20+
21+
/** Utility to map MIME types to file extensions. */
22+
final class MimeTypeMapper {
23+
private static final ImmutableMap<String, String> MIME_TO_EXT =
24+
ImmutableMap.<String, String>builder()
25+
// Images
26+
.put("image/jpeg", ".jpg")
27+
.put("image/png", ".png")
28+
.put("image/gif", ".gif")
29+
.put("image/webp", ".webp")
30+
.put("image/bmp", ".bmp")
31+
.put("image/tiff", ".tiff")
32+
// Audio
33+
.put("audio/mpeg", ".mp3")
34+
.put("audio/ogg", ".ogg")
35+
.put("audio/wav", ".wav")
36+
.put("audio/x-wav", ".wav")
37+
.put("audio/webm", ".webm")
38+
.put("audio/aac", ".aac")
39+
.put("audio/midi", ".mid")
40+
.put("audio/x-m4a", ".m4a")
41+
// Video
42+
.put("video/mp4", ".mp4")
43+
.put("video/mpeg", ".mpeg")
44+
.put("video/ogg", ".ogv")
45+
.put("video/webm", ".webm")
46+
.put("video/avi", ".avi")
47+
.put("video/x-msvideo", ".avi")
48+
.put("video/quicktime", ".mov")
49+
.buildOrThrow();
50+
51+
private MimeTypeMapper() {}
52+
53+
/**
54+
* Returns the file extension (including the dot) for the given MIME type. Returns an empty string
55+
* if the MIME type is unknown.
56+
*/
57+
static String getExtension(String mimeType) {
58+
return MIME_TO_EXT.getOrDefault(mimeType, "");
59+
}
60+
}

0 commit comments

Comments
 (0)