adk-typescript
Version:
TypeScript port of Google's Agent Development Kit (ADK)
181 lines (180 loc) • 6.72 kB
JavaScript
;
Object.defineProperty(exports, "__esModule", { value: true });
exports.GcsArtifactService = void 0;
/**
* An artifact service implementation using Google Cloud Storage (GCS).
*/
const storage_1 = require("@google-cloud/storage");
/**
* An artifact service implementation using Google Cloud Storage (GCS).
*/
class GcsArtifactService {
/**
* Initializes the GcsArtifactService.
*
* @param bucketName The name of the bucket to use
* @param options Optional configuration options for the Google Cloud Storage client
*/
constructor(bucketName, options = {}) {
this.bucketName = bucketName;
this.storage = new storage_1.Storage(options);
this.bucket = this.storage.bucket(this.bucketName);
}
/**
* Checks if the filename has a user namespace.
*
* @param filename The filename to check
* @returns True if the filename has a user namespace (starts with "user:"), false otherwise
*/
fileHasUserNamespace(filename) {
return filename.startsWith('user:');
}
/**
* Constructs the blob name in GCS.
*
* @param appName The name of the application
* @param userId The ID of the user
* @param sessionId The ID of the session
* @param filename The name of the artifact file
* @param version The version of the artifact
* @returns The constructed blob name in GCS
*/
getBlobName(appName, userId, sessionId, filename, version) {
if (this.fileHasUserNamespace(filename)) {
return `${appName}/${userId}/user/${filename}/${version}`;
}
return `${appName}/${userId}/${sessionId}/${filename}/${version}`;
}
/**
* Saves an artifact to the artifact service storage.
*
* @param params The artifact parameters
* @returns The revision ID
*/
async saveArtifact(params) {
const { appName, userId, sessionId, filename, artifact } = params;
if (!artifact) {
throw new Error('Cannot save empty artifact');
}
const versions = await this.listVersions(params);
const version = versions.length > 0 ? Math.max(...versions) + 1 : 0;
const blobName = this.getBlobName(appName, userId, sessionId, filename, version);
const file = this.bucket.file(blobName);
// Check if the artifact has inline data
if (!artifact.inlineData || !artifact.inlineData.data) {
throw new Error('Artifact must contain inline data');
}
// Upload the file content
await file.save(Buffer.from(artifact.inlineData.data, 'base64'), {
contentType: artifact.inlineData.mimeType
});
return version;
}
/**
* Gets an artifact from the artifact service storage.
*
* @param params The artifact parameters
* @returns The artifact or undefined if not found
*/
async loadArtifact(params) {
const { appName, userId, sessionId, filename, version } = params;
let versionToLoad = version;
if (versionToLoad === undefined) {
const versions = await this.listVersions(params);
if (versions.length === 0) {
return undefined;
}
versionToLoad = Math.max(...versions);
}
const blobName = this.getBlobName(appName, userId, sessionId, filename, versionToLoad);
const file = this.bucket.file(blobName);
try {
const [contents] = await file.download();
const [metadata] = await file.getMetadata();
return {
inlineData: {
data: contents.toString('base64'),
mimeType: metadata.contentType || 'application/octet-stream'
}
};
}
catch (error) {
console.error('Error loading artifact:', error);
return undefined;
}
}
/**
* Lists all the artifact filenames within a session.
*
* @param params The artifact parameters
* @returns A list of all artifact filenames within a session
*/
async listArtifactKeys(params) {
const { appName, userId, sessionId } = params;
const filenames = new Set();
// List files in the session-specific prefix
const sessionPrefix = `${appName}/${userId}/${sessionId}/`;
const [sessionFiles] = await this.storage.bucket(this.bucketName).getFiles({
prefix: sessionPrefix
});
for (const file of sessionFiles) {
const parts = file.name.split('/');
if (parts.length >= 4) {
filenames.add(parts[3]);
}
}
// List files in the user namespace prefix
const userNamespacePrefix = `${appName}/${userId}/user/`;
const [userNamespaceFiles] = await this.storage.bucket(this.bucketName).getFiles({
prefix: userNamespacePrefix
});
for (const file of userNamespaceFiles) {
const parts = file.name.split('/');
if (parts.length >= 4) {
filenames.add(parts[3]);
}
}
return Array.from(filenames).sort();
}
/**
* Deletes an artifact.
*
* @param params The artifact parameters
*/
async deleteArtifact(params) {
const { appName, userId, sessionId, filename } = params;
const versions = await this.listVersions(params);
const deletePromises = versions.map(version => {
const blobName = this.getBlobName(appName, userId, sessionId, filename, version);
return this.bucket.file(blobName).delete().catch(error => {
console.error(`Failed to delete version ${version} of ${filename}:`, error);
});
});
await Promise.all(deletePromises);
}
/**
* Lists all versions of an artifact.
*
* @param params The artifact parameters
* @returns A list of all available versions of the artifact
*/
async listVersions(params) {
const { appName, userId, sessionId, filename } = params;
const prefix = this.getBlobName(appName, userId, sessionId, filename, '');
const [files] = await this.bucket.getFiles({
prefix
});
const versions = [];
for (const file of files) {
const parts = file.name.split('/');
if (parts.length >= 5) {
const version = parseInt(parts[4], 10);
if (!isNaN(version)) {
versions.push(version);
}
}
}
return versions;
}
}
exports.GcsArtifactService = GcsArtifactService;