adk-typescript
Version:
TypeScript port of Google's Agent Development Kit (ADK)
709 lines (708 loc) • 26.9 kB
JavaScript
"use strict";
var __decorate = (this && this.__decorate) || function (decorators, target, key, desc) {
var c = arguments.length, r = c < 3 ? target : desc === null ? desc = Object.getOwnPropertyDescriptor(target, key) : desc, d;
if (typeof Reflect === "object" && typeof Reflect.decorate === "function") r = Reflect.decorate(decorators, target, key, desc);
else for (var i = decorators.length - 1; i >= 0; i--) if (d = decorators[i]) r = (c < 3 ? d(r) : c > 3 ? d(target, key, r) : d(target, key)) || r;
return c > 3 && r && Object.defineProperty(target, key, r), r;
};
var __metadata = (this && this.__metadata) || function (k, v) {
if (typeof Reflect === "object" && typeof Reflect.metadata === "function") return Reflect.metadata(k, v);
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.DatabaseSessionService = void 0;
// Make sure we import reflect-metadata first
require("reflect-metadata");
const uuid_1 = require("uuid");
const typeorm_1 = require("typeorm");
const BaseSessionService_1 = require("./BaseSessionService");
const State_1 = require("./State");
const sessionUtils_1 = require("./sessionUtils");
/**
* Default maximum length for key columns in the database
*/
const DEFAULT_MAX_KEY_LENGTH = 128;
/**
* Default maximum length for VARCHAR columns in the database
*/
const DEFAULT_MAX_VARCHAR_LENGTH = 256;
/**
* Extract state delta from a state object, categorizing into app, user, and session state
*/
function extractStateDelta(state) {
const appStateDelta = {};
const userStateDelta = {};
const sessionStateDelta = {};
if (state) {
for (const [key, value] of Object.entries(state)) {
if (key.startsWith(State_1.StatePrefix.APP_PREFIX)) {
appStateDelta[key.substring(State_1.StatePrefix.APP_PREFIX.length)] = value;
}
else if (key.startsWith(State_1.StatePrefix.USER_PREFIX)) {
userStateDelta[key.substring(State_1.StatePrefix.USER_PREFIX.length)] = value;
}
else if (!key.startsWith(State_1.StatePrefix.TEMP_PREFIX)) {
sessionStateDelta[key] = value;
}
}
}
return [appStateDelta, userStateDelta, sessionStateDelta];
}
/**
* Merge app, user, and session state into a single state object
*/
function mergeState(appState, userState, sessionState) {
const mergedState = new State_1.State(sessionState);
for (const [key, value] of Object.entries(appState)) {
mergedState.set(State_1.StatePrefix.APP_PREFIX + key, value);
}
for (const [key, value] of Object.entries(userState)) {
mergedState.set(State_1.StatePrefix.USER_PREFIX + key, value);
}
return mergedState;
}
/**
* Represents a session stored in the database
*/
let StorageSession = class StorageSession {
};
__decorate([
(0, typeorm_1.PrimaryColumn)({ length: DEFAULT_MAX_KEY_LENGTH }),
__metadata("design:type", String)
], StorageSession.prototype, "appName", void 0);
__decorate([
(0, typeorm_1.PrimaryColumn)({ length: DEFAULT_MAX_KEY_LENGTH }),
__metadata("design:type", String)
], StorageSession.prototype, "userId", void 0);
__decorate([
(0, typeorm_1.PrimaryColumn)({ length: DEFAULT_MAX_KEY_LENGTH }),
__metadata("design:type", String)
], StorageSession.prototype, "id", void 0);
__decorate([
(0, typeorm_1.Column)({ type: 'simple-json', default: '{}' }),
__metadata("design:type", Object)
], StorageSession.prototype, "state", void 0);
__decorate([
(0, typeorm_1.CreateDateColumn)(),
__metadata("design:type", Date)
], StorageSession.prototype, "createTime", void 0);
__decorate([
(0, typeorm_1.UpdateDateColumn)(),
__metadata("design:type", Date)
], StorageSession.prototype, "updateTime", void 0);
__decorate([
(0, typeorm_1.OneToMany)(() => StorageEvent, event => event.storageSession, { cascade: true }),
__metadata("design:type", Array)
], StorageSession.prototype, "storageEvents", void 0);
StorageSession = __decorate([
(0, typeorm_1.Entity)({ name: 'sessions' })
], StorageSession);
/**
* Represents an event stored in the database
*/
let StorageEvent = class StorageEvent {
/**
* Gets the set of long-running tool IDs
*/
get longRunningToolIds() {
return this.longRunningToolIdsJson
? new Set(JSON.parse(this.longRunningToolIdsJson))
: new Set();
}
/**
* Sets the long-running tool IDs
*/
set longRunningToolIds(value) {
if (value) {
this.longRunningToolIdsJson = JSON.stringify(Array.from(value));
}
else {
this.longRunningToolIdsJson = undefined;
}
}
};
__decorate([
(0, typeorm_1.PrimaryColumn)({ length: DEFAULT_MAX_KEY_LENGTH }),
__metadata("design:type", String)
], StorageEvent.prototype, "id", void 0);
__decorate([
(0, typeorm_1.PrimaryColumn)({ length: DEFAULT_MAX_KEY_LENGTH }),
__metadata("design:type", String)
], StorageEvent.prototype, "appName", void 0);
__decorate([
(0, typeorm_1.PrimaryColumn)({ length: DEFAULT_MAX_KEY_LENGTH }),
__metadata("design:type", String)
], StorageEvent.prototype, "userId", void 0);
__decorate([
(0, typeorm_1.PrimaryColumn)({ length: DEFAULT_MAX_KEY_LENGTH }),
__metadata("design:type", String)
], StorageEvent.prototype, "sessionId", void 0);
__decorate([
(0, typeorm_1.Column)({ length: DEFAULT_MAX_VARCHAR_LENGTH }),
__metadata("design:type", String)
], StorageEvent.prototype, "invocationId", void 0);
__decorate([
(0, typeorm_1.Column)({ length: DEFAULT_MAX_VARCHAR_LENGTH }),
__metadata("design:type", String)
], StorageEvent.prototype, "author", void 0);
__decorate([
(0, typeorm_1.Column)({ nullable: true, length: DEFAULT_MAX_VARCHAR_LENGTH }),
__metadata("design:type", String)
], StorageEvent.prototype, "branch", void 0);
__decorate([
(0, typeorm_1.CreateDateColumn)(),
__metadata("design:type", Date)
], StorageEvent.prototype, "timestamp", void 0);
__decorate([
(0, typeorm_1.Column)({ type: 'simple-json', nullable: true }),
__metadata("design:type", Object)
], StorageEvent.prototype, "content", void 0);
__decorate([
(0, typeorm_1.Column)({ type: 'simple-json' }),
__metadata("design:type", Object)
], StorageEvent.prototype, "actions", void 0);
__decorate([
(0, typeorm_1.Column)({ nullable: true }),
__metadata("design:type", String)
], StorageEvent.prototype, "longRunningToolIdsJson", void 0);
__decorate([
(0, typeorm_1.Column)({ type: 'simple-json', nullable: true }),
__metadata("design:type", Object)
], StorageEvent.prototype, "groundingMetadata", void 0);
__decorate([
(0, typeorm_1.Column)({ nullable: true }),
__metadata("design:type", Boolean)
], StorageEvent.prototype, "partial", void 0);
__decorate([
(0, typeorm_1.Column)({ nullable: true }),
__metadata("design:type", Boolean)
], StorageEvent.prototype, "turnComplete", void 0);
__decorate([
(0, typeorm_1.Column)({ nullable: true, length: DEFAULT_MAX_VARCHAR_LENGTH }),
__metadata("design:type", String)
], StorageEvent.prototype, "errorCode", void 0);
__decorate([
(0, typeorm_1.Column)({ nullable: true, length: 1024 }),
__metadata("design:type", String)
], StorageEvent.prototype, "errorMessage", void 0);
__decorate([
(0, typeorm_1.Column)({ nullable: true }),
__metadata("design:type", Boolean)
], StorageEvent.prototype, "interrupted", void 0);
__decorate([
(0, typeorm_1.ManyToOne)(() => StorageSession, session => session.storageEvents, {
onDelete: 'CASCADE'
}),
(0, typeorm_1.JoinColumn)([
{ name: 'appName', referencedColumnName: 'appName' },
{ name: 'userId', referencedColumnName: 'userId' },
{ name: 'sessionId', referencedColumnName: 'id' }
]),
__metadata("design:type", StorageSession)
], StorageEvent.prototype, "storageSession", void 0);
StorageEvent = __decorate([
(0, typeorm_1.Entity)({ name: 'events' })
], StorageEvent);
/**
* Represents app-wide state stored in the database
*/
let StorageAppState = class StorageAppState {
};
__decorate([
(0, typeorm_1.PrimaryColumn)({ length: DEFAULT_MAX_KEY_LENGTH }),
__metadata("design:type", String)
], StorageAppState.prototype, "appName", void 0);
__decorate([
(0, typeorm_1.Column)({ type: 'simple-json', default: '{}' }),
__metadata("design:type", Object)
], StorageAppState.prototype, "state", void 0);
__decorate([
(0, typeorm_1.UpdateDateColumn)(),
__metadata("design:type", Date)
], StorageAppState.prototype, "updateTime", void 0);
StorageAppState = __decorate([
(0, typeorm_1.Entity)({ name: 'app_states' })
], StorageAppState);
/**
* Represents user-specific state stored in the database
*/
let StorageUserState = class StorageUserState {
};
__decorate([
(0, typeorm_1.PrimaryColumn)({ length: DEFAULT_MAX_KEY_LENGTH }),
__metadata("design:type", String)
], StorageUserState.prototype, "appName", void 0);
__decorate([
(0, typeorm_1.PrimaryColumn)({ length: DEFAULT_MAX_KEY_LENGTH }),
__metadata("design:type", String)
], StorageUserState.prototype, "userId", void 0);
__decorate([
(0, typeorm_1.Column)({ type: 'simple-json', default: '{}' }),
__metadata("design:type", Object)
], StorageUserState.prototype, "state", void 0);
__decorate([
(0, typeorm_1.UpdateDateColumn)(),
__metadata("design:type", Date)
], StorageUserState.prototype, "updateTime", void 0);
StorageUserState = __decorate([
(0, typeorm_1.Entity)({ name: 'user_states' })
], StorageUserState);
/**
* Parse a database URL to determine the database type and configuration
*/
function parseDbUrl(dbUrl) {
if (dbUrl.startsWith('sqlite://')) {
const path = dbUrl.replace('sqlite://', '');
if (path === '/:memory:') {
// In-memory SQLite database
return {
type: 'sqlite',
database: ':memory:',
synchronize: true,
entities: [StorageSession, StorageEvent, StorageAppState, StorageUserState],
logging: ['error']
};
}
else {
// File-based SQLite database
return {
type: 'sqlite',
database: path,
synchronize: true,
entities: [StorageSession, StorageEvent, StorageAppState, StorageUserState],
logging: ['error']
};
}
}
// Other database types could be added here (PostgreSQL, MySQL, etc.)
throw new Error(`Unsupported database URL format: ${dbUrl}`);
}
/**
* DatabaseConnectionManager manages creating and reusing database connections
*/
class DatabaseConnectionManager {
/**
* Get (or create) a connection for the given database URL
*/
static async getConnection(dbUrl) {
if (this.connections.has(dbUrl)) {
const connection = this.connections.get(dbUrl);
// If the connection was closed, re-initialize it
if (!connection.isInitialized) {
await connection.initialize();
}
return connection;
}
// Create a new connection
const options = parseDbUrl(dbUrl);
const dataSource = new typeorm_1.DataSource(options);
try {
await dataSource.initialize();
this.connections.set(dbUrl, dataSource);
return dataSource;
}
catch (error) {
if (error instanceof Error) {
throw new Error(`Failed to initialize database connection: ${error.message}`);
}
throw new Error('Unknown error initializing database connection');
}
}
/**
* Close all open connections
*/
static async closeAllConnections() {
for (const connection of this.connections.values()) {
if (connection.isInitialized) {
await connection.destroy();
}
}
this.connections.clear();
}
}
DatabaseConnectionManager.connections = new Map();
/**
* A session service that uses a database for persistent storage.
* This implementation uses TypeORM with SQLite by default.
*/
class DatabaseSessionService extends BaseSessionService_1.BaseSessionService {
/**
* Creates a new DatabaseSessionService.
* @param dbUrl The database URL (e.g., 'sqlite:///:memory:' for in-memory SQLite)
*/
constructor(dbUrl) {
super();
this.dbUrl = dbUrl;
}
/**
* Closes the database connection for this instance.
* This is useful for testing to properly clean up resources.
*/
async closeConnection() {
if (this.connection && this.connection.isInitialized) {
await this.connection.destroy();
this.connection = undefined;
}
}
/**
* Closes all database connections.
* This static method is useful for testing to properly clean up resources.
*/
static async closeAllConnections() {
await DatabaseConnectionManager.closeAllConnections();
}
/**
* Ensures the database connection is established
*/
async ensureConnection() {
if (!this.connection) {
this.connection = await DatabaseConnectionManager.getConnection(this.dbUrl);
// Initialize repositories
this.sessionRepo = this.connection.getRepository(StorageSession);
this.eventRepo = this.connection.getRepository(StorageEvent);
this.appStateRepo = this.connection.getRepository(StorageAppState);
this.userStateRepo = this.connection.getRepository(StorageUserState);
}
}
/**
* Creates a new session.
*/
async createSession(options) {
await this.ensureConnection();
const { appName, userId, state = {} } = options;
const sessionId = options.sessionId || (0, uuid_1.v4)();
// Fetch app and user states from storage or create if they don't exist
let appState = await this.appStateRepo.findOneBy({ appName });
if (!appState) {
appState = new StorageAppState();
appState.appName = appName;
appState.state = {};
await this.appStateRepo.save(appState);
}
let userState = await this.userStateRepo.findOneBy({ appName, userId });
if (!userState) {
userState = new StorageUserState();
userState.appName = appName;
userState.userId = userId;
userState.state = {};
await this.userStateRepo.save(userState);
}
// Extract state deltas
const [appStateDelta, userStateDelta, sessionState] = extractStateDelta(state);
// Apply state deltas
if (Object.keys(appStateDelta).length > 0) {
Object.assign(appState.state, appStateDelta);
await this.appStateRepo.save(appState);
}
if (Object.keys(userStateDelta).length > 0) {
Object.assign(userState.state, userStateDelta);
await this.userStateRepo.save(userState);
}
// Create and save the session
const storageSession = new StorageSession();
storageSession.appName = appName;
storageSession.userId = userId;
storageSession.id = sessionId;
storageSession.state = sessionState;
await this.sessionRepo.save(storageSession);
// Return the session with merged state
const mergedState = mergeState(appState.state, userState.state, sessionState);
return {
id: sessionId,
appName,
userId,
state: mergedState,
events: []
};
}
/**
* Gets a session by its ID.
*/
async getSession(options) {
await this.ensureConnection();
const { appName, userId, sessionId } = options;
// Fetch the session
const storageSession = await this.sessionRepo.findOneBy({
appName,
userId,
id: sessionId
});
if (!storageSession) {
return null;
}
// Fetch app and user states
const appState = await this.appStateRepo.findOneBy({ appName });
const userState = await this.userStateRepo.findOneBy({ appName, userId });
// Fetch events for this session
const storageEvents = await this.eventRepo.findBy({
appName,
userId,
sessionId
});
// Create the session object with merged state
const session = {
id: sessionId,
appName,
userId,
state: mergeState(appState?.state || {}, userState?.state || {}, storageSession.state),
events: []
};
// Convert storage events to Event objects
session.events = storageEvents.map(storageEvent => ({
id: storageEvent.id,
invocationId: storageEvent.invocationId,
author: storageEvent.author,
content: (0, sessionUtils_1.decodeContent)(storageEvent.content),
actions: storageEvent.actions,
turnComplete: storageEvent.turnComplete,
partial: storageEvent.partial,
longRunningToolIds: storageEvent.longRunningToolIds,
errorCode: storageEvent.errorCode,
errorMessage: storageEvent.errorMessage,
interrupted: storageEvent.interrupted
}));
return session;
}
/**
* Lists all sessions for a user in an app.
*/
async listSessions(options) {
await this.ensureConnection();
const { appName, userId } = options;
// Get sessions but only those with the exact app name and user ID
const storageSessions = await this.sessionRepo.findBy({
appName: appName,
userId: userId
});
// Create session objects without events or full state
const sessions = storageSessions.map(storageSession => ({
id: storageSession.id,
appName,
userId,
state: new State_1.State(),
events: []
}));
return { sessions };
}
/**
* Deletes a session.
*/
async deleteSession(options) {
await this.ensureConnection();
const { appName, userId, sessionId } = options;
await this.sessionRepo.delete({
appName,
userId,
id: sessionId
});
}
/**
* Appends an event to a session.
*/
async appendEvent(options) {
await this.ensureConnection();
const { session, event } = options;
if (event.partial) {
// Don't persist partial events
session.events.push(event);
return;
}
// Make sure the session exists
const storageSession = await this.sessionRepo.findOneBy({
appName: session.appName,
userId: session.userId,
id: session.id
});
if (!storageSession) {
// Create the session if it doesn't exist
await this.createSession({
appName: session.appName,
userId: session.userId,
sessionId: session.id
});
}
// Fetch app and user states
let appState = await this.appStateRepo.findOneBy({ appName: session.appName });
if (!appState) {
appState = new StorageAppState();
appState.appName = session.appName;
appState.state = {};
await this.appStateRepo.save(appState);
}
let userState = await this.userStateRepo.findOneBy({
appName: session.appName,
userId: session.userId
});
if (!userState) {
userState = new StorageUserState();
userState.appName = session.appName;
userState.userId = session.userId;
userState.state = {};
await this.userStateRepo.save(userState);
}
// Handle state delta from the event
if (event.actions?.stateDelta) {
const [appStateDelta, userStateDelta, sessionStateDelta] = extractStateDelta(event.actions.stateDelta);
// Update app state
if (Object.keys(appStateDelta).length > 0) {
Object.assign(appState.state, appStateDelta);
await this.appStateRepo.save(appState);
}
// Update user state
if (Object.keys(userStateDelta).length > 0) {
Object.assign(userState.state, userStateDelta);
await this.userStateRepo.save(userState);
}
// Update session state
if (Object.keys(sessionStateDelta).length > 0) {
await this.sessionRepo.findOneBy({
appName: session.appName,
userId: session.userId,
id: session.id
}).then(storageSession => {
if (storageSession) {
Object.assign(storageSession.state, sessionStateDelta);
return this.sessionRepo.save(storageSession);
}
});
}
// Update the in-memory session state
session.state = mergeState(appState.state, userState.state, { ...session.state, ...sessionStateDelta });
}
// Create and save the event
const storageEvent = new StorageEvent();
storageEvent.id = event.id || (0, uuid_1.v4)();
storageEvent.appName = session.appName;
storageEvent.userId = session.userId;
storageEvent.sessionId = session.id;
storageEvent.invocationId = event.invocationId;
storageEvent.author = event.author;
storageEvent.content = event.content ? (0, sessionUtils_1.encodeContent)(event.content) : undefined;
storageEvent.actions = event.actions || {};
storageEvent.turnComplete = event.turnComplete;
storageEvent.partial = event.partial;
storageEvent.longRunningToolIds = event.longRunningToolIds;
storageEvent.errorCode = event.errorCode;
storageEvent.errorMessage = event.errorMessage;
storageEvent.interrupted = event.interrupted;
await this.eventRepo.save(storageEvent);
// Add event to the session
if (!event.id) {
event.id = storageEvent.id;
}
session.events.push(event);
}
/**
* Lists events in a session.
*/
async listEvents(options) {
await this.ensureConnection();
const { appName, userId, sessionId } = options;
// Fetch the session
const storageSession = await this.sessionRepo.findOneBy({
appName,
userId,
id: sessionId
});
if (!storageSession) {
return { events: [] };
}
// Fetch events for this session
const storageEvents = await this.eventRepo.findBy({
appName,
userId,
sessionId
});
// Convert storage events to Event objects
const events = storageEvents.map(storageEvent => ({
id: storageEvent.id,
invocationId: storageEvent.invocationId,
author: storageEvent.author,
content: (0, sessionUtils_1.decodeContent)(storageEvent.content),
actions: storageEvent.actions,
turnComplete: storageEvent.turnComplete,
partial: storageEvent.partial,
longRunningToolIds: storageEvent.longRunningToolIds,
errorCode: storageEvent.errorCode,
errorMessage: storageEvent.errorMessage,
interrupted: storageEvent.interrupted
}));
return { events };
}
/**
* Updates the state of a session.
*/
async updateSessionState(appName, userId, sessionId, stateDelta) {
await this.ensureConnection();
// Fetch the session
const storageSession = await this.sessionRepo.findOneBy({
appName,
userId,
id: sessionId
});
if (!storageSession) {
throw new Error(`Session ${sessionId} not found for user ${userId} in app ${appName}`);
}
// Fetch app and user states
let appState = await this.appStateRepo.findOneBy({ appName });
if (!appState) {
appState = new StorageAppState();
appState.appName = appName;
appState.state = {};
await this.appStateRepo.save(appState);
}
let userState = await this.userStateRepo.findOneBy({ appName, userId });
if (!userState) {
userState = new StorageUserState();
userState.appName = appName;
userState.userId = userId;
userState.state = {};
await this.userStateRepo.save(userState);
}
// Extract state deltas
const [appStateDelta, userStateDelta, sessionStateDelta] = extractStateDelta(stateDelta);
// Apply state deltas
if (Object.keys(appStateDelta).length > 0) {
Object.assign(appState.state, appStateDelta);
await this.appStateRepo.save(appState);
}
if (Object.keys(userStateDelta).length > 0) {
Object.assign(userState.state, userStateDelta);
await this.userStateRepo.save(userState);
}
// Update session state
if (Object.keys(sessionStateDelta).length > 0) {
Object.assign(storageSession.state, sessionStateDelta);
await this.sessionRepo.save(storageSession);
}
// Fetch events for this session
const storageEvents = await this.eventRepo.findBy({
appName,
userId,
sessionId
});
// Create the session object
const session = {
id: sessionId,
appName,
userId,
state: mergeState(appState.state, userState.state, storageSession.state),
events: []
};
// Convert storage events to Event objects
session.events = storageEvents.map(storageEvent => ({
id: storageEvent.id,
invocationId: storageEvent.invocationId,
author: storageEvent.author,
content: (0, sessionUtils_1.decodeContent)(storageEvent.content),
actions: storageEvent.actions,
turnComplete: storageEvent.turnComplete,
partial: storageEvent.partial,
longRunningToolIds: storageEvent.longRunningToolIds,
errorCode: storageEvent.errorCode,
errorMessage: storageEvent.errorMessage,
interrupted: storageEvent.interrupted
}));
return session;
}
}
exports.DatabaseSessionService = DatabaseSessionService;