CallAILangChainChatModel works

Author: treblereel

experimental/ai/impl/src/main/java/io/serverlessworkflow/impl/executors/ai/AIChatModelCallExecutor.java

@@ -18,6 +18,7 @@
 
 import io.serverlessworkflow.ai.api.types.CallAILangChainChatModel;
 import io.serverlessworkflow.api.types.TaskBase;
+import io.serverlessworkflow.api.types.ai.AbstractCallAIChatModelTask;
 import io.serverlessworkflow.api.types.ai.CallAIChatModel;
 import io.serverlessworkflow.impl.TaskContext;
 import io.serverlessworkflow.impl.WorkflowApplication;
@@ -28,10 +29,11 @@
 import io.serverlessworkflow.impl.resources.ResourceLoader;
 import java.util.concurrent.CompletableFuture;
 
-public class AIChatModelCallExecutor implements CallableTask<CallAIChatModel> {
+public class AIChatModelCallExecutor implements CallableTask<AbstractCallAIChatModelTask> {
 
   @Override
-  public void init(CallAIChatModel task, WorkflowApplication application, ResourceLoader loader) {}
+  public void init(
+      AbstractCallAIChatModelTask task, WorkflowApplication application, ResourceLoader loader) {}
 
   @Override
   public CompletableFuture<WorkflowModel> apply(
@@ -54,6 +56,6 @@ public CompletableFuture<WorkflowModel> apply(
 
   @Override
   public boolean accept(Class<? extends TaskBase> clazz) {
-    return CallAIChatModel.class.isAssignableFrom(clazz);
+    return AbstractCallAIChatModelTask.class.isAssignableFrom(clazz);
   }
 }

experimental/ai/impl/src/main/java/io/serverlessworkflow/impl/executors/ai/AbstractCallAIChatModelExecutor.java

@@ -16,46 +16,7 @@
 
 package io.serverlessworkflow.impl.executors.ai;
 
-import dev.langchain4j.data.message.AiMessage;
-import dev.langchain4j.model.chat.response.ChatResponse;
-import dev.langchain4j.model.output.FinishReason;
-import dev.langchain4j.model.output.TokenUsage;
-import java.util.Map;
-
 public abstract class AbstractCallAIChatModelExecutor<T> {
 
   public abstract Object apply(T callAIChatModel, Object javaObject);
-
-  protected Map<String, Object> prepareResponse(ChatResponse response, Object javaObject) {
-    String id = response.id();
-    String modelName = response.modelName();
-    TokenUsage tokenUsage = response.tokenUsage();
-    FinishReason finishReason = response.finishReason();
-    AiMessage aiMessage = response.aiMessage();
-
-    Map<String, Object> responseMap = (Map<String, Object>) javaObject;
-    if (response.id() != null) {
-      responseMap.put("id", id);
-    }
-
-    if (modelName != null) {
-      responseMap.put("modelName", modelName);
-    }
-
-    if (tokenUsage != null) {
-      responseMap.put("tokenUsage.inputTokenCount", tokenUsage.inputTokenCount());
-      responseMap.put("tokenUsage.outputTokenCount", tokenUsage.outputTokenCount());
-      responseMap.put("tokenUsage.totalTokenCount", tokenUsage.totalTokenCount());
-    }
-
-    if (finishReason != null) {
-      responseMap.put("finishReason", finishReason.name());
-    }
-
-    if (aiMessage != null) {
-      responseMap.put("text", aiMessage.text());
-    }
-
-    return responseMap;
-  }
 }

experimental/ai/impl/src/main/java/io/serverlessworkflow/impl/executors/ai/CallAIChatModelExecutor.java

@@ -16,10 +16,14 @@
 
 package io.serverlessworkflow.impl.executors.ai;
 
+import dev.langchain4j.data.message.AiMessage;
 import dev.langchain4j.data.message.ChatMessage;
 import dev.langchain4j.data.message.SystemMessage;
 import dev.langchain4j.data.message.UserMessage;
 import dev.langchain4j.model.chat.ChatModel;
+import dev.langchain4j.model.chat.response.ChatResponse;
+import dev.langchain4j.model.output.FinishReason;
+import dev.langchain4j.model.output.TokenUsage;
 import io.serverlessworkflow.api.types.ai.CallAIChatModel;
 import io.serverlessworkflow.impl.services.ChatModelService;
 import java.util.ArrayList;
@@ -105,4 +109,37 @@ private static Set<String> extractVariables(String template) {
   private void validate(CallAIChatModel callAIChatModel, Object javaObject) {
     // TODO
   }
+
+  protected Map<String, Object> prepareResponse(ChatResponse response, Object javaObject) {
+    String id = response.id();
+    String modelName = response.modelName();
+    TokenUsage tokenUsage = response.tokenUsage();
+    FinishReason finishReason = response.finishReason();
+    AiMessage aiMessage = response.aiMessage();
+
+    Map<String, Object> responseMap = (Map<String, Object>) javaObject;
+    if (response.id() != null) {
+      responseMap.put("id", id);
+    }
+
+    if (modelName != null) {
+      responseMap.put("modelName", modelName);
+    }
+
+    if (tokenUsage != null) {
+      responseMap.put("tokenUsage.inputTokenCount", tokenUsage.inputTokenCount());
+      responseMap.put("tokenUsage.outputTokenCount", tokenUsage.outputTokenCount());
+      responseMap.put("tokenUsage.totalTokenCount", tokenUsage.totalTokenCount());
+    }
+
+    if (finishReason != null) {
+      responseMap.put("finishReason", finishReason.name());
+    }
+
+    if (aiMessage != null) {
+      responseMap.put("text", aiMessage.text());
+    }
+
+    return responseMap;
+  }
 }

experimental/ai/impl/src/main/java/io/serverlessworkflow/impl/executors/ai/CallAILangChainChatModelExecutor.java

@@ -17,7 +17,6 @@
 package io.serverlessworkflow.impl.executors.ai;
 
 import dev.langchain4j.model.chat.ChatModel;
-import dev.langchain4j.model.chat.response.ChatResponse;
 import dev.langchain4j.service.AiServices;
 import dev.langchain4j.service.V;
 import io.serverlessworkflow.ai.api.types.CallAILangChainChatModel;
@@ -49,8 +48,10 @@ public Object apply(CallAILangChainChatModel callAIChatModel, Object javaObject)
       }
 
       Object response = method.invoke(aiServices, args);
-      if (response instanceof ChatResponse chatResponse) {
-        return prepareResponse(chatResponse, substitutions);
+
+      if (response instanceof String chatResponse) {
+        substitutions.put("text", chatResponse);
+        return substitutions;
       } else {
         throw new IllegalArgumentException(
             "Method " + method.getName() + " did not return a ChatResponse");
@@ -65,7 +66,7 @@ private void validate(
 
   private Method getMethod(Class<?> chatModelRequest, String methodName) {
     for (Method method : chatModelRequest.getMethods()) {
-      if (method.getName().equals("methodName")) {
+      if (method.getName().equals(methodName)) {
         return method;
       }
     }
@@ -76,20 +77,16 @@ private Method getMethod(Class<?> chatModelRequest, String methodName) {
   private List<String> resolvedParameters(Method method, Map<String, Object> substitutions) {
     List<String> resolvedParameters = new ArrayList<>();
     for (Parameter parameter : method.getParameters()) {
-      String paramName = resolveParameter(parameter);
-      if (substitutions.containsKey(paramName)) {
-        resolvedParameters.add(paramName);
-      } else {
-        throw new IllegalArgumentException("Missing substitution for parameter: " + paramName);
+      if (parameter.getAnnotation(V.class) != null) {
+        V v = parameter.getAnnotation(V.class);
+        String paramName = v.value();
+        if (substitutions.containsKey(paramName)) {
+          resolvedParameters.add(paramName);
+        } else {
+          throw new IllegalArgumentException("Missing substitution for parameter: " + paramName);
+        }
       }
     }
     return resolvedParameters;
   }
-
-  private String resolveParameter(Parameter parameter) {
-    if (parameter.getAnnotation(V.class) != null) {
-      return parameter.getName();
-    }
-    return parameter.getName();
-  }
 }