@bubblelab/bubble-core 0.1.10 → 0.1.11

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

@@ -1,20 +1,49 @@
 import { z } from 'zod';
 import { ServiceBubble } from '../../types/service-bubble-class.js';
-import { CredentialType, BUBBLE_CREDENTIAL_OPTIONS, } from '@bubblelab/shared-schemas';
+import { CredentialType, BUBBLE_CREDENTIAL_OPTIONS, RECOMMENDED_MODELS, } from '@bubblelab/shared-schemas';
 import { StateGraph, MessagesAnnotation } from '@langchain/langgraph';
 import { ChatOpenAI } from '@langchain/openai';
 import { ChatAnthropic } from '@langchain/anthropic';
-import { ChatGoogleGenerativeAI } from '@langchain/google-genai';
 import { HumanMessage, AIMessage, ToolMessage, AIMessageChunk, } from '@langchain/core/messages';
 import { DynamicStructuredTool } from '@langchain/core/tools';
 import { AvailableModels } from '@bubblelab/shared-schemas';
 import { AvailableTools, } from '../../types/available-tools.js';
 import { BubbleFactory } from '../../bubble-factory.js';
-import { extractAndStreamThinkingTokens, formatFinalResponse, } from '../../utils/agent-formatter.js';
+import { ConversationMessageSchema } from '@bubblelab/shared-schemas';
+import { extractAndStreamThinkingTokens, formatFinalResponse, generationsToMessageContent, } from '../../utils/agent-formatter.js';
 import { isAIMessage, isAIMessageChunk } from '@langchain/core/messages';
+import { HarmBlockThreshold, HarmCategory } from '@google/generative-ai';
+import { SafeGeminiChat } from '../../utils/safe-gemini-chat.js';
+import { zodSchemaToJsonString, buildJsonSchemaInstruction, } from '../../utils/zod-schema.js';
+// Define backup model configuration schema
+const BackupModelConfigSchema = z.object({
+    model: AvailableModels.describe('Backup AI model to use if the primary model fails (format: provider/model-name).'),
+    temperature: z
+        .number()
+        .min(0)
+        .max(2)
+        .optional()
+        .describe('Temperature for backup model. If not specified, uses primary model temperature.'),
+    maxTokens: z
+        .number()
+        .positive()
+        .optional()
+        .describe('Max tokens for backup model. If not specified, uses primary model maxTokens.'),
+    reasoningEffort: z
+        .enum(['low', 'medium', 'high'])
+        .optional()
+        .describe('Reasoning effort for backup model. If not specified, uses primary model reasoningEffort.'),
+    maxRetries: z
+        .number()
+        .int()
+        .min(0)
+        .max(10)
+        .optional()
+        .describe('Max retries for backup model. If not specified, uses primary model maxRetries.'),
+});
 // Define model configuration
 const ModelConfigSchema = z.object({
-    model: AvailableModels.default('google/gemini-2.5-flash').describe('AI model to use (format: provider/model-name).'),
+    model: AvailableModels.describe('AI model to use (format: provider/model-name).'),
     temperature: z
         .number()
         .min(0)
@@ -27,10 +56,30 @@ const ModelConfigSchema = z.object({
         .optional()
         .default(12800)
         .describe('Maximum number of tokens to generate in response, keep at default of 40000 unless the response is expected to be certain length'),
+    reasoningEffort: z
+        .enum(['low', 'medium', 'high'])
+        .optional()
+        .describe('Reasoning effort for model. If not specified, uses primary model reasoningEffort.'),
+    maxRetries: z
+        .number()
+        .int()
+        .min(0)
+        .max(10)
+        .default(3)
+        .describe('Maximum number of retries for API calls (default: 3). Useful for handling transient errors like 503 Service Unavailable.'),
+    provider: z
+        .array(z.string())
+        .optional()
+        .describe('Providers for ai agent (open router only).'),
     jsonMode: z
         .boolean()
         .default(false)
-        .describe('When true, strips markdown formatting and returns clean JSON response'),
+        .describe('When true, returns clean JSON response, you must provide the exact JSON schema in the system prompt'),
+    backupModel: BackupModelConfigSchema.default({
+        model: RECOMMENDED_MODELS.FAST,
+    })
+        .optional()
+        .describe('Backup model configuration to use if the primary model fails.'),
 });
 // Define tool configuration for pre-registered tools
 const ToolConfigSchema = z.object({
@@ -56,8 +105,11 @@ const CustomToolSchema = z.object({
         .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") }'),
+        .union([
+        z.record(z.string(), z.unknown()),
+        z.custom((val) => val && typeof val === 'object' && '_def' in val),
+    ])
+        .describe('Zod schema object defining the tool parameters. Can be either a plain object (e.g., { amount: z.number() }) or a Zod object directly (e.g., z.object({ amount: z.number() })).'),
     func: z
         .function()
         .args(z.record(z.string(), z.unknown()))
@@ -91,6 +143,11 @@ const ImageInputSchema = z.discriminatedUnion('type', [
     Base64ImageSchema,
     UrlImageSchema,
 ]);
+// Schema for the expected JSON output structure - accepts either a Zod schema or a JSON schema string
+const ExpectedOutputSchema = z.union([
+    z.custom((val) => val?._def !== undefined),
+    z.string(),
+]);
 // Define the parameters schema for the AI Agent bubble
 const AIAgentParamsSchema = z.object({
     message: z
@@ -101,6 +158,10 @@ const AIAgentParamsSchema = z.object({
         .array(ImageInputSchema)
         .default([])
         .describe('Array of base64 encoded images to include with the message (for multimodal AI models). Example: [{type: "base64", data: "base64...", mimeType: "image/png", description: "A beautiful image of a cat"}] or [{type: "url", url: "https://example.com/image.png", description: "A beautiful image of a cat"}]'),
+    conversationHistory: z
+        .array(ConversationMessageSchema)
+        .optional()
+        .describe('Previous conversation messages for multi-turn conversations. When provided, messages are sent as separate turns to enable KV cache optimization. Format: [{role: "user", content: "..."}, {role: "assistant", content: "..."}, ...]'),
     systemPrompt: z
         .string()
         .default('You are a helpful AI assistant')
@@ -111,21 +172,15 @@ const AIAgentParamsSchema = z.object({
         .optional()
         .describe('A friendly name for the AI agent'),
     model: ModelConfigSchema.default({
-        model: 'google/gemini-2.5-flash',
-        temperature: 0.7,
+        model: RECOMMENDED_MODELS.FAST,
+        temperature: 1,
         maxTokens: 50000,
+        maxRetries: 3,
         jsonMode: false,
-    }).describe('AI model configuration including provider, temperature, and tokens. For model unless otherwise specified, use google/gemini-2.5-flash as default. Use google/gemini-2.5-flash-image-preview to edit and generate images.'),
+    }).describe('AI model configuration including provider, temperature, and tokens, retries, and json mode. Always include this.'),
     tools: z
         .array(ToolConfigSchema)
-        .default([
-        {
-            name: 'web-search-tool',
-            config: {
-                maxResults: 5,
-            },
-        },
-    ])
+        .default([])
         .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)
@@ -135,9 +190,9 @@ const AIAgentParamsSchema = z.object({
     maxIterations: z
         .number()
         .positive()
-        .min(2)
-        .default(10)
-        .describe('Maximum number of iterations for the agent workflow, 2 iterations per turn of conversation'),
+        .min(4)
+        .default(40)
+        .describe('Maximum number of iterations for the agent workflow, 5 iterations per turn of conversation'),
     credentials: z
         .record(z.nativeEnum(CredentialType), z.string())
         .optional()
@@ -146,13 +201,14 @@ const AIAgentParamsSchema = z.object({
         .boolean()
         .default(false)
         .describe('Enable real-time streaming of tokens, tool calls, and iteration progress'),
+    expectedOutputSchema: ExpectedOutputSchema.optional().describe('Zod schema or JSON schema string that defines the expected structure of the AI response. When provided, automatically enables JSON mode and instructs the AI to output in the exact format. Example: z.object({ summary: z.string(), items: z.array(z.object({ name: z.string(), score: z.number() })) })'),
     // 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
         .string()
-        .describe('The AI agents final response to the user message. For text responses, returns plain text or JSON string. For image generation models (like gemini-2.5-flash-image-preview), returns base64-encoded image data with data URI format (data:image/png;base64,...)'),
+        .describe('The AI agents final response to the user message. For text responses, returns plain text. If JSON mode is enabled, returns a JSON string. For image generation models (like gemini-2.5-flash-image-preview), returns base64-encoded image data with data URI format (data:image/png;base64,...)'),
     toolCalls: z
         .array(z.object({
         tool: z.string().describe('Name of the tool that was called'),
@@ -190,15 +246,19 @@ export class AIAgentBubble extends ServiceBubble {
     factory;
     beforeToolCallHook;
     afterToolCallHook;
+    afterLLMCallHook;
     streamingCallback;
     shouldStopAfterTools = false;
+    shouldContinueToAgent = false;
     constructor(params = {
         message: 'Hello, how are you?',
         systemPrompt: 'You are a helpful AI assistant',
-    }, context) {
-        super(params, context);
+        model: { model: RECOMMENDED_MODELS.FAST },
+    }, context, instanceId) {
+        super(params, context, instanceId);
         this.beforeToolCallHook = params.beforeToolCall;
         this.afterToolCallHook = params.afterToolCall;
+        this.afterLLMCallHook = params.afterLLMCall;
         this.streamingCallback = params.streamingCallback;
         this.factory = new BubbleFactory();
     }
@@ -211,97 +271,89 @@ export class AIAgentBubble extends ServiceBubble {
         }
         return false;
     }
+    /**
+     * Build effective model config from primary and optional backup settings
+     */
+    buildModelConfig(primaryConfig, backupConfig) {
+        if (!backupConfig) {
+            return primaryConfig;
+        }
+        return {
+            model: backupConfig.model,
+            temperature: backupConfig.temperature ?? primaryConfig.temperature,
+            maxTokens: backupConfig.maxTokens ?? primaryConfig.maxTokens,
+            maxRetries: backupConfig.maxRetries ?? primaryConfig.maxRetries,
+            provider: primaryConfig.provider,
+            jsonMode: primaryConfig.jsonMode,
+            backupModel: undefined, // Don't chain backup models
+        };
+    }
+    /**
+     * Core execution logic for running the agent with a given model config
+     */
+    async executeWithModel(modelConfig) {
+        const { message, images, systemPrompt, tools, customTools, maxIterations, conversationHistory, } = this.params;
+        // Initialize the language model
+        const llm = this.initializeModel(modelConfig);
+        // 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
+        return this.executeAgent(graph, message, images, maxIterations, modelConfig, conversationHistory);
+    }
+    /**
+     * Modify params before execution - centralizes all param transformations
+     */
+    beforeAction() {
+        // Auto-enable JSON mode when expectedOutputSchema is provided
+        if (this.params.expectedOutputSchema) {
+            this.params.model.jsonMode = true;
+            // Enhance system prompt with JSON schema instructions
+            const schemaString = zodSchemaToJsonString(this.params.expectedOutputSchema);
+            this.params.systemPrompt = `${this.params.systemPrompt}\n\n${buildJsonSchemaInstruction(schemaString)}`;
+        }
+    }
     async performAction(context) {
         // Context is available but not currently used in this implementation
         void context;
-        const { message, images, systemPrompt, model, tools, customTools, maxIterations, } = this.params;
+        // Apply param transformations before execution
+        this.beforeAction();
         try {
-            // Initialize the language model
-            const llm = this.initializeModel(model);
-            // 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
-            const result = await this.executeAgent(graph, message, images, maxIterations, model.jsonMode);
-            return result;
+            return await this.executeWithModel(this.params.model);
         }
         catch (error) {
-            // Return error information but mark as recoverable
             const errorMessage = error instanceof Error ? error.message : 'Unknown error';
-            console.warn('[AIAgent] Execution error (continuing):', errorMessage);
+            console.warn('[AIAgent] Execution error:', errorMessage);
+            // Return error information but mark as recoverable
             return {
                 response: `Error: ${errorMessage}`,
-                success: false, // Still false but execution can continue
+                success: false,
                 toolCalls: [],
                 error: errorMessage,
                 iterations: 0,
             };
         }
     }
+    getCredentialType() {
+        return this.getCredentialTypeForModel(this.params.model.model);
+    }
     /**
-     * Execute the AI agent with streaming support for real-time feedback
+     * Get credential type for a specific model string
      */
-    async actionWithStreaming(streamingCallback, context) {
-        // Context is available but not currently used in this implementation
-        void context;
-        const { message, images, systemPrompt, model, tools, customTools, maxIterations, } = this.params;
-        const startTime = Date.now();
-        // Send start event
-        await streamingCallback({
-            type: 'start',
-            data: {
-                message: `Analyzing with ${this.params.name || 'AI Agent'}`,
-                maxIterations,
-                timestamp: new Date().toISOString(),
-            },
-        });
-        try {
-            // Send LLM start event
-            await streamingCallback({
-                type: 'llm_start',
-                data: {
-                    model: model.model,
-                    temperature: model.temperature,
-                },
-            });
-            // Initialize the language model
-            const llm = this.initializeModel(model);
-            // 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
-            const result = await this.executeAgentWithStreaming(graph, message, images, maxIterations, model.jsonMode, streamingCallback);
-            const totalDuration = Date.now() - startTime;
-            // Send completion event
-            await streamingCallback({
-                type: 'complete',
-                data: {
-                    result,
-                    totalDuration,
-                },
-            });
-            return result;
-        }
-        catch (error) {
-            const errorMessage = error instanceof Error ? error.message : 'Unknown error';
-            // Send error event as recoverable
-            await streamingCallback({
-                type: 'error',
-                data: {
-                    error: errorMessage,
-                    recoverable: true, // Mark as recoverable to continue execution
-                },
-            });
-            console.warn('[AIAgent] Streaming execution error (continuing):', errorMessage);
-            return {
-                response: `Error: ${errorMessage}`,
-                success: false, // Still false but execution can continue
-                toolCalls: [],
-                error: errorMessage,
-                iterations: 0,
-            };
+    getCredentialTypeForModel(model) {
+        const [provider] = model.split('/');
+        switch (provider) {
+            case 'openai':
+                return CredentialType.OPENAI_CRED;
+            case 'google':
+                return CredentialType.GOOGLE_GEMINI_CRED;
+            case 'anthropic':
+                return CredentialType.ANTHROPIC_CRED;
+            case 'openrouter':
+                return CredentialType.OPENROUTER_CRED;
+            default:
+                throw new Error(`Unsupported model provider: ${provider}`);
         }
     }
     chooseCredential() {
@@ -327,15 +379,40 @@ export class AIAgentBubble extends ServiceBubble {
         }
     }
     initializeModel(modelConfig) {
-        const { model, temperature, maxTokens } = modelConfig;
+        const { model, temperature, maxTokens, maxRetries } = modelConfig;
         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();
+        const reasoningEffort = modelConfig.reasoningEffort;
+        // Get credential based on the modelConfig's provider (not this.params.model)
+        const credentials = this.params.credentials;
+        if (!credentials || typeof credentials !== 'object') {
+            throw new Error(`No ${provider.toUpperCase()} credentials provided`);
+        }
+        let apiKey;
+        switch (provider) {
+            case 'openai':
+                apiKey = credentials[CredentialType.OPENAI_CRED];
+                break;
+            case 'google':
+                apiKey = credentials[CredentialType.GOOGLE_GEMINI_CRED];
+                break;
+            case 'anthropic':
+                apiKey = credentials[CredentialType.ANTHROPIC_CRED];
+                break;
+            case 'openrouter':
+                apiKey = credentials[CredentialType.OPENROUTER_CRED];
+                break;
+            default:
+                throw new Error(`Unsupported model provider: ${provider}`);
+        }
+        if (!apiKey) {
+            throw new Error(`No credential found for provider: ${provider}`);
+        }
         // Enable streaming if streamingCallback is provided
         const enableStreaming = !!this.streamingCallback;
+        // Default to 3 retries if not specified
+        const retries = maxRetries ?? 3;
         switch (provider) {
             case 'openai':
                 return new ChatOpenAI({
@@ -343,17 +420,70 @@ export class AIAgentBubble extends ServiceBubble {
                     temperature,
                     maxTokens,
                     apiKey,
+                    ...(reasoningEffort && {
+                        reasoning: {
+                            effort: reasoningEffort,
+                            summary: 'auto',
+                        },
+                    }),
                     streaming: enableStreaming,
+                    maxRetries: retries,
                 });
-            case 'google':
-                return new ChatGoogleGenerativeAI({
+            case 'google': {
+                const thinkingConfig = reasoningEffort
+                    ? {
+                        includeThoughts: reasoningEffort ? true : false,
+                        thinkingBudget: reasoningEffort === 'low'
+                            ? 1025
+                            : reasoningEffort === 'medium'
+                                ? 5000
+                                : 10000,
+                    }
+                    : undefined;
+                return new SafeGeminiChat({
                     model: modelName,
                     temperature,
                     maxOutputTokens: maxTokens,
+                    ...(thinkingConfig && { thinkingConfig }),
                     apiKey,
-                    streaming: enableStreaming,
+                    // 3.0 pro preview does breaks with streaming, disabled temporarily until fixed
+                    streaming: false,
+                    maxRetries: retries,
+                    // Disable all safety filters to prevent candidateContent.parts.reduce errors
+                    // when Gemini blocks content and returns candidates without content field
+                    safetySettings: [
+                        {
+                            category: HarmCategory.HARM_CATEGORY_HARASSMENT,
+                            threshold: HarmBlockThreshold.BLOCK_NONE,
+                        },
+                        {
+                            category: HarmCategory.HARM_CATEGORY_HATE_SPEECH,
+                            threshold: HarmBlockThreshold.BLOCK_NONE,
+                        },
+                        {
+                            category: HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
+                            threshold: HarmBlockThreshold.BLOCK_NONE,
+                        },
+                        {
+                            category: HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
+                            threshold: HarmBlockThreshold.BLOCK_NONE,
+                        },
+                    ],
                 });
-            case 'anthropic':
+            }
+            case 'anthropic': {
+                // Configure Anthropic "thinking" only when reasoning is enabled.
+                // Anthropic's API does not allow `budget_tokens` when thinking is disabled.
+                const thinkingConfig = reasoningEffort != null
+                    ? {
+                        type: 'enabled',
+                        budget_tokens: reasoningEffort === 'low'
+                            ? 1025
+                            : reasoningEffort === 'medium'
+                                ? 5000
+                                : 10000,
+                    }
+                    : undefined;
                 return new ChatAnthropic({
                     model: modelName,
                     temperature,
@@ -361,7 +491,10 @@ export class AIAgentBubble extends ServiceBubble {
                     maxTokens,
                     streaming: enableStreaming,
                     apiKey,
+                    ...(thinkingConfig && { thinking: thinkingConfig }),
+                    maxRetries: retries,
                 });
+            }
             case 'openrouter':
                 console.log('openrouter', modelName);
                 return new ChatOpenAI({
@@ -371,12 +504,16 @@ export class AIAgentBubble extends ServiceBubble {
                     maxTokens,
                     apiKey,
                     streaming: enableStreaming,
+                    maxRetries: retries,
                     configuration: {
                         baseURL: 'https://openrouter.ai/api/v1',
                     },
                     modelKwargs: {
+                        provider: {
+                            order: this.params.model.provider,
+                        },
                         reasoning: {
-                            effort: 'medium',
+                            effort: reasoningEffort ?? 'medium',
                             exclude: false,
                         },
                     },
@@ -392,10 +529,22 @@ export class AIAgentBubble extends ServiceBubble {
         for (const customTool of customToolConfigs) {
             try {
                 console.log(`🛠️ [AIAgent] Initializing custom tool: ${customTool.name}`);
+                // Handle both plain object and Zod object schemas
+                let schema;
+                if (customTool.schema &&
+                    typeof customTool.schema === 'object' &&
+                    '_def' in customTool.schema) {
+                    // Already a Zod schema object, use it directly
+                    schema = customTool.schema;
+                }
+                else {
+                    // Plain object, convert to Zod object
+                    schema = z.object(customTool.schema);
+                }
                 const dynamicTool = new DynamicStructuredTool({
                     name: customTool.name,
                     description: customTool.description,
-                    schema: z.object(customTool.schema),
+                    schema: schema,
                     func: customTool.func,
                 });
                 tools.push(dynamicTool);
@@ -475,12 +624,35 @@ export class AIAgentBubble extends ServiceBubble {
             const tool = tools.find((t) => t.name === toolCall.name);
             if (!tool) {
                 console.warn(`Tool ${toolCall.name} not found`);
+                const errorContent = `Error: Tool ${toolCall.name} not found`;
+                const startTime = Date.now();
+                // Send tool_start event
+                this.streamingCallback?.({
+                    type: 'tool_start',
+                    data: {
+                        tool: toolCall.name,
+                        input: toolCall.args,
+                        callId: toolCall.id,
+                    },
+                });
+                // Send tool_complete event with error
+                this.streamingCallback?.({
+                    type: 'tool_complete',
+                    data: {
+                        callId: toolCall.id,
+                        input: toolCall.args,
+                        tool: toolCall.name,
+                        output: { error: errorContent },
+                        duration: Date.now() - startTime,
+                    },
+                });
                 toolMessages.push(new ToolMessage({
-                    content: `Error: Tool ${toolCall.name} not found`,
+                    content: errorContent,
                     tool_call_id: toolCall.id,
                 }));
                 continue;
             }
+            const startTime = Date.now();
             try {
                 // Call beforeToolCall hook if provided
                 const hookResult_before = await this.beforeToolCallHook?.({
@@ -488,7 +660,6 @@ export class AIAgentBubble extends ServiceBubble {
                     toolInput: toolCall.args,
                     messages: currentMessages,
                 });
-                const startTime = Date.now();
                 this.streamingCallback?.({
                     type: 'tool_start',
                     data: {
@@ -545,17 +716,30 @@ export class AIAgentBubble extends ServiceBubble {
             }
             catch (error) {
                 console.error(`Error executing tool ${toolCall.name}:`, error);
+                const errorContent = `Error: ${error instanceof Error ? error.message : 'Unknown error'}`;
                 const errorMessage = new ToolMessage({
-                    content: `Error: ${error instanceof Error ? error.message : 'Unknown error'}`,
+                    content: errorContent,
                     tool_call_id: toolCall.id,
                 });
                 toolMessages.push(errorMessage);
                 currentMessages = [...currentMessages, errorMessage];
+                // Send tool_complete event even on failure so frontend can track it properly
+                this.streamingCallback?.({
+                    type: 'tool_complete',
+                    data: {
+                        callId: toolCall.id,
+                        input: toolCall.args,
+                        tool: toolCall.name,
+                        output: { error: errorContent },
+                        duration: Date.now() - startTime,
+                    },
+                });
             }
         }
         // Return the updated messages
         // If hooks modified messages, use those; otherwise use the original messages + tool messages
         if (currentMessages.length !== messages.length + toolMessages.length) {
+            console.error('[AIAgent] Current messages length does not match expected length', currentMessages.length, messages.length, toolMessages.length);
             return { messages: currentMessages };
         }
         return { messages: toolMessages };
@@ -563,63 +747,180 @@ export class AIAgentBubble extends ServiceBubble {
     async createAgentGraph(llm, tools, systemPrompt) {
         // Define the agent node
         const agentNode = async ({ messages }) => {
-            // Enhance system prompt for JSON mode
+            // systemPrompt is already enhanced by beforeAction() if expectedOutputSchema was provided
             const systemMessage = new HumanMessage(systemPrompt);
             const allMessages = [systemMessage, ...messages];
-            // If we have tools, bind them to the LLM
-            const modelWithTools = tools.length > 0 ? llm.bindTools(tools) : llm;
-            // 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) {
+            // Helper function for exponential backoff with jitter
+            const exponentialBackoff = (attemptNumber) => {
+                // Base delay: 1 second, exponentially increases (1s, 2s, 4s, 8s, ...)
+                const baseDelay = 1000;
+                const maxDelay = 32000; // Cap at 32 seconds
+                const delay = Math.min(baseDelay * Math.pow(2, attemptNumber - 1), maxDelay);
+                // Add jitter (random ±25% variation) to prevent thundering herd
+                const jitter = delay * 0.25 * (Math.random() - 0.5);
+                const finalDelay = delay + jitter;
+                return new Promise((resolve) => setTimeout(resolve, finalDelay));
+            };
+            // Shared onFailedAttempt callback to avoid duplication
+            const onFailedAttempt = async (error) => {
+                const attemptNumber = error.attemptNumber;
+                const retriesLeft = error.retriesLeft;
+                // Check if this is a candidateContent error
+                const errorMessage = error.message || String(error);
+                if (errorMessage.includes('candidateContent') ||
+                    errorMessage.includes('parts.reduce') ||
+                    errorMessage.includes('undefined is not an object')) {
+                    this.context?.logger?.error(`[AIAgent] Gemini candidateContent error detected (attempt ${attemptNumber}). This indicates blocked/empty content from Gemini API.`);
+                }
+                this.context?.logger?.warn(`[AIAgent] LLM call failed (attempt ${attemptNumber}/${this.params.model.maxRetries}). Retries left: ${retriesLeft}. Error: ${error.message}`);
+                // Optionally emit streaming event for retry
+                if (this.streamingCallback) {
+                    await this.streamingCallback({
+                        type: 'error',
+                        data: {
+                            error: `Retry attempt ${attemptNumber}/${this.params.model.maxRetries}: ${error.message}`,
+                            recoverable: retriesLeft > 0,
+                        },
+                    });
+                }
+                // Wait with exponential backoff before retrying
+                if (retriesLeft > 0) {
+                    await exponentialBackoff(attemptNumber);
+                }
+            };
+            // If we have tools, bind them to the LLM, then add retry logic
+            // IMPORTANT: Must bind tools FIRST, then add retry - not the other way around
+            const modelWithTools = tools.length > 0
+                ? llm.bindTools(tools).withRetry({
+                    stopAfterAttempt: this.params.model.maxRetries,
+                    onFailedAttempt,
+                })
+                : llm.withRetry({
+                    stopAfterAttempt: this.params.model.maxRetries,
+                    onFailedAttempt,
+                });
+            try {
+                // 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,
+                                            },
+                                        });
+                                    }
+                                    const content = formatFinalResponse(generationsToMessageContent(output.generations.flat()), this.params.model.model).response;
                                     await this.streamingCallback?.({
-                                        type: 'think',
+                                        type: 'llm_complete',
                                         data: {
-                                            content: thinking,
                                             messageId,
+                                            content: content,
+                                            totalTokens: output.llmOutput?.usage_metadata?.total_tokens,
                                         },
                                     });
-                                }
-                                await this.streamingCallback?.({
-                                    type: 'llm_complete',
-                                    data: {
-                                        messageId,
-                                        totalTokens: output.llmOutput?.usage_metadata?.total_tokens,
-                                    },
-                                });
+                                },
                             },
-                        },
-                    ],
-                });
-                return { messages: [response] };
+                        ],
+                    });
+                    return { messages: [response] };
+                }
+                else {
+                    // Non-streaming fallback
+                    const response = await modelWithTools.invoke(allMessages);
+                    return { messages: [response] };
+                }
             }
-            else {
-                // Non-streaming fallback
-                const response = await modelWithTools.invoke(allMessages);
-                return { messages: [response] };
+            catch (error) {
+                // Catch candidateContent errors that slip through SafeGeminiChat
+                const errorMessage = error instanceof Error ? error.message : String(error);
+                if (errorMessage.includes('candidateContent') ||
+                    errorMessage.includes('parts.reduce') ||
+                    errorMessage.includes('undefined is not an object')) {
+                    console.error('[AIAgent] Caught candidateContent error in agentNode:', errorMessage);
+                    // Return error as AIMessage instead of crashing
+                    return {
+                        messages: [
+                            new AIMessage({
+                                content: `[Gemini Error] Unable to generate response due to content filtering. Error: ${errorMessage}`,
+                                additional_kwargs: {
+                                    finishReason: 'ERROR',
+                                    error: errorMessage,
+                                },
+                            }),
+                        ],
+                    };
+                }
+                // Rethrow other errors
+                throw error;
             }
         };
-        // Define conditional edge function
-        const shouldContinue = ({ messages }) => {
+        // Node that runs after agent to check afterLLMCall hook before routing
+        const afterLLMCheckNode = async ({ messages, }) => {
+            // Reset the flag at the start
+            this.shouldContinueToAgent = false;
+            // Get the last AI message
             const lastMessage = messages[messages.length - 1];
-            // Check if the last message has tool calls
-            if (lastMessage.tool_calls && lastMessage.tool_calls.length > 0) {
+            const hasToolCalls = !!(lastMessage.tool_calls && lastMessage.tool_calls.length > 0);
+            // Only call hook if we're about to end (no tool calls) and hook is provided
+            if (!hasToolCalls && this.afterLLMCallHook) {
+                console.log('[AIAgent] No tool calls detected, calling afterLLMCall hook');
+                const hookResult = await this.afterLLMCallHook({
+                    messages,
+                    lastAIMessage: lastMessage,
+                    hasToolCalls,
+                });
+                // If hook wants to continue to agent, set flag and return modified messages
+                if (hookResult.continueToAgent) {
+                    console.log('[AIAgent] afterLLMCall hook requested retry to agent');
+                    this.shouldContinueToAgent = true;
+                    // Return the modified messages from the hook
+                    // We need to return only the new messages to append
+                    const newMessages = hookResult.messages.slice(messages.length);
+                    return { messages: newMessages };
+                }
+            }
+            // No modifications needed
+            return { messages: [] };
+        };
+        // Define conditional edge function after LLM check
+        const shouldContinueAfterLLMCheck = ({ messages, }) => {
+            // First check if afterLLMCall hook requested continuing to agent
+            if (this.shouldContinueToAgent) {
+                return 'agent';
+            }
+            // Find the last AI message (could be followed by human messages from hook)
+            const aiMessages = [];
+            for (const msg of messages) {
+                if (isAIMessage(msg)) {
+                    aiMessages.push(msg);
+                }
+                else if ('tool_calls' in msg &&
+                    msg.constructor?.name === 'AIMessageChunk') {
+                    aiMessages.push(msg);
+                }
+            }
+            const lastAIMessage = aiMessages[aiMessages.length - 1];
+            // Check if the last AI message has tool calls
+            if (lastAIMessage?.tool_calls && lastAIMessage.tool_calls.length > 0) {
                 return 'tools';
             }
             return '__end__';
@@ -642,23 +943,57 @@ export class AIAgentBubble extends ServiceBubble {
             };
             graph
                 .addNode('tools', toolNode)
+                .addNode('afterLLMCheck', afterLLMCheckNode)
                 .addEdge('__start__', 'agent')
-                .addConditionalEdges('agent', shouldContinue)
+                .addEdge('agent', 'afterLLMCheck')
+                .addConditionalEdges('afterLLMCheck', shouldContinueAfterLLMCheck)
                 .addConditionalEdges('tools', shouldContinueAfterTools);
         }
         else {
-            graph.addEdge('__start__', 'agent').addEdge('agent', '__end__');
+            // Even without tools, add the afterLLMCheck node for hook support
+            graph
+                .addNode('afterLLMCheck', afterLLMCheckNode)
+                .addEdge('__start__', 'agent')
+                .addEdge('agent', 'afterLLMCheck')
+                .addConditionalEdges('afterLLMCheck', shouldContinueAfterLLMCheck);
         }
         return graph.compile();
     }
-    async executeAgent(graph, message, images, maxIterations, jsonMode) {
+    async executeAgent(graph, message, images, maxIterations, modelConfig, conversationHistory) {
+        const jsonMode = modelConfig.jsonMode;
         const toolCalls = [];
         let iterations = 0;
         console.log('[AIAgent] Starting execution with message:', message.substring(0, 100) + '...');
         console.log('[AIAgent] Max iterations:', maxIterations);
         try {
             console.log('[AIAgent] Invoking graph...');
-            // Create human message with text and optional images
+            // Build messages array starting with conversation history (for KV cache optimization)
+            const initialMessages = [];
+            // Convert conversation history to LangChain messages if provided
+            // This enables KV cache optimization by keeping previous turns as separate messages
+            if (conversationHistory && conversationHistory.length > 0) {
+                for (const historyMsg of conversationHistory) {
+                    switch (historyMsg.role) {
+                        case 'user':
+                            initialMessages.push(new HumanMessage(historyMsg.content));
+                            break;
+                        case 'assistant':
+                            initialMessages.push(new AIMessage(historyMsg.content));
+                            break;
+                        case 'tool':
+                            // Tool messages require a tool_call_id
+                            if (historyMsg.toolCallId) {
+                                initialMessages.push(new ToolMessage({
+                                    content: historyMsg.content,
+                                    tool_call_id: historyMsg.toolCallId,
+                                    name: historyMsg.name,
+                                }));
+                            }
+                            break;
+                    }
+                }
+            }
+            // Create the current human message with text and optional images
             let humanMessage;
             if (images && images.length > 0) {
                 console.log('[AIAgent] Creating multimodal message with', images.length, 'images');
@@ -709,7 +1044,9 @@ export class AIAgentBubble extends ServiceBubble {
                 // Text-only message
                 humanMessage = new HumanMessage(message);
             }
-            const result = await graph.invoke({ messages: [humanMessage] }, { recursionLimit: maxIterations });
+            // Add the current message to the conversation
+            initialMessages.push(humanMessage);
+            const result = await graph.invoke({ messages: initialMessages }, { recursionLimit: maxIterations });
             console.log('[AIAgent] Graph execution completed');
             console.log('[AIAgent] Total messages:', result.messages.length);
             iterations = result.messages.length;
@@ -760,6 +1097,9 @@ export class AIAgentBubble extends ServiceBubble {
             const aiMessages = result.messages.filter((msg) => isAIMessage(msg) || isAIMessageChunk(msg));
             console.log('[AIAgent] Found', aiMessages.length, 'AI messages');
             const finalMessage = aiMessages[aiMessages.length - 1];
+            if (finalMessage?.additional_kwargs?.finishReason === 'SAFETY_BLOCKED') {
+                throw new Error(`[Gemini Error] Unable to generate a response. Please increase maxTokens in model configuration or try again with a different model.`);
+            }
             // Check for MAX_TOKENS finish reason
             if (finalMessage?.additional_kwargs?.finishReason === 'MAX_TOKENS') {
                 throw new Error('Response was truncated due to max tokens limit. Please increase maxTokens in model configuration.');
@@ -785,19 +1125,29 @@ export class AIAgentBubble extends ServiceBubble {
             }
             if (totalTokensSum > 0 && this.context && this.context.logger) {
                 this.context.logger.logTokenUsage({
-                    inputTokens: totalInputTokens,
-                    outputTokens: totalOutputTokens,
-                    totalTokens: totalTokensSum,
-                    modelName: this.params.model.model,
-                }, `LLM completion: ${totalInputTokens} input + ${totalOutputTokens} output = ${totalTokensSum} total tokens`, {
+                    usage: totalInputTokens,
+                    service: this.getCredentialTypeForModel(modelConfig.model),
+                    unit: 'input_tokens',
+                    subService: modelConfig.model,
+                }, `LLM completion: ${totalInputTokens} input`, {
+                    bubbleName: 'ai-agent',
+                    variableId: this.context?.variableId,
+                    operationType: 'bubble_execution',
+                });
+                this.context.logger.logTokenUsage({
+                    usage: totalOutputTokens,
+                    service: this.getCredentialTypeForModel(modelConfig.model),
+                    unit: 'output_tokens',
+                    subService: modelConfig.model,
+                }, `LLM completion: ${totalOutputTokens} output`, {
                     bubbleName: 'ai-agent',
                     variableId: this.context?.variableId,
                     operationType: 'bubble_execution',
                 });
             }
-            const response = finalMessage?.content || 'No response generated';
+            const response = finalMessage?.content || '';
             // Use shared formatting method
-            const formattedResult = await formatFinalResponse(response, this.params.model.model, jsonMode);
+            const formattedResult = formatFinalResponse(response, modelConfig.model, jsonMode);
             // If there's an error from formatting (e.g., invalid JSON), return early
             if (formattedResult.error) {
                 return {
@@ -828,6 +1178,21 @@ export class AIAgentBubble extends ServiceBubble {
             console.warn('[AIAgent] Execution error (continuing):', error);
             console.log('[AIAgent] Tool calls before error:', toolCalls.length);
             console.log('[AIAgent] Iterations before error:', iterations);
+            // Model fallback logic - only retry if this config has a backup model
+            if (modelConfig.backupModel) {
+                console.log(`[AIAgent] Retrying with backup model: ${modelConfig.backupModel.model}`);
+                this.context?.logger?.warn(`Primary model ${modelConfig.model} failed: ${error instanceof Error ? error.message : 'Unknown error'}. Retrying with backup model... ${modelConfig.backupModel.model}`);
+                this.streamingCallback?.({
+                    type: 'error',
+                    data: {
+                        error: `Primary model ${modelConfig.model} failed: ${error instanceof Error ? error.message : 'Unknown error'}. Retrying with backup model... ${modelConfig.backupModel.model}`,
+                        recoverable: true,
+                    },
+                });
+                const backupModelConfig = this.buildModelConfig(modelConfig, modelConfig.backupModel);
+                const backupResult = await this.executeWithModel(backupModelConfig);
+                return backupResult;
+            }
             const errorMessage = error instanceof Error ? error.message : 'Unknown error';
             // Return partial results to allow execution to continue
             // Include any tool calls that were completed before the error
@@ -840,236 +1205,5 @@ export class AIAgentBubble extends ServiceBubble {
             };
         }
     }
-    /**
-     * Execute agent with streaming support using LangGraph streamEvents
-     */
-    async executeAgentWithStreaming(graph, message, images, maxIterations, jsonMode, streamingCallback) {
-        const toolCalls = [];
-        let iterations = 0;
-        let currentMessageId = '';
-        console.log('[AIAgent] Starting streaming execution with message:', message.substring(0, 100) + '...');
-        try {
-            // Create human message with text and optional images
-            let humanMessage;
-            if (images && images.length > 0) {
-                console.log('[AIAgent] Creating multimodal message with', images.length, 'images');
-                // Create multimodal content array
-                const content = [{ type: 'text', text: message }];
-                // Add images to content
-                for (const image of images) {
-                    let imageUrl;
-                    if (image.type === 'base64') {
-                        // Base64 encoded image
-                        imageUrl = `data:${image.mimeType};base64,${image.data}`;
-                    }
-                    else {
-                        // URL image - fetch and convert to base64 for Google Gemini compatibility
-                        try {
-                            console.log('[AIAgent] Fetching image from URL:', image.url);
-                            const response = await fetch(image.url);
-                            if (!response.ok) {
-                                throw new Error(`Failed to fetch image: ${response.status} ${response.statusText}`);
-                            }
-                            const arrayBuffer = await response.arrayBuffer();
-                            const base64Data = Buffer.from(arrayBuffer).toString('base64');
-                            // Detect MIME type from response or default to PNG
-                            const contentType = response.headers.get('content-type') || 'image/png';
-                            imageUrl = `data:${contentType};base64,${base64Data}`;
-                            console.log('[AIAgent] Successfully converted URL image to base64');
-                        }
-                        catch (error) {
-                            console.error('[AIAgent] Error fetching image from URL:', error);
-                            throw new Error(`Failed to load image from URL ${image.url}: ${error instanceof Error ? error.message : 'Unknown error'}`);
-                        }
-                    }
-                    content.push({
-                        type: 'image_url',
-                        image_url: { url: imageUrl },
-                    });
-                    // Add image description if provided
-                    if (image.description) {
-                        content.push({
-                            type: 'text',
-                            text: `Image description: ${image.description}`,
-                        });
-                    }
-                }
-                humanMessage = new HumanMessage({ content });
-            }
-            else {
-                // Text-only message
-                humanMessage = new HumanMessage(message);
-            }
-            // Stream events from the graph
-            const eventStream = graph.streamEvents({ messages: [humanMessage] }, {
-                version: 'v2',
-                recursionLimit: maxIterations,
-            });
-            let currentIteration = 0;
-            const toolCallMap = new Map();
-            let accumulatedContent = '';
-            // Track processed events to prevent duplicates
-            const processedIterationEvents = new Set();
-            for await (const event of eventStream) {
-                if (!event || typeof event !== 'object')
-                    continue;
-                // Handle different types of streaming events
-                switch (event.event) {
-                    case 'on_chat_model_start':
-                        currentIteration++;
-                        currentMessageId = `msg-${Date.now()}-${currentIteration}`;
-                        if (streamingCallback) {
-                            await streamingCallback({
-                                type: 'iteration_start',
-                                data: { iteration: currentIteration },
-                            });
-                        }
-                        break;
-                    case 'on_chat_model_stream':
-                        // Stream individual tokens
-                        if (event.data?.chunk?.content && streamingCallback) {
-                            const content = event.data.chunk.content;
-                            accumulatedContent += content;
-                            await streamingCallback({
-                                type: 'token',
-                                data: {
-                                    content,
-                                    messageId: currentMessageId,
-                                },
-                            });
-                        }
-                        break;
-                    case 'on_chat_model_end':
-                        if (streamingCallback) {
-                            const usageMetadata = event.data?.output?.usage_metadata;
-                            const totalTokens = usageMetadata?.total_tokens;
-                            // Track token usage if available
-                            if (usageMetadata &&
-                                this.context != null &&
-                                this.context.logger != null) {
-                                const tokenUsage = {
-                                    inputTokens: usageMetadata.input_tokens || 0,
-                                    outputTokens: usageMetadata.output_tokens || 0,
-                                    totalTokens: totalTokens || 0,
-                                    modelName: this.params.model.model,
-                                };
-                                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',
-                                });
-                            }
-                            await streamingCallback({
-                                type: 'llm_complete',
-                                data: {
-                                    messageId: currentMessageId,
-                                    totalTokens,
-                                },
-                            });
-                        }
-                        break;
-                    case 'on_tool_start':
-                        if (event.name && event.data?.input && streamingCallback) {
-                            const callId = `tool-${Date.now()}-${Math.random().toString(36).substring(2, 9)}`;
-                            toolCallMap.set(callId, {
-                                name: event.name,
-                                args: event.data.input,
-                                startTime: Date.now(),
-                            });
-                            await streamingCallback({
-                                type: 'tool_start',
-                                data: {
-                                    tool: event.name,
-                                    input: event.data.input,
-                                    callId,
-                                },
-                            });
-                        }
-                        break;
-                    case 'on_tool_end':
-                        if (event.name && event.data?.output && streamingCallback) {
-                            // Find matching tool call
-                            const matchingCall = Array.from(toolCallMap.entries()).find(([, callData]) => callData.name === event.name);
-                            if (matchingCall) {
-                                const [callId, callData] = matchingCall;
-                                const duration = Date.now() - callData.startTime;
-                                toolCalls.push({
-                                    tool: callData.name,
-                                    input: callData.args,
-                                    output: event.data.output,
-                                });
-                                await streamingCallback({
-                                    type: 'tool_complete',
-                                    data: {
-                                        callId,
-                                        input: callData.args,
-                                        tool: callData.name,
-                                        output: event.data.output,
-                                        duration,
-                                    },
-                                });
-                                toolCallMap.delete(callId);
-                            }
-                        }
-                        break;
-                    case 'on_chain_end':
-                        // This indicates the completion of the entire graph
-                        if (event.data?.output) {
-                            iterations = currentIteration;
-                            // Prevent duplicate iteration_complete events
-                            const iterationKey = `iteration_${currentIteration}`;
-                            if (streamingCallback &&
-                                !processedIterationEvents.has(iterationKey)) {
-                                processedIterationEvents.add(iterationKey);
-                                await streamingCallback({
-                                    type: 'iteration_complete',
-                                    data: {
-                                        iteration: currentIteration,
-                                        hasToolCalls: toolCalls.length > 0,
-                                    },
-                                });
-                            }
-                        }
-                        break;
-                }
-            }
-            // Process final result
-            const accumulatedResponse = accumulatedContent || 'No response generated';
-            // Use shared formatting method
-            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 {
-                    response: formattedResult.response,
-                    toolCalls: toolCalls.length > 0 ? toolCalls : [],
-                    iterations,
-                    error: formattedResult.error,
-                    success: false,
-                };
-            }
-            const finalResponse = formattedResult.response;
-            console.log('[AIAgent] Streaming execution completed with', iterations, 'iterations and', toolCalls.length, 'tool calls');
-            return {
-                response: typeof finalResponse === 'string'
-                    ? finalResponse
-                    : JSON.stringify(finalResponse),
-                toolCalls: toolCalls.length > 0 ? toolCalls : [],
-                iterations,
-                error: '',
-                success: true,
-            };
-        }
-        catch (error) {
-            console.warn('[AIAgent] Streaming execution error (continuing):', error);
-            const errorMessage = error instanceof Error ? error.message : 'Unknown error';
-            return {
-                response: `Execution error: ${errorMessage}`,
-                success: false, // Still false but don't completely halt execution
-                iterations,
-                toolCalls: toolCalls.length > 0 ? toolCalls : [], // Preserve completed tool calls
-                error: errorMessage,
-            };
-        }
-    }
 }
 //# sourceMappingURL=ai-agent.js.map