@jupyterlite/ai 0.8.1 → 0.9.0-a1

package/lib/agent.js

@@ -0,0 +1,627 @@
+import { Signal } from '@lumino/signaling';
+import { Agent, Runner, user } from '@openai/agents';
+import { BrowserMCPServerStreamableHttp } from './mcp/browser';
+import { createModel } from './providers/models';
+import { SECRETS_NAMESPACE } from './tokens';
+/**
+ * Manages the AI agent lifecycle and execution loop.
+ * Provides agent initialization, tool management, MCP server integration,
+ * and handles the complete agent execution cycle including tool approvals.
+ * Emits events for UI updates instead of directly manipulating the chat interface.
+ */
+export class AgentManager {
+    /**
+     * Creates a new AgentManager instance.
+     * @param options Configuration options for the agent manager
+     */
+    constructor(options) {
+        Private.setToken(options.token);
+        this._settingsModel = options.settingsModel;
+        this._toolRegistry = options.toolRegistry;
+        this._chatProviderRegistry = options.chatProviderRegistry;
+        this._secretsManager = options.secretsManager;
+        this._selectedToolNames = [];
+        this._agent = null;
+        this._runner = new Runner({ tracingDisabled: true });
+        this._history = [];
+        this._mcpServers = [];
+        this._isInitializing = false;
+        this._controller = null;
+        this._pendingApprovals = new Map();
+        this._interruptedState = null;
+        this._agentEvent = new Signal(this);
+        this._mcpConnectionChanged = new Signal(this);
+        this._tokenUsage = { inputTokens: 0, outputTokens: 0 };
+        this._tokenUsageChanged = new Signal(this);
+        // Initialize selected tools to all available tools by default
+        if (this._toolRegistry) {
+            this._selectedToolNames = Object.keys(this._toolRegistry.tools);
+        }
+        // Initialize agent on construction
+        this._initializeAgent().catch(error => console.warn('Failed to initialize agent in constructor:', error));
+        // Listen for settings changes
+        this._settingsModel.stateChanged.connect(this._onSettingsChanged, this);
+    }
+    /**
+     * Signal emitted when agent events occur
+     */
+    get agentEvent() {
+        return this._agentEvent;
+    }
+    /**
+     * Signal emitted when MCP connection status changes
+     */
+    get mcpConnectionChanged() {
+        return this._mcpConnectionChanged;
+    }
+    /**
+     * Gets the current token usage statistics.
+     */
+    get tokenUsage() {
+        return this._tokenUsage;
+    }
+    /**
+     * Signal emitted when token usage statistics change.
+     */
+    get tokenUsageChanged() {
+        return this._tokenUsageChanged;
+    }
+    /**
+     * Sets the selected tools by name and reinitializes the agent.
+     * @param toolNames Array of tool names to select
+     */
+    setSelectedTools(toolNames) {
+        this._selectedToolNames = [...toolNames];
+        this._initializeAgent().catch(error => console.warn('Failed to initialize agent on tools change:', error));
+    }
+    /**
+     * Gets the currently selected tools as OpenAI agents tools.
+     * @returns Array of selected tools formatted for OpenAI agents
+     */
+    get selectedAgentTools() {
+        if (!this._toolRegistry) {
+            return [];
+        }
+        const result = [];
+        for (const name of this._selectedToolNames) {
+            const tool = this._toolRegistry.get(name);
+            if (tool) {
+                result.push(tool);
+            }
+        }
+        return result;
+    }
+    /**
+     * Checks if a specific MCP server is connected by server name.
+     * @param serverName The name of the MCP server to check
+     * @returns True if the server is connected, false otherwise
+     */
+    isMCPServerConnected(serverName) {
+        return this._mcpServers.some(server => server.name === serverName);
+    }
+    /**
+     * Checks if the current configuration is valid for agent operations.
+     * Uses the provider registry to determine if an API key is required.
+     * @returns True if the configuration is valid, false otherwise
+     */
+    hasValidConfig() {
+        const activeProvider = this._settingsModel.getActiveProvider();
+        if (!activeProvider) {
+            return false;
+        }
+        if (!activeProvider.model) {
+            return false;
+        }
+        if (this._chatProviderRegistry) {
+            const providerInfo = this._chatProviderRegistry.getProviderInfo(activeProvider.provider);
+            if (providerInfo?.requiresApiKey) {
+                return !!activeProvider.apiKey;
+            }
+        }
+        return true;
+    }
+    /**
+     * Clears conversation history and resets agent state.
+     * Removes all conversation history, pending approvals, and interrupted state.
+     */
+    clearHistory() {
+        this._history = [];
+        this._runner = new Runner();
+        this._pendingApprovals.clear();
+        this._interruptedState = null;
+        // Reset token usage
+        this._tokenUsage = { inputTokens: 0, outputTokens: 0 };
+        this._tokenUsageChanged.emit(this._tokenUsage);
+    }
+    /**
+     * Stops the current streaming response by aborting the request.
+     */
+    stopStreaming() {
+        this._controller?.abort();
+    }
+    /**
+     * Generates AI response to user message using the agent.
+     * Handles the complete execution cycle including tool calls and approvals.
+     * @param message The user message to respond to (may include processed attachment content)
+     */
+    async generateResponse(message) {
+        const config = this._settingsModel.config;
+        this._controller = new AbortController();
+        try {
+            // Ensure we have an agent
+            if (!this._agent) {
+                await this._initializeAgent();
+            }
+            if (!this._agent) {
+                throw new Error('Failed to initialize agent');
+            }
+            const shouldUseTools = config.toolsEnabled &&
+                this._selectedToolNames.length > 0 &&
+                this._toolRegistry &&
+                Object.keys(this._toolRegistry.tools).length > 0 &&
+                this._supportsToolCalling();
+            // Add user message to history
+            this._history.push(user(message));
+            // Main agentic loop
+            let result = await this._runner.run(this._agent, this._history, {
+                stream: true,
+                signal: this._controller.signal,
+                ...(shouldUseTools && { maxTurns: config.maxTurns })
+            });
+            await this._processRunResult(result);
+            let hasInterruptions = result.interruptions && result.interruptions.length > 0;
+            while (hasInterruptions) {
+                this._interruptedState = result;
+                const interruptions = result.interruptions;
+                if (interruptions.length > 1) {
+                    await this._handleGroupedToolApprovals(interruptions);
+                }
+                else {
+                    await this._handleSingleToolApproval(interruptions[0]);
+                }
+                // Wait for all approvals to be resolved
+                while (this._pendingApprovals.size > 0) {
+                    await new Promise(resolve => setTimeout(resolve, 100));
+                }
+                // Continue execution
+                result = await this._runner.run(this._agent, result.state, {
+                    stream: true,
+                    signal: this._controller.signal,
+                    maxTurns: config.maxTurns
+                });
+                await this._processRunResult(result);
+                hasInterruptions =
+                    result.interruptions && result.interruptions.length > 0;
+            }
+            // Clear interrupted state
+            this._interruptedState = null;
+            this._history = result.history;
+        }
+        catch (error) {
+            this._agentEvent.emit({
+                type: 'error',
+                data: { error: error }
+            });
+        }
+        finally {
+            this._controller = null;
+        }
+    }
+    /**
+     * Approves a tool call by interruption ID.
+     * @param interruptionId The interruption ID to approve
+     */
+    async approveToolCall(interruptionId) {
+        const pending = this._pendingApprovals.get(interruptionId);
+        if (!pending) {
+            console.warn(`No pending approval found for interruption ${interruptionId}`);
+            return;
+        }
+        if (this._interruptedState) {
+            this._interruptedState.state.approve(pending.interruption);
+        }
+        pending.approved = true;
+        this._pendingApprovals.delete(interruptionId);
+    }
+    /**
+     * Rejects a tool call by interruption ID.
+     * @param interruptionId The interruption ID to reject
+     */
+    async rejectToolCall(interruptionId) {
+        const pending = this._pendingApprovals.get(interruptionId);
+        if (!pending) {
+            console.warn(`No pending approval found for interruption ${interruptionId}`);
+            return;
+        }
+        if (this._interruptedState) {
+            this._interruptedState.state.reject(pending.interruption);
+        }
+        pending.approved = false;
+        this._pendingApprovals.delete(interruptionId);
+    }
+    /**
+     * Approves all tools in a group by group ID.
+     * @param groupId The group ID containing the tool calls
+     * @param interruptionIds Array of interruption IDs to approve
+     */
+    async approveGroupedToolCalls(groupId, interruptionIds) {
+        for (const interruptionId of interruptionIds) {
+            const pending = this._pendingApprovals.get(interruptionId);
+            if (pending && pending.groupId === groupId) {
+                if (this._interruptedState) {
+                    this._interruptedState.state.approve(pending.interruption);
+                }
+                pending.approved = true;
+                this._pendingApprovals.delete(interruptionId);
+            }
+        }
+    }
+    /**
+     * Rejects all tools in a group by group ID.
+     * @param groupId The group ID containing the tool calls
+     * @param interruptionIds Array of interruption IDs to reject
+     */
+    async rejectGroupedToolCalls(groupId, interruptionIds) {
+        for (const interruptionId of interruptionIds) {
+            const pending = this._pendingApprovals.get(interruptionId);
+            if (pending && pending.groupId === groupId) {
+                if (this._interruptedState) {
+                    this._interruptedState.state.reject(pending.interruption);
+                }
+                pending.approved = false;
+                this._pendingApprovals.delete(interruptionId);
+            }
+        }
+    }
+    /**
+     * Handles settings changes and reinitializes the agent.
+     */
+    _onSettingsChanged() {
+        this._initializeAgent().catch(error => console.warn('Failed to initialize agent on settings change:', error));
+    }
+    /**
+     * Initializes the AI agent with current settings and tools.
+     * Sets up the agent with model configuration, tools, and MCP servers.
+     */
+    async _initializeAgent() {
+        if (this._isInitializing) {
+            return;
+        }
+        this._isInitializing = true;
+        try {
+            const config = this._settingsModel.config;
+            const model = await this._createModel();
+            const shouldUseTools = config.toolsEnabled &&
+                this._selectedToolNames.length > 0 &&
+                this._toolRegistry &&
+                Object.keys(this._toolRegistry.tools).length > 0 &&
+                this._supportsToolCalling();
+            await this._initializeMCPServers();
+            const tools = shouldUseTools ? this.selectedAgentTools : [];
+            const mcpServers = this._mcpServers.filter(server => server !== null);
+            this._agent = new Agent({
+                name: 'Assistant',
+                instructions: shouldUseTools
+                    ? this._getEnhancedSystemPrompt(config.systemPrompt || '')
+                    : config.systemPrompt || 'You are a helpful assistant.',
+                model: model,
+                mcpServers,
+                tools,
+                ...(config.temperature && {
+                    modelSettings: {
+                        temperature: config.temperature,
+                        ...(config.maxTokens && { maxTokens: config.maxTokens })
+                    }
+                })
+            });
+        }
+        catch (error) {
+            console.warn('Failed to initialize agent:', error);
+            this._agent = null;
+        }
+        finally {
+            this._isInitializing = false;
+        }
+    }
+    /**
+     * Initializes MCP (Model Context Protocol) servers based on current settings.
+     * Closes existing servers and connects to enabled servers from configuration.
+     */
+    async _initializeMCPServers() {
+        const config = this._settingsModel.config;
+        const enabledServers = config.mcpServers.filter(server => server.enabled);
+        let connectionChanged = false;
+        // Close existing servers
+        for (const server of this._mcpServers) {
+            try {
+                await server.close();
+                connectionChanged = true;
+            }
+            catch (error) {
+                console.warn('Error closing MCP server:', error);
+            }
+        }
+        this._mcpServers = [];
+        // Initialize new servers
+        for (const serverConfig of enabledServers) {
+            try {
+                const mcpServer = new BrowserMCPServerStreamableHttp({
+                    url: serverConfig.url,
+                    name: serverConfig.name
+                });
+                await mcpServer.connect();
+                this._mcpServers.push(mcpServer);
+                connectionChanged = true;
+            }
+            catch (error) {
+                console.warn(`Failed to connect to MCP server "${serverConfig.name}" at ${serverConfig.url}:`, error);
+            }
+        }
+        // Emit connection change signal if there were any changes
+        if (connectionChanged) {
+            this._mcpConnectionChanged.emit(this._mcpServers.length > 0);
+        }
+    }
+    /**
+     * Processes the result stream from agent execution.
+     * Handles message streaming, tool calls, and emits appropriate events.
+     * @param result The async iterable result from agent execution
+     */
+    async _processRunResult(result) {
+        let fullResponse = '';
+        let currentMessageId = null;
+        for await (const event of result) {
+            if (event.type === 'raw_model_stream_event') {
+                const data = event.data;
+                if (data.type === 'response_started') {
+                    currentMessageId = `msg-${Date.now()}-${Math.random()}`;
+                    fullResponse = '';
+                    this._agentEvent.emit({
+                        type: 'message_start',
+                        data: { messageId: currentMessageId }
+                    });
+                }
+                else if (data.type === 'output_text_delta') {
+                    if (currentMessageId) {
+                        const chunk = data.delta || '';
+                        fullResponse += chunk;
+                        this._agentEvent.emit({
+                            type: 'message_chunk',
+                            data: {
+                                messageId: currentMessageId,
+                                chunk,
+                                fullContent: fullResponse
+                            }
+                        });
+                    }
+                }
+                else if (data.type === 'response_done') {
+                    if (currentMessageId) {
+                        this._agentEvent.emit({
+                            type: 'message_complete',
+                            data: {
+                                messageId: currentMessageId,
+                                content: fullResponse
+                            }
+                        });
+                        currentMessageId = null;
+                    }
+                }
+                else if (data.type === 'model') {
+                    const modelEvent = data.event;
+                    if (modelEvent.type === 'tool-call') {
+                        this._handleToolCallStart(modelEvent);
+                    }
+                }
+            }
+            else if (event.type === 'run_item_stream_event') {
+                if (event.name === 'tool_output') {
+                    this._handleToolOutput(event);
+                }
+            }
+        }
+    }
+    /**
+     * Handles the start of a tool call from the model event.
+     * @param modelEvent The model event containing tool call information
+     */
+    _handleToolCallStart(modelEvent) {
+        const toolCallId = modelEvent.toolCallId;
+        const toolName = modelEvent.toolName;
+        const toolInput = modelEvent.input;
+        let parsedToolInput;
+        try {
+            parsedToolInput = JSON.parse(toolInput);
+        }
+        catch (error) {
+            parsedToolInput = {};
+        }
+        this._agentEvent.emit({
+            type: 'tool_call_start',
+            data: {
+                callId: toolCallId,
+                toolName,
+                input: parsedToolInput
+            }
+        });
+    }
+    /**
+     * Handles tool execution output and completion.
+     * @param event The tool output event containing result information
+     */
+    _handleToolOutput(event) {
+        const toolEvent = event;
+        const toolCallOutput = toolEvent.item;
+        const callId = toolCallOutput.rawItem.callId;
+        const resultText = typeof toolCallOutput.output === 'string'
+            ? toolCallOutput.output
+            : JSON.stringify(toolCallOutput.output, null, 2);
+        const isError = toolCallOutput.rawItem.type === 'function_call_result' &&
+            toolCallOutput.rawItem.error;
+        const toolName = toolCallOutput.rawItem.type === 'function_call_result'
+            ? toolCallOutput.rawItem.name
+            : 'Unknown Tool';
+        this._agentEvent.emit({
+            type: 'tool_call_complete',
+            data: {
+                callId,
+                toolName,
+                output: resultText,
+                isError
+            }
+        });
+    }
+    /**
+     * Handles approval request for a single tool call.
+     * @param interruption The tool approval interruption item
+     */
+    async _handleSingleToolApproval(interruption) {
+        const toolName = interruption.rawItem?.name || 'Unknown Tool';
+        const toolInput = interruption.rawItem?.arguments || {};
+        const interruptionId = `int-${Date.now()}-${Math.random()}`;
+        const callId = interruption.rawItem?.callId;
+        this._pendingApprovals.set(interruptionId, { interruption });
+        this._agentEvent.emit({
+            type: 'tool_approval_required',
+            data: {
+                interruptionId,
+                toolName,
+                toolInput,
+                callId
+            }
+        });
+    }
+    /**
+     * Handles approval requests for multiple grouped tool calls.
+     * @param interruptions Array of tool approval interruption items
+     */
+    async _handleGroupedToolApprovals(interruptions) {
+        const groupId = `group-${Date.now()}-${Math.random()}`;
+        const approvals = interruptions.map(interruption => {
+            const toolName = interruption.rawItem?.name || 'Unknown Tool';
+            const toolInput = interruption.rawItem?.arguments || {};
+            const interruptionId = `int-${Date.now()}-${Math.random()}`;
+            this._pendingApprovals.set(interruptionId, { interruption, groupId });
+            return {
+                interruptionId,
+                toolName,
+                toolInput
+            };
+        });
+        this._agentEvent.emit({
+            type: 'grouped_approval_required',
+            data: {
+                groupId,
+                approvals
+            }
+        });
+    }
+    /**
+     * Checks if the current provider supports tool calling.
+     * @returns True if the provider supports tool calling, false otherwise
+     */
+    _supportsToolCalling() {
+        const activeProvider = this._settingsModel.getActiveProvider();
+        if (!activeProvider || !this._chatProviderRegistry) {
+            return false;
+        }
+        const providerInfo = this._chatProviderRegistry.getProviderInfo(activeProvider.provider);
+        // Default to true if supportsToolCalling is not specified
+        return providerInfo?.supportsToolCalling !== false;
+    }
+    /**
+     * Creates a model instance based on current settings.
+     * @returns The configured model instance for the agent
+     */
+    async _createModel() {
+        const activeProvider = this._settingsModel.getActiveProvider();
+        if (!activeProvider) {
+            throw new Error('No active provider configured');
+        }
+        const provider = activeProvider.provider;
+        const model = activeProvider.model;
+        let apiKey;
+        if (this._secretsManager && this._settingsModel.config.useSecretsManager) {
+            apiKey =
+                (await this._secretsManager.get(Private.getToken(), SECRETS_NAMESPACE, `${provider}:apiKey`))?.value ?? '';
+        }
+        else {
+            apiKey = this._settingsModel.getApiKey(activeProvider.id);
+        }
+        return createModel({
+            provider,
+            model,
+            apiKey
+        }, this._chatProviderRegistry);
+    }
+    /**
+     * Enhances the base system prompt with tool usage guidelines.
+     * @param baseSystemPrompt The base system prompt from settings
+     * @returns The enhanced system prompt with tool usage instructions
+     */
+    _getEnhancedSystemPrompt(baseSystemPrompt) {
+        const progressReportingPrompt = `
+
+IMPORTANT: Follow this message flow pattern for better user experience:
+
+1. FIRST: Explain what you're going to do and your approach
+2. THEN: Execute tools (these will show automatically with step numbers)
+3. FINALLY: Provide a concise summary of what was accomplished
+
+Example flow:
+- "I'll help you create a notebook with example cells. Let me first create the file structure, then add Python and Markdown cells."
+- [Tool executions happen with automatic step display]
+- "Successfully created your notebook with 3 cells: a title, code example, and visualization cell."
+
+Guidelines:
+- Start responses with your plan/approach before tool execution
+- Let the system handle tool execution display (don't duplicate details)
+- End with a brief summary of accomplishments
+- Use natural, conversational tone throughout
+
+COMMAND DISCOVERY:
+- When you want to execute JupyterLab commands, ALWAYS use the 'discover_commands' tool first to find available commands and their metadata.
+- This ensures you have complete information about command IDs, descriptions, and required arguments before attempting to execute them. Only after discovering the available commands should you use the 'execute_command' tool with the correct command ID and arguments.
+
+TOOL SELECTION GUIDELINES:
+- For file operations (create, read, write, modify files and directories): Use dedicated file manipulation tools
+- For general JupyterLab UI interactions (opening panels, running commands, navigating interface): Use the general command tool (execute_command)
+- Examples of file operations: Creating notebooks, editing code files, managing project structure
+- Examples of UI interactions: Opening terminal, switching tabs, running notebook cells, accessing menus
+`;
+        return baseSystemPrompt + progressReportingPrompt;
+    }
+    // Private attributes
+    _settingsModel;
+    _toolRegistry;
+    _chatProviderRegistry;
+    _secretsManager;
+    _selectedToolNames;
+    _agent;
+    _runner;
+    _history;
+    _mcpServers;
+    _isInitializing;
+    _controller;
+    _pendingApprovals;
+    _interruptedState;
+    _agentEvent;
+    _mcpConnectionChanged;
+    _tokenUsage;
+    _tokenUsageChanged;
+}
+var Private;
+(function (Private) {
+    /**
+     * The token to use with the secrets manager, setter and getter.
+     */
+    let secretsToken;
+    function setToken(value) {
+        secretsToken = value;
+    }
+    Private.setToken = setToken;
+    function getToken() {
+        return secretsToken;
+    }
+    Private.getToken = getToken;
+})(Private || (Private = {}));