@bubblelab/bubble-core 0.1.8 → 0.1.9

package/dist/bubbles/service-bubble/ai-agent.js

@@ -2,16 +2,16 @@ import { z } from 'zod';
 import { ServiceBubble } from '../../types/service-bubble-class.js';
 import { CredentialType, BUBBLE_CREDENTIAL_OPTIONS, } from '@bubblelab/shared-schemas';
 import { StateGraph, MessagesAnnotation } from '@langchain/langgraph';
-import { ToolNode } from '@langchain/langgraph/prebuilt';
 import { ChatOpenAI } from '@langchain/openai';
 import { ChatAnthropic } from '@langchain/anthropic';
 import { ChatGoogleGenerativeAI } from '@langchain/google-genai';
-import { HumanMessage, AIMessage, ToolMessage } from '@langchain/core/messages';
+import { HumanMessage, AIMessage, ToolMessage, AIMessageChunk, } from '@langchain/core/messages';
 import { DynamicStructuredTool } from '@langchain/core/tools';
-import { AvailableModels } from '../../types/ai-models.js';
-import { AvailableTools } from '../../types/available-tools.js';
+import { AvailableModels } from '@bubblelab/shared-schemas';
+import { AvailableTools, } from '../../types/available-tools.js';
 import { BubbleFactory } from '../../bubble-factory.js';
-import { parseJsonWithFallbacks } from '../../utils/json-parsing.js';
+import { extractAndStreamThinkingTokens, formatFinalResponse, } from '../../utils/agent-formatter.js';
+import { isAIMessage, isAIMessageChunk } from '@langchain/core/messages';
 // Define model configuration
 const ModelConfigSchema = z.object({
     model: AvailableModels.default('google/gemini-2.5-flash').describe('AI model to use (format: provider/model-name).'),
@@ -19,20 +19,20 @@ const ModelConfigSchema = z.object({
         .number()
         .min(0)
         .max(2)
-        .default(0.7)
+        .default(1)
         .describe('Temperature for response randomness (0 = deterministic, 2 = very random)'),
     maxTokens: z
         .number()
         .positive()
         .optional()
         .default(40000)
-        .describe('Maximum number of tokens to generate in response'),
+        .describe('Maximum number of tokens to generate in response, keep at default of 40000 unless the response is expected to be certain length'),
     jsonMode: z
         .boolean()
         .default(false)
         .describe('When true, strips markdown formatting and returns clean JSON response'),
 });
-// Define tool configuration
+// Define tool configuration for pre-registered tools
 const ToolConfigSchema = z.object({
     name: AvailableTools.describe('Name of the tool type or tool bubble to enable for the AI agent'),
     credentials: z
@@ -45,6 +45,25 @@ const ToolConfigSchema = z.object({
         .optional()
         .describe('Configuration for the tool or tool bubble'),
 });
+// Define custom tool schema for runtime-defined tools
+const CustomToolSchema = z.object({
+    name: z
+        .string()
+        .min(1)
+        .describe('Unique name for your custom tool (e.g., "calculate-tax")'),
+    description: z
+        .string()
+        .min(1)
+        .describe('Description of what the tool does - helps the AI know when to use it'),
+    schema: z
+        .record(z.string(), z.unknown())
+        .describe('Zod schema object defining the tool parameters. Example: { amount: z.number().describe("Amount to calculate tax on"), rate: z.number().describe("Tax rate") }'),
+    func: z
+        .function()
+        .args(z.record(z.string(), z.unknown()))
+        .returns(z.promise(z.unknown()))
+        .describe('Async function that executes the tool logic. Receives params matching the schema and returns a result.'),
+});
 // Define image input schemas - supports both base64 data and URLs
 const Base64ImageSchema = z.object({
     type: z.literal('base64').default('base64'),
@@ -107,12 +126,18 @@ const AIAgentParamsSchema = z.object({
             },
         },
     ])
-        .describe('Array of tools the AI agent can use. Can be tool types (web-search-tool, web-scrape-tool, web-crawl-tool, web-extract-tool). If using image models, set the tools to []'),
+        .describe('Array of pre-registered tools the AI agent can use. Can be tool types (web-search-tool, web-scrape-tool, web-crawl-tool, web-extract-tool, instagram-tool). If using image models, set the tools to []'),
+    customTools: z
+        .array(CustomToolSchema)
+        .default([])
+        .optional()
+        .describe('Array of custom runtime-defined tools with their own schemas and functions. Use this to add domain-specific tools without pre-registration. Example: [{ name: "calculate-tax", description: "Calculates sales tax", schema: { amount: z.number() }, func: async (input) => {...} }]'),
     maxIterations: z
         .number()
         .positive()
+        .min(2)
         .default(10)
-        .describe('Maximum number of iterations for the agent workflow'),
+        .describe('Maximum number of iterations for the agent workflow, 2 iterations per turn of conversation'),
     credentials: z
         .record(z.nativeEnum(CredentialType), z.string())
         .optional()
@@ -121,6 +146,8 @@ const AIAgentParamsSchema = z.object({
         .boolean()
         .default(false)
         .describe('Enable real-time streaming of tokens, tool calls, and iteration progress'),
+    // Note: beforeToolCall and afterToolCall are function hooks added via TypeScript interface
+    // They cannot be part of the Zod schema but are available in the params
 });
 const AIAgentResultSchema = z.object({
     response: z
@@ -161,11 +188,18 @@ export class AIAgentBubble extends ServiceBubble {
   `;
     static alias = 'agent';
     factory;
+    beforeToolCallHook;
+    afterToolCallHook;
+    streamingCallback;
+    shouldStopAfterTools = false;
     constructor(params = {
         message: 'Hello, how are you?',
         systemPrompt: 'You are a helpful AI assistant',
     }, context) {
         super(params, context);
+        this.beforeToolCallHook = params.beforeToolCall;
+        this.afterToolCallHook = params.afterToolCall;
+        this.streamingCallback = params.streamingCallback;
         this.factory = new BubbleFactory();
     }
     async testCredential() {
@@ -180,12 +214,12 @@ export class AIAgentBubble extends ServiceBubble {
     async performAction(context) {
         // Context is available but not currently used in this implementation
         void context;
-        const { message, images, systemPrompt, model, tools, maxIterations } = this.params;
+        const { message, images, systemPrompt, model, tools, customTools, maxIterations, } = this.params;
         try {
             // Initialize the language model
             const llm = this.initializeModel(model);
-            // Initialize tools
-            const agentTools = await this.initializeTools(tools);
+            // Initialize tools (both pre-registered and custom)
+            const agentTools = await this.initializeTools(tools, customTools);
             // Create the agent graph
             const graph = await this.createAgentGraph(llm, agentTools, systemPrompt);
             // Execute the agent
@@ -211,7 +245,7 @@ export class AIAgentBubble extends ServiceBubble {
     async actionWithStreaming(streamingCallback, context) {
         // Context is available but not currently used in this implementation
         void context;
-        const { message, images, systemPrompt, model, tools, maxIterations } = this.params;
+        const { message, images, systemPrompt, model, tools, customTools, maxIterations, } = this.params;
         const startTime = Date.now();
         // Send start event
         await streamingCallback({
@@ -233,8 +267,8 @@ export class AIAgentBubble extends ServiceBubble {
             });
             // Initialize the language model
             const llm = this.initializeModel(model);
-            // Initialize tools
-            const agentTools = await this.initializeTools(tools);
+            // Initialize tools (both pre-registered and custom)
+            const agentTools = await this.initializeTools(tools, customTools);
             // Create the agent graph
             const graph = await this.createAgentGraph(llm, agentTools, systemPrompt);
             // Execute the agent with streaming
@@ -292,70 +326,16 @@ export class AIAgentBubble extends ServiceBubble {
                 throw new Error(`Unsupported model provider: ${provider}`);
         }
     }
-    /**
-     * Format final response with special handling for Gemini image models and JSON mode
-     */
-    async formatFinalResponse(response, modelConfig, jsonMode) {
-        let finalResponse = typeof response === 'string' ? response : JSON.stringify(response);
-        // Special handling for Gemini image models that return images in inlineData format
-        if (modelConfig.model.includes('gemini') &&
-            modelConfig.model.includes('image')) {
-            finalResponse = this.formatGeminiImageResponse(finalResponse);
-        }
-        else if (jsonMode && typeof finalResponse === 'string') {
-            // Handle JSON mode: use the improved utility function
-            const result = parseJsonWithFallbacks(finalResponse);
-            if (!result.success) {
-                return {
-                    response: result.response,
-                    error: `${this.params.name || 'AI Agent'} failed to generate valid JSON. Post-processing attempted but JSON is still malformed. Original response: ${finalResponse}`,
-                };
-            }
-            return { response: result.response };
-        }
-        return { response: finalResponse };
-    }
-    /**
-     * Convert Gemini's inlineData format to LangChain-compatible data URI format
-     */
-    formatGeminiImageResponse(response) {
-        if (typeof response !== 'string') {
-            return String(response);
-        }
-        try {
-            console.log('[AIAgent] Formatting Gemini image response...');
-            // Look for Gemini's inlineData format in the response
-            const inlineDataRegex = /\{\s*"inlineData"\s*:\s*\{\s*"mimeType"\s*:\s*"([^"]+)"\s*,\s*"data"\s*:\s*"([^"]+)"\s*\}\s*\}/;
-            const match = response.match(inlineDataRegex);
-            if (match) {
-                const [, mimeType, data] = match;
-                const dataUri = `data:${mimeType};base64,${data}`;
-                console.log(`[AIAgent] Extracted first data URI from Gemini inlineData: ${mimeType}`);
-                return dataUri;
-            }
-            // Also check for the more complex format with text
-            const complexInlineDataRegex = /\{\s*"inlineData"\s*:\s*\{\s*"mimeType"\s*:\s*"([^"]+)"\s*,\s*"data"\s*:\s*"([^"]+)"/;
-            const complexMatch = response.match(complexInlineDataRegex);
-            if (complexMatch) {
-                const [, mimeType, data] = complexMatch;
-                const dataUri = `data:${mimeType};base64,${data}`;
-                console.log(`[AIAgent] Extracted first data URI from complex Gemini inlineData: ${mimeType}`);
-                return dataUri;
-            }
-            // If no inlineData found, return original response
-            return response;
-        }
-        catch (error) {
-            console.warn('[AIAgent] Error formatting Gemini image response:', error);
-            return response;
-        }
-    }
     initializeModel(modelConfig) {
         const { model, temperature, maxTokens } = modelConfig;
-        const [provider, modelName] = model.split('/');
+        const slashIndex = model.indexOf('/');
+        const provider = model.substring(0, slashIndex);
+        const modelName = model.substring(slashIndex + 1);
         // Use chooseCredential to get the appropriate credential
         // This will throw immediately if credentials are missing
         const apiKey = this.chooseCredential();
+        // Enable streaming if streamingCallback is provided
+        const enableStreaming = !!this.streamingCallback;
         switch (provider) {
             case 'openai':
                 return new ChatOpenAI({
@@ -363,6 +343,7 @@ export class AIAgentBubble extends ServiceBubble {
                     temperature,
                     maxTokens,
                     apiKey,
+                    streaming: enableStreaming,
                 });
             case 'google':
                 return new ChatGoogleGenerativeAI({
@@ -370,39 +351,70 @@ export class AIAgentBubble extends ServiceBubble {
                     temperature,
                     maxOutputTokens: maxTokens,
                     apiKey,
+                    streaming: enableStreaming,
                 });
             case 'anthropic':
                 return new ChatAnthropic({
                     model: modelName,
                     temperature,
                     anthropicApiKey: apiKey,
-                    maxTokens: 1000,
-                    streaming: false,
+                    maxTokens,
+                    streaming: enableStreaming,
                     apiKey,
                 });
             case 'openrouter':
+                console.log('openrouter', modelName);
                 return new ChatOpenAI({
                     model: modelName,
+                    __includeRawResponse: true,
                     temperature,
                     maxTokens,
                     apiKey,
+                    streaming: enableStreaming,
                     configuration: {
                         baseURL: 'https://openrouter.ai/api/v1',
                     },
+                    modelKwargs: {
+                        reasoning: {
+                            effort: 'medium',
+                            exclude: false,
+                        },
+                    },
                 });
             default:
                 throw new Error(`Unsupported model provider: ${provider}`);
         }
     }
-    async initializeTools(toolConfigs) {
+    async initializeTools(toolConfigs, customToolConfigs = []) {
         const tools = [];
         await this.factory.registerDefaults();
+        // First, initialize custom tools
+        for (const customTool of customToolConfigs) {
+            try {
+                console.log(`🛠️ [AIAgent] Initializing custom tool: ${customTool.name}`);
+                const dynamicTool = new DynamicStructuredTool({
+                    name: customTool.name,
+                    description: customTool.description,
+                    schema: z.object(customTool.schema),
+                    func: customTool.func,
+                });
+                tools.push(dynamicTool);
+            }
+            catch (error) {
+                console.error(`Error initializing custom tool '${customTool.name}':`, error);
+                // Continue with other tools even if one fails
+                continue;
+            }
+        }
+        // Then, initialize pre-registered tools from factory
         for (const toolConfig of toolConfigs) {
             try {
-                // Get the tool bubble class from the factory
                 const ToolBubbleClass = this.factory.get(toolConfig.name);
                 if (!ToolBubbleClass) {
-                    console.warn(`Tool bubble '${toolConfig.name}' not found in factory`);
+                    if (this.context && this.context.logger) {
+                        this.context.logger.warn(`Tool bubble '${toolConfig.name}' not found in factory. This tool will not be used.`);
+                    }
+                    console.warn(`Tool bubble '${toolConfig.name}' not found in factory. This tool will not be used.`);
                     continue;
                 }
                 // Check if it's a tool bubble (has toAgentTool method)
@@ -447,6 +459,106 @@ export class AIAgentBubble extends ServiceBubble {
         }
         return tools;
     }
+    /**
+     * Custom tool execution node that supports hooks
+     */
+    async executeToolsWithHooks(state, tools) {
+        const { messages } = state;
+        const lastMessage = messages[messages.length - 1];
+        const toolCalls = lastMessage.tool_calls || [];
+        const toolMessages = [];
+        let currentMessages = [...messages];
+        // Reset stop flag at the start of tool execution
+        this.shouldStopAfterTools = false;
+        // Execute each tool call
+        for (const toolCall of toolCalls) {
+            const tool = tools.find((t) => t.name === toolCall.name);
+            if (!tool) {
+                console.warn(`Tool ${toolCall.name} not found`);
+                toolMessages.push(new ToolMessage({
+                    content: `Error: Tool ${toolCall.name} not found`,
+                    tool_call_id: toolCall.id,
+                }));
+                continue;
+            }
+            try {
+                // Call beforeToolCall hook if provided
+                const hookResult_before = await this.beforeToolCallHook?.({
+                    toolName: toolCall.name,
+                    toolInput: toolCall.args,
+                    messages: currentMessages,
+                });
+                const startTime = Date.now();
+                this.streamingCallback?.({
+                    type: 'tool_start',
+                    data: {
+                        tool: toolCall.name,
+                        input: toolCall.args,
+                        callId: toolCall.id,
+                    },
+                });
+                // If hook returns modified messages/toolInput, apply them
+                if (hookResult_before) {
+                    if (hookResult_before.messages) {
+                        currentMessages = hookResult_before.messages;
+                    }
+                    toolCall.args = hookResult_before.toolInput;
+                }
+                // Execute the tool
+                const toolOutput = await tool.invoke(toolCall.args);
+                // Create tool message
+                const toolMessage = new ToolMessage({
+                    content: typeof toolOutput === 'string'
+                        ? toolOutput
+                        : JSON.stringify(toolOutput),
+                    tool_call_id: toolCall.id,
+                });
+                toolMessages.push(toolMessage);
+                currentMessages = [...currentMessages, toolMessage];
+                // Call afterToolCall hook if provided
+                const hookResult_after = await this.afterToolCallHook?.({
+                    toolName: toolCall.name,
+                    toolInput: toolCall.args,
+                    toolOutput,
+                    messages: currentMessages,
+                });
+                // If hook returns modified messages, update current messages
+                if (hookResult_after) {
+                    if (hookResult_after.messages) {
+                        currentMessages = hookResult_after.messages;
+                    }
+                    // Check if hook wants to stop execution
+                    if (hookResult_after.shouldStop === true) {
+                        this.shouldStopAfterTools = true;
+                    }
+                }
+                this.streamingCallback?.({
+                    type: 'tool_complete',
+                    data: {
+                        callId: toolCall.id,
+                        tool: toolCall.name,
+                        output: toolOutput,
+                        duration: Date.now() - startTime,
+                    },
+                });
+            }
+            catch (error) {
+                console.error(`Error executing tool ${toolCall.name}:`, error);
+                const errorMessage = new ToolMessage({
+                    content: `Error: ${error instanceof Error ? error.message : 'Unknown error'}`,
+                    tool_call_id: toolCall.id,
+                });
+                toolMessages.push(errorMessage);
+                currentMessages = [...currentMessages, errorMessage];
+            }
+        }
+        // Return the updated messages
+        // If hooks modified messages, use those; otherwise use the original messages + tool messages
+        if (currentMessages.length !== messages.length + toolMessages.length) {
+            return { messages: currentMessages };
+        }
+        return { messages: toolMessages };
+    }
     async createAgentGraph(llm, tools, systemPrompt) {
         // Define the agent node
         const agentNode = async ({ messages }) => {
@@ -455,8 +567,52 @@ export class AIAgentBubble extends ServiceBubble {
             const allMessages = [systemMessage, ...messages];
             // If we have tools, bind them to the LLM
             const modelWithTools = tools.length > 0 ? llm.bindTools(tools) : llm;
-            const response = await modelWithTools.invoke(allMessages);
-            return { messages: [response] };
+            // Use streaming if streamingCallback is provided
+            if (this.streamingCallback) {
+                const messageId = `msg-${Date.now()}-${Math.random().toString(36).substring(2, 9)}`;
+                // Use invoke with callbacks for streaming
+                const response = await modelWithTools.invoke(allMessages, {
+                    callbacks: [
+                        {
+                            handleLLMStart: async () => {
+                                await this.streamingCallback?.({
+                                    type: 'llm_start',
+                                    data: {
+                                        model: this.params.model.model,
+                                        temperature: this.params.model.temperature,
+                                    },
+                                });
+                            },
+                            handleLLMEnd: async (output) => {
+                                // Extract thinking tokens from different model providers
+                                const thinking = extractAndStreamThinkingTokens(output);
+                                if (thinking) {
+                                    await this.streamingCallback?.({
+                                        type: 'think',
+                                        data: {
+                                            content: thinking,
+                                            messageId,
+                                        },
+                                    });
+                                }
+                                await this.streamingCallback?.({
+                                    type: 'llm_complete',
+                                    data: {
+                                        messageId,
+                                        totalTokens: output.llmOutput?.usage_metadata?.total_tokens,
+                                    },
+                                });
+                            },
+                        },
+                    ],
+                });
+                return { messages: [response] };
+            }
+            else {
+                // Non-streaming fallback
+                const response = await modelWithTools.invoke(allMessages);
+                return { messages: [response] };
+            }
         };
         // Define conditional edge function
         const shouldContinue = ({ messages }) => {
@@ -467,16 +623,27 @@ export class AIAgentBubble extends ServiceBubble {
             }
             return '__end__';
         };
+        // Define conditional edge after tools to check if we should stop
+        const shouldContinueAfterTools = () => {
+            // Check if the afterToolCall hook requested stopping
+            if (this.shouldStopAfterTools) {
+                return '__end__';
+            }
+            // Otherwise continue back to agent
+            return 'agent';
+        };
         // Build the graph
         const graph = new StateGraph(MessagesAnnotation).addNode('agent', agentNode);
         if (tools.length > 0) {
-            // Use the official ToolNode for tool execution
-            const toolNode = new ToolNode(tools);
+            // Use custom tool node with hooks support
+            const toolNode = async (state) => {
+                return await this.executeToolsWithHooks(state, tools);
+            };
             graph
                 .addNode('tools', toolNode)
                 .addEdge('__start__', 'agent')
                 .addConditionalEdges('agent', shouldContinue)
-                .addEdge('tools', 'agent');
+                .addConditionalEdges('tools', shouldContinueAfterTools);
         }
         else {
             graph.addEdge('__start__', 'agent').addEdge('agent', '__end__');
@@ -545,15 +712,16 @@ export class AIAgentBubble extends ServiceBubble {
             console.log('[AIAgent] Graph execution completed');
             console.log('[AIAgent] Total messages:', result.messages.length);
             iterations = result.messages.length;
-            // Extract tool calls from messages
+            // Extract tool calls from messages and track individual LLM calls
             // Store tool calls temporarily to match with their responses
             const toolCallMap = new Map();
             for (let i = 0; i < result.messages.length; i++) {
                 const msg = result.messages[i];
-                if (msg instanceof AIMessage && msg.tool_calls) {
+                if (msg instanceof AIMessage ||
+                    (msg instanceof AIMessageChunk && msg.tool_calls)) {
                     const typedToolCalls = msg.tool_calls;
                     // Log and track tool calls
-                    for (const toolCall of typedToolCalls) {
+                    for (const toolCall of typedToolCalls || []) {
                         toolCallMap.set(toolCall.id, {
                             name: toolCall.name,
                             args: toolCall.args,
@@ -588,7 +756,7 @@ export class AIAgentBubble extends ServiceBubble {
             }
             // Get the final AI message response
             console.log('[AIAgent] Filtering AI messages...');
-            const aiMessages = result.messages.filter((msg) => msg instanceof AIMessage);
+            const aiMessages = result.messages.filter((msg) => isAIMessage(msg) || isAIMessageChunk(msg));
             console.log('[AIAgent] Found', aiMessages.length, 'AI messages');
             const finalMessage = aiMessages[aiMessages.length - 1];
             // Check for MAX_TOKENS finish reason
@@ -601,10 +769,17 @@ export class AIAgentBubble extends ServiceBubble {
             let totalOutputTokens = 0;
             let totalTokensSum = 0;
             for (const msg of result.messages) {
-                if (msg instanceof AIMessage && msg.usage_metadata) {
-                    totalInputTokens += msg.usage_metadata.input_tokens || 0;
-                    totalOutputTokens += msg.usage_metadata.output_tokens || 0;
-                    totalTokensSum += msg.usage_metadata.total_tokens || 0;
+                if (msg instanceof AIMessage ||
+                    (msg instanceof AIMessageChunk && msg.usage_metadata)) {
+                    totalInputTokens +=
+                        msg.usage_metadata?.input_tokens ||
+                            0;
+                    totalOutputTokens +=
+                        msg.usage_metadata?.output_tokens ||
+                            0;
+                    totalTokensSum +=
+                        msg.usage_metadata?.total_tokens ||
+                            0;
                 }
             }
             if (totalTokensSum > 0 && this.context && this.context.logger) {
@@ -615,12 +790,13 @@ export class AIAgentBubble extends ServiceBubble {
                     modelName: this.params.model.model,
                 }, `LLM completion: ${totalInputTokens} input + ${totalOutputTokens} output = ${totalTokensSum} total tokens`, {
                     bubbleName: 'ai-agent',
+                    variableId: this.context?.variableId,
                     operationType: 'bubble_execution',
                 });
             }
             const response = finalMessage?.content || 'No response generated';
             // Use shared formatting method
-            const formattedResult = await this.formatFinalResponse(response, this.params.model, jsonMode);
+            const formattedResult = await formatFinalResponse(response, this.params.model.model, jsonMode);
             // If there's an error from formatting (e.g., invalid JSON), return early
             if (formattedResult.error) {
                 return {
@@ -778,6 +954,7 @@ export class AIAgentBubble extends ServiceBubble {
                                 };
                                 this.context.logger.logTokenUsage(tokenUsage, `LLM completion: ${tokenUsage.inputTokens} input + ${tokenUsage.outputTokens} output = ${tokenUsage.totalTokens} total tokens`, {
                                     bubbleName: 'ai-agent',
+                                    variableId: this.context?.variableId,
                                     operationType: 'bubble_execution',
                                 });
                             }
@@ -824,6 +1001,7 @@ export class AIAgentBubble extends ServiceBubble {
                                     type: 'tool_complete',
                                     data: {
                                         callId,
+                                        tool: callData.name,
                                         output: event.data.output,
                                         duration,
                                     },
@@ -856,7 +1034,7 @@ export class AIAgentBubble extends ServiceBubble {
             // Process final result
             const accumulatedResponse = accumulatedContent || 'No response generated';
             // Use shared formatting method
-            const formattedResult = await this.formatFinalResponse(accumulatedResponse, this.params.model, jsonMode);
+            const formattedResult = await formatFinalResponse(accumulatedResponse, this.params.model.model, jsonMode);
             // If there's an error from formatting (e.g., invalid JSON), return early with consistent behavior
             if (formattedResult.error) {
                 return {