diff --git a/experimental/ai/impl/src/main/java/io/serverlessworkflow/impl/executors/ai/AIChatModelCallExecutor.java b/experimental/ai/impl/src/main/java/io/serverlessworkflow/impl/executors/ai/AIChatModelCallExecutor.java
index 7d64bb9a..28d80eff 100644
--- a/experimental/ai/impl/src/main/java/io/serverlessworkflow/impl/executors/ai/AIChatModelCallExecutor.java
+++ b/experimental/ai/impl/src/main/java/io/serverlessworkflow/impl/executors/ai/AIChatModelCallExecutor.java
@@ -29,20 +29,19 @@
 
 public class AIChatModelCallExecutor implements CallableTask<CallAgentAI> {
 
+  private CallAgentAIExecutor executor;
+  
   @Override
-  public void init(CallAgentAI task, WorkflowApplication application, ResourceLoader loader) {}
+  public void init(CallAgentAI task, WorkflowApplication application, ResourceLoader loader) {
+	  executor = new CallAgentAIExecutor(task);
+  }
 
   @Override
   public CompletableFuture<WorkflowModel> apply(
       WorkflowContext workflowContext, TaskContext taskContext, WorkflowModel input) {
     WorkflowModelFactory modelFactory = workflowContext.definition().application().modelFactory();
-    if (taskContext.task() instanceof CallAgentAI agenticAI) {
-      return CompletableFuture.completedFuture(
-          modelFactory.fromAny(new CallAgentAIExecutor().execute(agenticAI, input.asJavaObject())));
-    }
-    throw new IllegalArgumentException(
-        "AIChatModelCallExecutor can only process CallAgentAI tasks, but received: "
-            + taskContext.task().getClass().getName());
+    return CompletableFuture.completedFuture(
+          modelFactory.fromAny(executor.execute(input.asJavaObject())));
   }
 
   @Override
diff --git a/experimental/ai/impl/src/main/java/io/serverlessworkflow/impl/executors/ai/CallAgentAIExecutor.java b/experimental/ai/impl/src/main/java/io/serverlessworkflow/impl/executors/ai/CallAgentAIExecutor.java
index b13ab474..1f916468 100644
--- a/experimental/ai/impl/src/main/java/io/serverlessworkflow/impl/executors/ai/CallAgentAIExecutor.java
+++ b/experimental/ai/impl/src/main/java/io/serverlessworkflow/impl/executors/ai/CallAgentAIExecutor.java
@@ -24,9 +24,15 @@
 
 public class CallAgentAIExecutor {
 
+  private final CallAgentAI callAgentAI;
+  
   private static final Cognisphere cognisphere = new Cognisphere();
 
-  public Object execute(CallAgentAI callAgentAI, Object input) {
+  public CallAgentAIExecutor(CallAgentAI callAgentAI) {
+	  this.callAgentAI = callAgentAI;
+}
+
+  public Object execute( Object input) {
     AgentExecutor agentExecutor = AgentUtil.agentToExecutor(callAgentAI.getAgentInstance());
 
     Map<String, Object> output = (Map<String, Object>) input;