adk-typescript
Version:
TypeScript port of Google's Agent Development Kit (ADK)
195 lines (194 loc) • 6.76 kB
JavaScript
"use strict";
/**
* The persistent context used to configure the code executor.
*/
Object.defineProperty(exports, "__esModule", { value: true });
exports.CodeExecutorContext = void 0;
const CONTEXT_KEY = '_code_execution_context';
const SESSION_ID_KEY = 'execution_session_id';
const PROCESSED_FILE_NAMES_KEY = 'processed_input_files';
const INPUT_FILE_KEY = '_code_executor_input_files';
const ERROR_COUNT_KEY = '_code_executor_error_counts';
const CODE_EXECUTION_RESULTS_KEY = '_code_execution_results';
/**
* The persistent context used to configure the code executor.
*/
class CodeExecutorContext {
/**
* Initializes the code executor context.
*
* @param sessionState - The session state to get the code executor context from.
*/
constructor(sessionState) {
this.context = this.getCodeExecutorContext(sessionState);
this.sessionState = sessionState;
}
/**
* Gets the state delta to update in the persistent session state.
*
* @returns The state delta to update in the persistent session state.
*/
getStateDelta() {
const contextToUpdate = JSON.parse(JSON.stringify(this.context));
return { [CONTEXT_KEY]: contextToUpdate };
}
/**
* Gets the session ID for the code executor.
*
* @returns The session ID for the code executor context.
*/
getExecutionId() {
if (!(SESSION_ID_KEY in this.context)) {
return undefined;
}
return this.context[SESSION_ID_KEY];
}
/**
* Sets the session ID for the code executor.
*
* @param sessionId - The session ID for the code executor.
*/
setExecutionId(sessionId) {
this.context[SESSION_ID_KEY] = sessionId;
}
/**
* Gets the processed file names from the session state.
*
* @returns A list of processed file names in the code executor context.
*/
getProcessedFileNames() {
if (!(PROCESSED_FILE_NAMES_KEY in this.context)) {
return [];
}
return this.context[PROCESSED_FILE_NAMES_KEY];
}
/**
* Adds the processed file name to the session state.
*
* @param fileNames - The processed file names to add to the session state.
*/
addProcessedFileNames(fileNames) {
if (!(PROCESSED_FILE_NAMES_KEY in this.context)) {
this.context[PROCESSED_FILE_NAMES_KEY] = [];
}
this.context[PROCESSED_FILE_NAMES_KEY].push(...fileNames);
}
/**
* Gets the code executor input file names from the session state.
*
* @returns A list of input files in the code executor context.
*/
getInputFiles() {
if (!this.sessionState.has(INPUT_FILE_KEY)) {
return [];
}
return this.sessionState.get(INPUT_FILE_KEY).map(file => file);
}
/**
* Adds the input files to the code executor context.
*
* @param inputFiles - The input files to add to the code executor context.
*/
addInputFiles(inputFiles) {
if (!this.sessionState.has(INPUT_FILE_KEY)) {
this.sessionState.set(INPUT_FILE_KEY, []);
}
const currentFiles = this.sessionState.get(INPUT_FILE_KEY);
for (const inputFile of inputFiles) {
currentFiles.push(inputFile);
}
this.sessionState.set(INPUT_FILE_KEY, currentFiles);
}
/**
* Removes the input files and processed file names to the code executor context.
*/
clearInputFiles() {
if (this.sessionState.has(INPUT_FILE_KEY)) {
this.sessionState.set(INPUT_FILE_KEY, []);
}
if (PROCESSED_FILE_NAMES_KEY in this.context) {
this.context[PROCESSED_FILE_NAMES_KEY] = [];
}
}
/**
* Gets the error count from the session state.
*
* @param invocationId - The invocation ID to get the error count for.
* @returns The error count for the given invocation ID.
*/
getErrorCount(invocationId) {
if (!this.sessionState.has(ERROR_COUNT_KEY)) {
return 0;
}
const errorCounts = this.sessionState.get(ERROR_COUNT_KEY);
return errorCounts[invocationId] || 0;
}
/**
* Increments the error count from the session state.
*
* @param invocationId - The invocation ID to increment the error count for.
*/
incrementErrorCount(invocationId) {
let errorCounts = {};
if (this.sessionState.has(ERROR_COUNT_KEY)) {
errorCounts = this.sessionState.get(ERROR_COUNT_KEY);
}
errorCounts[invocationId] = this.getErrorCount(invocationId) + 1;
this.sessionState.set(ERROR_COUNT_KEY, errorCounts);
}
/**
* Resets the error count from the session state.
*
* @param invocationId - The invocation ID to reset the error count for.
*/
resetErrorCount(invocationId) {
if (!this.sessionState.has(ERROR_COUNT_KEY)) {
return;
}
const errorCounts = this.sessionState.get(ERROR_COUNT_KEY);
if (invocationId in errorCounts) {
delete errorCounts[invocationId];
this.sessionState.set(ERROR_COUNT_KEY, errorCounts);
}
}
/**
* Updates the code execution result.
*
* @param invocationId - The invocation ID to update the code execution result for.
* @param code - The code to execute.
* @param resultStdout - The standard output of the code execution.
* @param resultStderr - The standard error of the code execution.
*/
updateCodeExecutionResult(invocationId, code, resultStdout, resultStderr) {
let results = {};
if (this.sessionState.has(CODE_EXECUTION_RESULTS_KEY)) {
results = this.sessionState.get(CODE_EXECUTION_RESULTS_KEY);
}
else {
this.sessionState.set(CODE_EXECUTION_RESULTS_KEY, results);
}
if (!(invocationId in results)) {
results[invocationId] = [];
}
results[invocationId].push({
code,
result_stdout: resultStdout,
result_stderr: resultStderr,
timestamp: Math.floor(Date.now() / 1000),
});
this.sessionState.set(CODE_EXECUTION_RESULTS_KEY, results);
}
/**
* Gets the code executor context from the session state.
*
* @param sessionState - The session state to get the code executor context from.
* @returns A dict of code executor context.
*/
getCodeExecutorContext(sessionState) {
if (!sessionState.has(CONTEXT_KEY)) {
sessionState.set(CONTEXT_KEY, {});
}
return sessionState.get(CONTEXT_KEY);
}
}
exports.CodeExecutorContext = CodeExecutorContext;