adk-typescript
Version:
TypeScript port of Google's Agent Development Kit (ADK)
735 lines (734 loc) • 32 kB
JavaScript
;
var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
if (k2 === undefined) k2 = k;
var desc = Object.getOwnPropertyDescriptor(m, k);
if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
desc = { enumerable: true, get: function() { return m[k]; } };
}
Object.defineProperty(o, k2, desc);
}) : (function(o, m, k, k2) {
if (k2 === undefined) k2 = k;
o[k2] = m[k];
}));
var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
Object.defineProperty(o, "default", { enumerable: true, value: v });
}) : function(o, v) {
o["default"] = v;
});
var __importStar = (this && this.__importStar) || (function () {
var ownKeys = function(o) {
ownKeys = Object.getOwnPropertyNames || function (o) {
var ar = [];
for (var k in o) if (Object.prototype.hasOwnProperty.call(o, k)) ar[ar.length] = k;
return ar;
};
return ownKeys(o);
};
return function (mod) {
if (mod && mod.__esModule) return mod;
var result = {};
if (mod != null) for (var k = ownKeys(mod), i = 0; i < k.length; i++) if (k[i] !== "default") __createBinding(result, mod, k[i]);
__setModuleDefault(result, mod);
return result;
};
})();
var __importDefault = (this && this.__importDefault) || function (mod) {
return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.createApiServer = createApiServer;
const express_1 = __importDefault(require("express"));
const http_1 = __importDefault(require("http"));
const path_1 = __importDefault(require("path"));
const fs_1 = __importDefault(require("fs"));
const socket_io_1 = require("socket.io");
const cors_1 = __importDefault(require("cors"));
const RunConfig_1 = require("../agents/RunConfig");
const InMemoryArtifactService_1 = require("../artifacts/InMemoryArtifactService");
const InMemoryMemoryService_1 = require("../memory/InMemoryMemoryService");
const DatabaseSessionService_1 = require("../sessions/DatabaseSessionService");
const InMemorySessionService_1 = require("../sessions/InMemorySessionService");
const runners_1 = require("../runners");
const utils_1 = require("./utils");
// Import the LiveRequestQueue and LiveRequest interfaces/classes
const LiveRequestQueue_1 = require("../agents/LiveRequestQueue");
// Import agent_graph
const agentGraph = __importStar(require("./agentGraph"));
// Constant for eval session ID prefix
const EVAL_SESSION_ID_PREFIX = 'eval_';
// Constant for eval set file extension
const EVAL_SET_FILE_EXTENSION = '.evalset.json';
/**
* Creates an Express app that serves as an API server for agents
*
* @param options Configuration options for the server
* @returns The configured Express app and server
*/
function createApiServer(options) {
const { agentDir, sessionDbUrl = '', allowOrigins = ['*'], web = false, traceToCloud = false, port = 8000 } = options;
// Trace dictionary for storing trace information
const traceDict = {};
// Create the Express app
const app = (0, express_1.default)();
const server = http_1.default.createServer(app);
const io = new socket_io_1.Server(server, {
cors: {
origin: allowOrigins,
methods: ['GET', 'POST'],
credentials: true
}
});
// Add middleware
app.use(express_1.default.json());
app.use((0, cors_1.default)({
origin: allowOrigins,
credentials: true
}));
// Add the agent directory to the module search path
if (!process.env.NODE_PATH?.includes(agentDir)) {
process.env.NODE_PATH = (process.env.NODE_PATH || '') + path_1.default.delimiter + agentDir;
// Force Node.js to reload the module paths
require('module').Module._initPaths();
}
// Initialize services
const runnerDict = {};
const rootAgentDict = {};
// Build the Artifact service
const artifactService = new InMemoryArtifactService_1.InMemoryArtifactService();
const memoryService = new InMemoryMemoryService_1.InMemoryMemoryService();
// Build the Session service
const agentEngineId = '';
let sessionService; // Temporary any type to fix linter error
if (sessionDbUrl) {
if (sessionDbUrl.startsWith('agentengine://')) {
// TODO: Implement VertexAI session service for TypeScript version
throw new Error('VertexAI session service not implemented in TypeScript version yet');
}
else {
sessionService = new DatabaseSessionService_1.DatabaseSessionService(sessionDbUrl);
}
}
else {
sessionService = new InMemorySessionService_1.InMemorySessionService();
}
// Define API endpoints
app.get('/list-apps', (req, res) => {
const basePath = path_1.default.resolve(agentDir);
if (!fs_1.default.existsSync(basePath)) {
return res.status(404).json({ error: 'Path not found' });
}
if (!fs_1.default.statSync(basePath).isDirectory()) {
return res.status(400).json({ error: 'Not a directory' });
}
const agentNames = fs_1.default.readdirSync(basePath)
.filter(x => {
const fullPath = path_1.default.join(basePath, x);
return fs_1.default.statSync(fullPath).isDirectory() &&
!x.startsWith('.') &&
x !== 'node_modules';
})
.sort();
res.json(agentNames);
});
app.get('/debug/trace/:eventId', (req, res) => {
const { eventId } = req.params;
const eventDict = traceDict[eventId];
if (!eventDict) {
return res.status(404).json({ error: 'Trace not found' });
}
res.json(eventDict);
});
// Session management endpoints
app.get('/apps/:appName/users/:userId/sessions/:sessionId', (req, res) => {
const { appName, userId, sessionId } = req.params;
// Connect to managed session if agent_engine_id is set
const effectiveAppName = agentEngineId || appName;
const session = sessionService.getSession({
appName: effectiveAppName,
userId,
sessionId
});
if (!session) {
return res.status(404).json({ error: 'Session not found' });
}
res.json(session);
});
app.get('/apps/:appName/users/:userId/sessions', (req, res) => {
const { appName, userId } = req.params;
// Connect to managed session if agent_engine_id is set
const effectiveAppName = agentEngineId || appName;
const sessions = sessionService.listSessions({
appName: effectiveAppName,
userId
}).sessions.filter((session) =>
// Remove sessions that were generated as a part of Eval
!session.id.startsWith(EVAL_SESSION_ID_PREFIX));
res.json(sessions);
});
app.post('/apps/:appName/users/:userId/sessions/:sessionId', (req, res) => {
const { appName, userId, sessionId } = req.params;
const state = req.body?.state || undefined;
// Connect to managed session if agent_engine_id is set
const effectiveAppName = agentEngineId || appName;
if (sessionService.getSession({
appName: effectiveAppName,
userId,
sessionId
})) {
return res.status(400).json({ error: `Session already exists: ${sessionId}` });
}
const session = sessionService.createSession({
appName: effectiveAppName,
userId,
state,
sessionId
});
console.log(`New session created: ${sessionId}`);
res.json(session);
});
app.post('/apps/:appName/users/:userId/sessions', (req, res) => {
const { appName, userId } = req.params;
const state = req.body?.state || undefined;
// Connect to managed session if agent_engine_id is set
const effectiveAppName = agentEngineId || appName;
const session = sessionService.createSession({
appName: effectiveAppName,
userId,
state
});
console.log('New session created');
res.json(session);
});
app.delete('/apps/:appName/users/:userId/sessions/:sessionId', (req, res) => {
const { appName, userId, sessionId } = req.params;
// Connect to managed session if agent_engine_id is set
const effectiveAppName = agentEngineId || appName;
sessionService.deleteSession({
appName: effectiveAppName,
userId,
sessionId
});
res.status(204).send();
});
// Artifact management endpoints
app.get('/apps/:appName/users/:userId/sessions/:sessionId/artifacts/:artifactName', async (req, res) => {
const { appName, userId, sessionId, artifactName } = req.params;
const version = req.query.version ? parseInt(req.query.version) : undefined;
// Connect to managed session if agent_engine_id is set
const effectiveAppName = agentEngineId || appName;
try {
const artifactResult = artifactService.loadArtifact({
appName: effectiveAppName,
userId,
sessionId,
filename: artifactName,
version
});
// Handle both synchronous and asynchronous cases
const artifact = artifactResult instanceof Promise ? await artifactResult : artifactResult;
if (!artifact) {
return res.status(404).json({ error: 'Artifact not found' });
}
res.json(artifact);
}
catch (error) {
console.error('Error loading artifact:', error);
res.status(500).json({ error: 'Error loading artifact' });
}
});
app.get('/apps/:appName/users/:userId/sessions/:sessionId/artifacts/:artifactName/versions/:versionId', async (req, res) => {
const { appName, userId, sessionId, artifactName, versionId } = req.params;
const version = parseInt(versionId);
// Connect to managed session if agent_engine_id is set
const effectiveAppName = agentEngineId || appName;
try {
const artifactResult = artifactService.loadArtifact({
appName: effectiveAppName,
userId,
sessionId,
filename: artifactName,
version
});
// Handle both synchronous and asynchronous cases
const artifact = artifactResult instanceof Promise ? await artifactResult : artifactResult;
if (!artifact) {
return res.status(404).json({ error: 'Artifact not found' });
}
res.json(artifact);
}
catch (error) {
console.error('Error loading artifact version:', error);
res.status(500).json({ error: 'Error loading artifact version' });
}
});
app.get('/apps/:appName/users/:userId/sessions/:sessionId/artifacts', async (req, res) => {
const { appName, userId, sessionId } = req.params;
// Connect to managed session if agent_engine_id is set
const effectiveAppName = agentEngineId || appName;
try {
// Fix for listArtifactKeys - no filename parameter needed
const artifactNamesResult = artifactService.listArtifactKeys({
appName: effectiveAppName,
userId,
sessionId
}); // Using 'as any' to bypass TypeScript check as this is how the method is implemented
// Handle both synchronous and asynchronous cases
const artifactNames = artifactNamesResult instanceof Promise ? await artifactNamesResult : artifactNamesResult;
res.json(artifactNames);
}
catch (error) {
console.error('Error listing artifacts:', error);
res.status(500).json({ error: 'Error listing artifacts' });
}
});
app.get('/apps/:appName/users/:userId/sessions/:sessionId/artifacts/:artifactName/versions', async (req, res) => {
const { appName, userId, sessionId, artifactName } = req.params;
// Connect to managed session if agent_engine_id is set
const effectiveAppName = agentEngineId || appName;
try {
const versionsResult = artifactService.listVersions({
appName: effectiveAppName,
userId,
sessionId,
filename: artifactName
});
// Handle both synchronous and asynchronous cases
const versions = versionsResult instanceof Promise ? await versionsResult : versionsResult;
res.json(versions);
}
catch (error) {
console.error('Error listing versions:', error);
res.status(500).json({ error: 'Error listing versions' });
}
});
app.delete('/apps/:appName/users/:userId/sessions/:sessionId/artifacts/:artifactName', (req, res) => {
const { appName, userId, sessionId, artifactName } = req.params;
// Connect to managed session if agent_engine_id is set
const effectiveAppName = agentEngineId || appName;
artifactService.deleteArtifact({
appName: effectiveAppName,
userId,
sessionId,
filename: artifactName
});
res.status(204).send();
});
// Agent run endpoint
app.post('/run', async (req, res) => {
const runRequest = req.body;
// Connect to managed session if agent_engine_id is set
const appId = agentEngineId || runRequest.appName;
const session = sessionService.getSession({
appName: appId,
userId: runRequest.userId,
sessionId: runRequest.sessionId
});
if (!session) {
return res.status(404).json({ error: 'Session not found' });
}
try {
const runner = await getRunner(runRequest.appName);
const events = [];
// Collect all events from the runner
for await (const event of runner.runAsync({
userId: runRequest.userId,
sessionId: runRequest.sessionId,
newMessage: runRequest.newMessage
})) {
events.push(event);
}
console.log(`Generated ${events.length} events in agent run:`, events);
res.json(events);
}
catch (error) {
console.error('Error in agent run:', error);
res.status(500).json({ error: 'Error running agent' });
}
});
// SSE endpoint for streaming responses
app.post('/run_sse', async (req, res) => {
const runRequest = req.body;
// Connect to managed session if agent_engine_id is set
const appId = agentEngineId || runRequest.appName;
const session = sessionService.getSession({
appName: appId,
userId: runRequest.userId,
sessionId: runRequest.sessionId
});
if (!session) {
return res.status(404).json({ error: 'Session not found' });
}
// Set up SSE
res.setHeader('Content-Type', 'text/event-stream');
res.setHeader('Cache-Control', 'no-cache');
res.setHeader('Connection', 'keep-alive');
// Helper to send SSE data
const sendEvent = (data) => {
res.write(`data: ${JSON.stringify(data)}\n\n`);
};
try {
const runner = await getRunner(runRequest.appName);
const streamingMode = runRequest.streaming ? RunConfig_1.StreamingMode.SSE : RunConfig_1.StreamingMode.NONE;
// Run the agent and stream events
for await (const event of runner.runAsync({
userId: runRequest.userId,
sessionId: runRequest.sessionId,
newMessage: runRequest.newMessage,
runConfig: { streamingMode }
})) {
console.log('Generated event in agent run streaming:', event);
sendEvent(event);
}
// End the response
res.end();
}
catch (error) {
console.error('Error in SSE streaming:', error);
sendEvent({ error: String(error) });
res.end();
}
});
// Set up WebSocket for live interaction
io.on('connection', async (socket) => {
console.log(`New WebSocket connection: ${socket.id}`);
// Handle run_live event
socket.on('run_live', async (data) => {
try {
const { appName, userId, sessionId, modalities = ['TEXT'] } = data;
// Connect to managed session if agent_engine_id is set
const effectiveAppName = agentEngineId || appName;
const session = sessionService.getSession({
appName: effectiveAppName,
userId,
sessionId
});
if (!session) {
socket.emit('error', { error: 'Session not found' });
socket.disconnect(true);
return;
}
// Create a live request queue for this socket
const liveRequestQueue = new LiveRequestQueue_1.LiveRequestQueue();
// Set up two tasks - one for forwarding events and one for processing messages
// Task 1: Forward events from the runner to the client
const forwardEvents = async () => {
try {
const runner = await getRunner(appName);
for await (const event of runner.runLive({
session,
liveRequestQueue
})) {
// Send event to client
socket.emit('event', event);
}
}
catch (error) {
console.error('Error in forwardEvents:', error);
socket.emit('error', { error: String(error) });
}
};
// Task 2: Process incoming messages from the client
const processMessages = async () => {
// Set up message handler
socket.on('message', (messageData) => {
try {
// Validate and forward to the live request queue
// Use appropriate method on LiveRequestQueue based on message type
if (messageData.content) {
liveRequestQueue.sendContent(messageData.content);
}
else if (messageData.blob) {
liveRequestQueue.sendBlob(messageData.blob);
}
else if (messageData.close) {
liveRequestQueue.sendClose();
}
}
catch (error) {
console.error('Error processing message:', error);
socket.emit('error', { error: String(error) });
}
});
};
// Start both tasks
forwardEvents().catch(error => {
console.error('Error in forward events task:', error);
socket.emit('error', { error: String(error) });
});
processMessages().catch(error => {
console.error('Error in process messages task:', error);
socket.emit('error', { error: String(error) });
});
// Handle disconnect - cleanup resources
socket.on('disconnect', () => {
console.log(`WebSocket client disconnected: ${socket.id}`);
// Any cleanup needed for the live request queue
});
}
catch (error) {
console.error('Error in run_live setup:', error);
socket.emit('error', { error: String(error) });
}
});
});
// Add more endpoints to match Python's fast_api.py functionality
// Add session to eval set endpoint
app.post('/apps/:appName/eval_sets/:evalSetId/add_session', async (req, res) => {
const { appName, evalSetId } = req.params;
const requestData = req.body;
// Validate eval ID format
const pattern = /^[a-zA-Z0-9_]+$/;
if (!pattern.test(requestData.evalId)) {
return res.status(400).json({
error: `Invalid eval id. Eval id should have the \`${pattern}\` format`
});
}
// Get the session
const session = sessionService.getSession({
appName,
userId: requestData.userId,
sessionId: requestData.sessionId
});
if (!session) {
return res.status(404).json({ error: 'Session not found' });
}
// Load the eval set file data
const evalSetFilePath = getEvalSetFilePath(appName, evalSetId);
if (!fs_1.default.existsSync(evalSetFilePath)) {
return res.status(404).json({ error: 'Eval set not found' });
}
try {
const evalSetDataRaw = fs_1.default.readFileSync(evalSetFilePath, 'utf-8');
const evalSetData = JSON.parse(evalSetDataRaw);
// Check if eval ID already exists
if (evalSetData.some((item) => item.name === requestData.evalId)) {
return res.status(400).json({
error: `Eval id \`${requestData.evalId}\` already exists in \`${evalSetId}\` eval set.`
});
}
// Convert session to evaluation format
const testData = (0, utils_1.convertSessionToEvalFormat)(session);
// Get root agent for initial session state
const rootAgent = await getRootAgent(appName);
const initialSessionState = (0, utils_1.createEmptyState)(rootAgent);
// Add to eval set
evalSetData.push({
name: requestData.evalId,
data: testData,
initial_session: {
state: initialSessionState,
app_name: appName,
user_id: requestData.userId
}
});
// Write updated eval set back to file
fs_1.default.writeFileSync(evalSetFilePath, JSON.stringify(evalSetData, null, 2));
res.status(201).json({ status: 'success' });
}
catch (error) {
console.error('Error adding session to eval set:', error);
res.status(500).json({ error: 'Failed to add session to eval set' });
}
});
// List evals in eval set endpoint
app.get('/apps/:appName/eval_sets/:evalSetId/evals', (req, res) => {
const { appName, evalSetId } = req.params;
// Load the eval set file data
const evalSetFilePath = getEvalSetFilePath(appName, evalSetId);
if (!fs_1.default.existsSync(evalSetFilePath)) {
return res.status(404).json({ error: 'Eval set not found' });
}
try {
const evalSetDataRaw = fs_1.default.readFileSync(evalSetFilePath, 'utf-8');
const evalSetData = JSON.parse(evalSetDataRaw);
const evalIds = evalSetData.map((item) => item.name);
res.json(evalIds.sort());
}
catch (error) {
console.error('Error listing evals:', error);
res.status(500).json({ error: 'Failed to list evals' });
}
});
// Run eval endpoint
app.post('/apps/:appName/eval_sets/:evalSetId/run_eval', async (req, res) => {
const { appName, evalSetId } = req.params;
const requestData = req.body;
// This is a placeholder implementation as the full eval runner would require more code
// In a complete implementation, this would use the cliEval module
try {
const evalSetFilePath = getEvalSetFilePath(appName, evalSetId);
if (!fs_1.default.existsSync(evalSetFilePath)) {
return res.status(404).json({ error: 'Eval set not found' });
}
// Get root agent
const rootAgent = await getRootAgent(appName);
// Mock response - in a real implementation, this would run the actual evaluations
const results = requestData.evalIds.map(evalId => ({
evalSetId,
evalId,
finalEvalStatus: 'SUCCESS',
evalMetricResults: [],
sessionId: `eval_${evalId}_${Date.now()}`
}));
res.json(results);
}
catch (error) {
console.error('Error running eval:', error);
res.status(500).json({ error: 'Failed to run eval' });
}
});
// Event graph endpoint
app.get('/apps/:appName/users/:userId/sessions/:sessionId/events/:eventId/graph', async (req, res) => {
const { appName, userId, sessionId, eventId } = req.params;
// Connect to managed session if agent_engine_id is set
const effectiveAppName = agentEngineId || appName;
const session = sessionService.getSession({
appName: effectiveAppName,
userId,
sessionId
});
if (!session || !session.events) {
return res.json({});
}
const event = session.events.find((e) => e.id === eventId);
if (!event) {
return res.json({});
}
try {
const rootAgent = await getRootAgent(appName);
let dotGraph = null;
// Check for function calls
const functionCalls = event.getFunctionCalls ? event.getFunctionCalls() : [];
const functionResponses = event.getFunctionResponses ? event.getFunctionResponses() : [];
if (functionCalls && functionCalls.length > 0) {
const functionCallHighlights = functionCalls.map((call) => [event.author, call.name]);
dotGraph = agentGraph.getAgentGraph(rootAgent, functionCallHighlights);
}
else if (functionResponses && functionResponses.length > 0) {
const functionResponseHighlights = functionResponses.map((response) => [response.name, event.author]);
dotGraph = agentGraph.getAgentGraph(rootAgent, functionResponseHighlights);
}
else {
dotGraph = agentGraph.getAgentGraph(rootAgent, [[event.author, '']]);
}
if (dotGraph) {
// For TypeScript, we'll return the dot source directly
return res.json({ dot_src: dotGraph.to_string() });
}
else {
return res.json({});
}
}
catch (error) {
console.error('Error generating agent graph:', error);
return res.status(500).json({ error: 'Error generating agent graph' });
}
});
// Helper functions for eval set management
function getEvalSetFilePath(appName, evalSetId) {
return path_1.default.join(agentDir, appName, evalSetId + EVAL_SET_FILE_EXTENSION);
}
// Eval set endpoints
app.post('/apps/:appName/eval_sets/:evalSetId', (req, res) => {
const { appName, evalSetId } = req.params;
// Validate eval set ID
const pattern = /^[a-zA-Z0-9_]+$/;
if (!pattern.test(evalSetId)) {
return res.status(400).json({
error: `Invalid eval set id. Eval set id should have the ${pattern} format`
});
}
// Define the file path
const newEvalSetPath = getEvalSetFilePath(appName, evalSetId);
console.log(`Creating eval set file ${newEvalSetPath}`);
if (!fs_1.default.existsSync(newEvalSetPath)) {
// Write the JSON string to the file
console.log("Eval set file doesn't exist, we will create a new one.");
fs_1.default.writeFileSync(newEvalSetPath, JSON.stringify([], null, 2));
}
res.status(201).send();
});
app.get('/apps/:appName/eval_sets', (req, res) => {
const { appName } = req.params;
const evalSetFilePath = path_1.default.join(agentDir, appName);
if (!fs_1.default.existsSync(evalSetFilePath)) {
return res.status(404).json({ error: 'App directory not found' });
}
const evalSets = fs_1.default.readdirSync(evalSetFilePath)
.filter(file => file.endsWith(EVAL_SET_FILE_EXTENSION))
.map(file => path_1.default.basename(file).replace(EVAL_SET_FILE_EXTENSION, ''));
res.json(evalSets.sort());
});
/**
* Helper function to get the root agent for an app
*/
async function getRootAgent(appName) {
if (rootAgentDict[appName]) {
return rootAgentDict[appName];
}
try {
// Dynamically import the agent module
const agentModule = require(path_1.default.join(agentDir, appName));
if (!agentModule.agent?.rootAgent) {
throw new Error(`Unable to find "rootAgent" from ${appName}.`);
}
const rootAgent = agentModule.agent.rootAgent;
rootAgentDict[appName] = rootAgent;
return rootAgent;
}
catch (error) {
console.error(`Error loading root agent for ${appName}:`, error);
throw error;
}
}
/**
* Helper function to get a runner for an app
*/
async function getRunner(appName) {
if (runnerDict[appName]) {
return runnerDict[appName];
}
// Load environment variables for the agent
(0, utils_1.loadDotenvForAgent)('', agentDir);
const rootAgent = await getRootAgent(appName);
const runner = new runners_1.Runner({
appName: agentEngineId || appName,
agent: rootAgent,
artifactService,
sessionService
});
runnerDict[appName] = runner;
return runner;
}
// If web UI is enabled, serve static files
if (web) {
const BASE_DIR = path_1.default.dirname(__filename);
const ANGULAR_DIST_PATH = path_1.default.join(BASE_DIR, 'browser');
// Root redirect to dev-ui
app.get('/', (req, res) => {
res.redirect('/dev-ui');
});
// Dev UI page - serve the index.html file directly
app.get('/dev-ui', (req, res) => {
const indexHtmlPath = path_1.default.join(ANGULAR_DIST_PATH, 'index.html');
if (fs_1.default.existsSync(indexHtmlPath)) {
res.sendFile(indexHtmlPath);
}
else {
// Fallback if index.html not found
res.status(404).send('Web UI not found. Please make sure the browser directory is properly installed.');
}
});
// Serve all static files from the browser directory
app.use('/', express_1.default.static(ANGULAR_DIST_PATH));
console.log(`Serving web UI from ${ANGULAR_DIST_PATH}`);
}
// Start the server if a port was provided
if (port) {
server.listen(port, () => {
console.log(`API server running at http://localhost:${port}`);
});
}
// Return both the app and server
return { app, server };
}