adk-typescript/dist/runners.js

TypeScript port of Google's Agent Development Kit (ADK)

"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.InMemoryRunner = exports.Runner = void 0;
const uuid_1 = require("uuid");
const InvocationContext_1 = require("./agents/InvocationContext");
const LlmAgent_1 = require("./agents/LlmAgent");
const LiveRequestQueue_1 = require("./agents/LiveRequestQueue");
const RunConfig_1 = require("./agents/RunConfig");
const InMemoryArtifactService_1 = require("./artifacts/InMemoryArtifactService");
const Event_1 = require("./events/Event");
const InMemoryMemoryService_1 = require("./memory/InMemoryMemoryService");
const InMemorySessionService_1 = require("./sessions/InMemorySessionService");
const telemetry_1 = require("./telemetry");
const BuiltInCodeExecutionTool_1 = require("./tools/BuiltInCodeExecutionTool");
// Logger is a placeholder - implement with proper logging library
const logger = {
    warning: (message, ...args) => console.warn(message, ...args),
    info: (message, ...args) => console.info(message, ...args),
    error: (message, ...args) => console.error(message, ...args)
};
/**
 * The Runner class is used to run agents.
 *
 * It manages the execution of an agent within a session, handling message
 * processing, event generation, and interaction with various services like
 * artifact storage, session management, and memory.
 */
class Runner {
    /**
     * Initializes the Runner.
     *
     * @param params The parameters for the runner.
     * @param params.appName The application name of the runner.
     * @param params.agent The root agent to run.
     * @param params.artifactService The artifact service for the runner.
     * @param params.sessionService The session service for the runner.
     * @param params.memoryService The memory service for the runner.
     */
    constructor(params) {
        this.appName = params.appName;
        this.agent = params.agent;
        this.artifactService = params.artifactService;
        this.sessionService = params.sessionService;
        this.memoryService = params.memoryService;
    }
    /**
     * Runs the agent.
     *
     * NOTE: This sync interface is only for local testing and convenience purpose.
     * Consider using `runAsync` for production usage.
     *
     * @param params The parameters for the run.
     * @param params.userId The user ID of the session.
     * @param params.sessionId The session ID of the session.
     * @param params.newMessage A new message to append to the session.
     * @param params.runConfig The run config for the agent.
     * @returns A generator that yields the events generated by the agent.
     */
    async *run(params) {
        const { userId, sessionId, newMessage, runConfig = new RunConfig_1.RunConfig() } = params;
        // In JavaScript, we can just use the async generator directly
        for await (const event of this.runAsync({
            userId,
            sessionId,
            newMessage,
            runConfig
        })) {
            yield event;
        }
    }
    /**
     * Main entry method to run the agent in this runner.
     *
     * @param params The parameters for the run.
     * @param params.userId The user ID of the session.
     * @param params.sessionId The session ID of the session.
     * @param params.newMessage A new message to append to the session.
     * @param params.runConfig The run config for the agent.
     * @returns An async generator that yields events generated by the agent.
     */
    async *runAsync(params) {
        const { userId, sessionId, newMessage, runConfig = new RunConfig_1.RunConfig() } = params;
        const span = telemetry_1.tracer.startAsCurrentSpan('invocation');
        try {
            const session = await this.sessionService.getSession({
                appName: this.appName,
                userId: userId,
                sessionId: sessionId
            });
            if (!session) {
                throw new Error(`Session not found: ${sessionId}`);
            }
            const invocationContext = this._newInvocationContext({
                session: session,
                newMessage,
                runConfig
            });
            const rootAgent = this.agent;
            if (newMessage) {
                await this._appendNewMessageToSession({
                    session: session,
                    newMessage,
                    invocationContext,
                    saveInputBlobsAsArtifacts: runConfig.saveInputBlobsAsArtifacts
                });
            }
            invocationContext.agent = this._findAgentToRun(session, rootAgent);
            // Use invoke method which is guaranteed to exist
            for await (const event of invocationContext.agent.invoke(invocationContext)) {
                if (!event.partial) {
                    await this.sessionService.appendEvent({
                        session: session,
                        event: event
                    });
                }
                yield event;
            }
        }
        finally {
            span.end();
        }
    }
    /**
     * Appends a new message to a session.
     *
     * @param params The parameters for appending a message.
     * @param params.session The session to append to.
     * @param params.newMessage The message to append.
     * @param params.invocationContext The invocation context.
     * @param params.saveInputBlobsAsArtifacts Whether to save input blobs as artifacts.
     * @private
     */
    async _appendNewMessageToSession(params) {
        const { session, newMessage, invocationContext, saveInputBlobsAsArtifacts = false } = params;
        if (this.artifactService && saveInputBlobsAsArtifacts) {
            // The runner directly saves the artifacts (if applicable) in the
            // user message and replaces the artifact data with a file name
            // placeholder.
            for (let i = 0; i < newMessage.parts.length; i++) {
                const part = newMessage.parts[i];
                if (!part.inlineData) {
                    continue;
                }
                const fileName = `artifact_${invocationContext.invocationId}_${i}`;
                try {
                    const saveResult = this.artifactService.saveArtifact({
                        appName: this.appName,
                        userId: session.userId,
                        sessionId: session.id,
                        filename: fileName,
                        artifact: part
                    });
                    // Handle both synchronous and asynchronous cases
                    if (saveResult instanceof Promise) {
                        await saveResult;
                    }
                    newMessage.parts[i] = {
                        text: `Uploaded file: ${fileName}. It is saved into artifacts`
                    };
                }
                catch (error) {
                    console.error('Error saving artifact:', error);
                    // Continue with the message even if saving artifact fails
                }
            }
        }
        // Appends only. We do not yield the event because it's not from the model.
        const event = new Event_1.Event({
            invocationId: invocationContext.invocationId,
            author: 'user',
            content: newMessage
        });
        await this.sessionService.appendEvent({
            session: session,
            event: event
        });
    }
    /**
     * Runs the agent in live mode.
     *
     * @param params The parameters for live mode.
     * @param params.session The session to use.
     * @param params.liveRequestQueue The queue for live requests.
     * @param params.runConfig The run config for the agent.
     * @returns An async generator of events.
     *
     * @experimental This feature is **experimental** and its API or behavior may change
     *               in future releases.
     */
    async *runLive(params) {
        const { session, liveRequestQueue, runConfig = new RunConfig_1.RunConfig() } = params;
        // TODO: right now, only works for a single audio agent without FC.
        const invocationContext = this._newInvocationContextForLive({
            session: session,
            liveRequestQueue,
            runConfig
        });
        const rootAgent = this.agent;
        invocationContext.agent = this._findAgentToRun(session, rootAgent);
        invocationContext.activeStreamingTools = new Map();
        // Get tools from the agent, assuming it has a canonicalTools property
        const tools = invocationContext.agent.canonicalTools || [];
        // Initialize streaming tools that rely on LiveRequestQueue
        for (const tool of tools) {
            // Check if tool requires LiveRequestQueue
            // This is a bit complex to port directly - in TypeScript we would need
            // to use reflection or decorator metadata to check parameter types
            if (tool.usesLiveQueue) {
                if (!invocationContext.activeStreamingTools) {
                    invocationContext.activeStreamingTools = new Map();
                }
                // Create an object with appropriate structure for ActiveStreamingTool
                // Based on inspected constructor, ActiveStreamingTool expects name, args, id
                const streamObj = {
                    name: tool.name,
                    args: {},
                    id: (0, uuid_1.v4)(),
                    stream: new LiveRequestQueue_1.LiveRequestQueue()
                };
                // Add to map
                invocationContext.activeStreamingTools.set(tool.name, 
                // Using any type assertion since we may not have the correct constructor
                // signature but we know the object structure should work
                streamObj);
            }
        }
        // Set invocation context to live mode
        invocationContext.live = true;
        // Use invoke method which will call the appropriate implementation based on live flag
        for await (const event of invocationContext.agent.invoke(invocationContext)) {
            await this.sessionService.appendEvent({
                session: session,
                event: event
            });
            yield event;
        }
    }
    /**
     * Closes a session and adds it to the memory service.
     *
     * @param session The session to close.
     * @experimental This feature is **experimental** and its API or behavior may change
     *               in future releases.
     */
    async closeSession(session) {
        if (this.memoryService) {
            await this.memoryService.addSessionToMemory(session);
        }
        await this.sessionService.closeSession({
            session: session
        });
    }
    /**
     * Finds the agent to run to continue the session.
     *
     * A qualified agent must be either of:
     * - The root agent;
     * - An LlmAgent who replied last and is capable to transfer to any other agent
     *   in the agent hierarchy.
     *
     * @param session The session to find the agent for.
     * @param rootAgent The root agent of the runner.
     * @returns The agent of the last message in the session or the root agent.
     * @private
     */
    _findAgentToRun(session, rootAgent) {
        // Make sure session has required properties
        if (!session.agents) {
            session.agents = new Map();
        }
        // Make sure the root agent is registered in the session
        if (!session.agents.has(rootAgent.name)) {
            session.agents.set(rootAgent.name, rootAgent);
        }
        // Add getAgent method if not present
        if (!session.getAgent) {
            session.getAgent = function (name) {
                return this.agents.get(name);
            };
        }
        // Filter for non-user events and process them in reverse order
        const nonUserEvents = session.events.filter(e => e.author !== 'user');
        for (let i = nonUserEvents.length - 1; i >= 0; i--) {
            const event = nonUserEvents[i];
            if (event.author === rootAgent.name) {
                // Found root agent
                return rootAgent;
            }
            const agent = rootAgent.findSubAgent(event.author);
            if (!agent) {
                // Agent not found, continue looking
                logger.warning(`Event from an unknown agent: ${event.author}, event id: ${event.id}`);
                continue;
            }
            if (this._isTransferableAcrossAgentTree(agent)) {
                return agent;
            }
        }
        // Falls back to root agent if no suitable agents are found in the session
        return rootAgent;
    }
    /**
     * Whether the agent to run can transfer to any other agent in the agent tree.
     *
     * This typically means all agent_to_run's parent through root agent can
     * transfer to their parent_agent.
     *
     * @param agentToRun The agent to check for transferability.
     * @returns True if the agent can transfer, False otherwise.
     * @private
     */
    _isTransferableAcrossAgentTree(agentToRun) {
        let agent = agentToRun;
        while (agent) {
            if (!(agent instanceof LlmAgent_1.LlmAgent)) {
                // Only LLM-based Agent can provide agent transfer capability
                return false;
            }
            const llmAgent = agent;
            // Check if transfers to peers are allowed - in TypeScript we use disallowTransferToPeers
            if (llmAgent.disallowTransferToPeers === true) {
                return false;
            }
            agent = agent.parentAgent;
        }
        return true;
    }
    /**
     * Creates a new invocation context.
     *
     * @param params The parameters for the invocation context.
     * @param params.session The session for the context.
     * @param params.newMessage The new message for the context.
     * @param params.liveRequestQueue The live request queue for the context.
     * @param params.runConfig The run config for the context.
     * @returns The new invocation context.
     * @private
     */
    _newInvocationContext(params) {
        const { session, newMessage, liveRequestQueue, runConfig = new RunConfig_1.RunConfig() } = params;
        const invocationId = (0, uuid_1.v4)();
        // Get the LLM model to use
        let llm;
        if (this.agent instanceof LlmAgent_1.LlmAgent) {
            try {
                llm = this.agent.canonicalModel;
            }
            catch (error) {
                console.error('Error getting canonicalModel:', error);
            }
        }
        if (runConfig.supportCfc && this.agent instanceof LlmAgent_1.LlmAgent) {
            // Get the agent's model name using canonicalModel
            const llmAgent = this.agent;
            const modelName = llmAgent.canonicalModel.model || 'unknown';
            if (!modelName.startsWith('gemini-2')) {
                throw new Error(`CFC is not supported for model: ${modelName} in agent: ${this.agent.name}`);
            }
            // Check if built-in code execution tool is already included
            const tools = llmAgent.canonicalTools || [];
            const hasCodeExecutionTool = tools.some((tool) => tool instanceof BuiltInCodeExecutionTool_1.BuiltInCodeExecutionTool);
            if (!hasCodeExecutionTool) {
                llmAgent.canonicalTools.push(new BuiltInCodeExecutionTool_1.BuiltInCodeExecutionTool());
            }
        }
        // Create context with the model set
        const context = new InvocationContext_1.InvocationContext({
            artifactService: this.artifactService,
            sessionService: this.sessionService,
            memoryService: this.memoryService,
            invocationId: invocationId,
            agent: this.agent,
            session: session,
            userContent: newMessage,
            liveRequestQueue: liveRequestQueue,
            runConfig: runConfig,
            llm: llm
        });
        // Debug log if model is missing
        if (!context.llm) {
            console.warn(`No LLM model set in invocation context for agent: ${this.agent.name}`);
        }
        return context;
    }
    /**
     * Creates a new invocation context for live multi-agent.
     *
     * @param params The parameters for the invocation context.
     * @param params.session The session for the context.
     * @param params.liveRequestQueue The live request queue for the context.
     * @param params.runConfig The run config for the context.
     * @returns The new invocation context.
     * @private
     */
    _newInvocationContextForLive(params) {
        const { session, liveRequestQueue, runConfig = new RunConfig_1.RunConfig() } = params;
        // For live multi-agent, we need model's text transcription as context for
        // next agent
        if (this.agent.subAgents && this.agent.subAgents.length > 0 && liveRequestQueue) {
            if (!runConfig.responseModalities || runConfig.responseModalities.length === 0) {
                // default
                runConfig.responseModalities = ['AUDIO'];
                if (!runConfig.outputAudioTranscription) {
                    runConfig.outputAudioTranscription = {};
                }
            }
            else if (!runConfig.responseModalities.includes('TEXT')) {
                if (!runConfig.outputAudioTranscription) {
                    runConfig.outputAudioTranscription = {};
                }
            }
        }
        return this._newInvocationContext({
            session,
            liveRequestQueue,
            runConfig
        });
    }
}
exports.Runner = Runner;
/**
 * An in-memory Runner for testing and development.
 *
 * This runner uses in-memory implementations for artifact, session, and memory
 * services, providing a lightweight and self-contained environment for agent
 * execution.
 */
class InMemoryRunner extends Runner {
    /**
     * Initializes the InMemoryRunner.
     *
     * @param agent The root agent to run.
     * @param appName The application name of the runner. Defaults to 'InMemoryRunner'.
     */
    constructor(agent, appName = 'InMemoryRunner') {
        super({
            appName: appName,
            agent: agent,
            artifactService: new InMemoryArtifactService_1.InMemoryArtifactService(),
            sessionService: new InMemorySessionService_1.InMemorySessionService(),
            memoryService: new InMemoryMemoryService_1.InMemoryMemoryService()
        });
    }
}
exports.InMemoryRunner = InMemoryRunner;