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 @@ -17,6 +17,7 @@ import {
convertCodeExecutionParts,
extractCodeAndTruncateContent,
File,
FileContentEncoding,
} from '../../code_executors/code_execution_utils.js';
import {CodeExecutorContext} from '../../code_executors/code_executor_context.js';
import {createEvent, Event} from '../../events/event.js';
Expand Down Expand Up @@ -427,6 +428,7 @@ function extractAndReplaceInlineFiles(
const file: File = {
name: fileName,
content: base64Decode(part.inlineData.data!),
contentEncoding: FileContentEncoding.UTF8,
mimeType,
};

Expand Down
6 changes: 4 additions & 2 deletions core/src/code_executors/agent_engine_sandbox_code_executor.ts
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
import {Client} from '@google-cloud/vertexai';
import {Language} from '@google-cloud/vertexai/build/src/genai/types.js';
import {experimental} from '../utils/experimental.js';
import {guessMimeType} from '../utils/file_utils.js';
import {decodeFileContent, guessMimeType} from '../utils/file_utils.js';

interface LocalChunk {
data?: string;
Expand All @@ -30,6 +30,7 @@ import {
CodeExecutionLanguage,
CodeExecutionResult,
File,
FileContentEncoding,
} from './code_execution_utils.js';

const DEFAULT_MAX_ATTEMPTS = 180;
Expand Down Expand Up @@ -166,7 +167,7 @@ export class AgentEngineSandboxCodeExecutor extends BaseCodeExecutor {
for (const file of codeExecutionInput.inputFiles) {
inputs.push({
mimeType: file.mimeType,
data: file.content, // Assumed to be already base64 encoded based on CodeExecutionInput definition
data: decodeFileContent(file).toString('base64'),
metadata: {
attributes: {
file_name: Buffer.from(file.name).toString('base64'),
Expand Down Expand Up @@ -226,6 +227,7 @@ export class AgentEngineSandboxCodeExecutor extends BaseCodeExecutor {
outputFiles.push({
name: name,
content: output.data || '',
contentEncoding: FileContentEncoding.BASE64,
mimeType: mimeType,
});
}
Expand Down
5 changes: 3 additions & 2 deletions core/src/code_executors/code_execution_utils.ts
Original file line number Diff line number Diff line change
Expand Up @@ -23,12 +23,13 @@ export interface File {
name: string;

/**
* The encoded bytes of the file content.
* The file content, encoded as described by `contentEncoding`.
* */
content: string;

/**
* The encoding of the file content.
* How `content` is encoded. Defaults to `FileContentEncoding.BASE64` when
* absent, so binary content survives a `File` that omits the field.
*/
contentEncoding?: FileContentEncoding;

Expand Down
29 changes: 23 additions & 6 deletions core/src/utils/file_utils.ts
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,26 @@

import * as fs from 'node:fs/promises';
import * as path from 'node:path';
import {File} from '../code_executors/code_execution_utils.js';
import {
File,
FileContentEncoding,
} from '../code_executors/code_execution_utils.js';

/**
* Decodes a file's content into its raw bytes.
*
* `contentEncoding` is optional, and an absent value means
* `FileContentEncoding.BASE64`: `File` is the only representation of a binary
* payload in the code-executor API, and arbitrary bytes do not survive being
* read back as text. Every producer in this package sets the field
* explicitly, so the default applies only to externally built files.
*/
export function decodeFileContent(file: File): Buffer {
return Buffer.from(
file.content,
file.contentEncoding ?? FileContentEncoding.BASE64,
);
}

/**
* Reports whether resolvedPath is resolvedBaseDir itself, or a path nested
Expand All @@ -27,7 +46,8 @@ function isInsideDir(resolvedPath: string, resolvedBaseDir: string): boolean {

/**
* Creates files with the given paths in the current working directory.
* @param files The files to materialize.
* @param files The files to materialize. Each file is written as the bytes
* {@link decodeFileContent} decodes its content into.
*/
export async function materializeFiles(
files: File[],
Expand Down Expand Up @@ -75,10 +95,7 @@ export async function materializeFiles(
}

await fs.mkdir(path.dirname(finalPath), {recursive: true});
await fs.writeFile(
finalPath,
Buffer.from(file.content, file.contentEncoding),
);
await fs.writeFile(finalPath, decodeFileContent(file));

createdFiles.push({
...file,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import {Client} from '@google-cloud/vertexai';
import {
AgentEngineSandboxCodeExecutor,
CodeExecutionLanguage,
FileContentEncoding,
InvocationContext,
} from '@google/adk';
import {beforeEach, describe, expect, it, vi} from 'vitest';
Expand Down Expand Up @@ -285,6 +286,98 @@ describe('AgentEngineSandboxCodeExecutor', () => {
expect(result.outputFiles[0].mimeType).toBe('image/png');
});

it('marks output files as base64', async () => {
mockClient.agentEnginesInternal.sandboxes.executeCodeInternal.mockResolvedValue(
{
outputs: [
{
mimeType: 'image/png',
data: 'base64data',
metadata: {
attributes: {
file_name: Buffer.from('plot.png').toString('base64'),
},
},
},
],
},
);

const result = await executor.executeCode({
invocationContext,
codeExecutionInput: {
code: 'print("hello")',
language: CodeExecutionLanguage.PYTHON,
inputFiles: [],
},
});

expect(result.outputFiles[0].contentEncoding).toBe(
FileContentEncoding.BASE64,
);
});

it('base64-encodes an input file that declares utf-8', async () => {
await executor.executeCode({
invocationContext,
codeExecutionInput: {
code: 'print("hello")',
language: CodeExecutionLanguage.PYTHON,
inputFiles: [
{
name: 'a.md',
content: '# hi',
contentEncoding: FileContentEncoding.UTF8,
mimeType: 'text/markdown',
},
],
},
});

expect(
mockClient.agentEnginesInternal.sandboxes.executeCodeInternal,
).toHaveBeenCalledWith(
expect.objectContaining({
inputs: expect.arrayContaining([
expect.objectContaining({
mimeType: 'text/markdown',
data: Buffer.from('# hi').toString('base64'),
}),
]),
}),
);
});

it('passes a base64 input file through unchanged', async () => {
const encoded = Buffer.from('a,b,c').toString('base64');

await executor.executeCode({
invocationContext,
codeExecutionInput: {
code: 'print("hello")',
language: CodeExecutionLanguage.PYTHON,
inputFiles: [
{
name: 'data.csv',
content: encoded,
contentEncoding: FileContentEncoding.BASE64,
mimeType: 'text/csv',
},
],
},
});

expect(
mockClient.agentEnginesInternal.sandboxes.executeCodeInternal,
).toHaveBeenCalledWith(
expect.objectContaining({
inputs: expect.arrayContaining([
expect.objectContaining({mimeType: 'text/csv', data: encoded}),
]),
}),
);
});

it('guesses mime type if missing in output', async () => {
mockClient.agentEnginesInternal.sandboxes.executeCodeInternal.mockResolvedValue(
{
Expand Down
53 changes: 53 additions & 0 deletions core/test/utils/file_utils_test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -181,5 +181,58 @@ describe('file_utils', () => {
);
expect(content3).toBe('third');
});

// The PNG signature is not valid utf-8, so a utf-8 write cannot reproduce
// it. Assert on the Buffer: a string comparison passes on mojibake.
const pngSignature = Buffer.from([
0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a,
]);

it('should write a file with no contentEncoding as decoded base64 bytes', async () => {
const files = [
{
name: 'plot.png',
content: pngSignature.toString('base64'),
mimeType: 'image/png',
},
];

await materializeFiles(files, tempDir);

const written = await fs.readFile(path.join(tempDir, 'plot.png'));
expect(written).toEqual(pngSignature);
});

it('should write a file that declares base64 as decoded bytes', async () => {
const files = [
{
name: 'plot.png',
content: pngSignature.toString('base64'),
contentEncoding: FileContentEncoding.BASE64,
mimeType: 'image/png',
},
];

await materializeFiles(files, tempDir);

const written = await fs.readFile(path.join(tempDir, 'plot.png'));
expect(written).toEqual(pngSignature);
});

it('should write a file that declares utf-8 verbatim', async () => {
const files = [
{
name: 'data.csv',
content: 'a,b,c\nhello',
contentEncoding: FileContentEncoding.UTF8,
mimeType: 'text/csv',
},
];

await materializeFiles(files, tempDir);

const written = await fs.readFile(path.join(tempDir, 'data.csv'), 'utf8');
expect(written).toBe('a,b,c\nhello');
});
});
});
Loading