Wrap up integration with langchain4j - everything is cog

Signed-off-by: Ricardo Zanini <ricardozanini@gmail.com>

Author: ricardozanini

experimental/agentic/src/main/java/io/serverlessworkflow/impl/expressions/agentic/AgenticModelFactory.java

@@ -21,28 +21,44 @@
 import io.serverlessworkflow.impl.WorkflowModel;
 import io.serverlessworkflow.impl.WorkflowModelCollection;
 import io.serverlessworkflow.impl.WorkflowModelFactory;
+import io.serverlessworkflow.impl.expressions.agentic.langchain4j.CognisphereRegistryAssessor;
 import io.serverlessworkflow.impl.expressions.func.JavaModel;
 import java.time.OffsetDateTime;
 import java.util.Map;
 
 class AgenticModelFactory implements WorkflowModelFactory {
 
+  /**
+   * Applies any change to the model after running as task. We will always set it to
+   * a @DefaultCognisphere object since @AgentExecutor is always adding the output to the
+   * cognisphere. We just have to make sure that cognisphere is always passed to the next input
+   * task.
+   *
+   * @param prev the global Cognisphere object getting updated by the workflow context
+   * @param obj the same Cognisphere object updated by the AgentExecutor
+   * @return the workflow context model holding the cognisphere object.
+   */
   @Override
   public WorkflowModel fromAny(WorkflowModel prev, Object obj) {
-    ((AgenticModel) prev).setObject(obj);
+    // We ignore `obj` since it's already included in `prev` within the Cognisphere instance
     return prev;
   }
 
   @Override
   public WorkflowModel combine(Map<String, WorkflowModel> workflowVariables) {
-    throw new UnsupportedOperationException();
+    // TODO: create a new cognisphere object in the CognisphereRegistryAssessor per branch
+    // TODO: Since we share the same cognisphere object, both branches are updating the same
+    // instance, so for now we return the first key.
+    return workflowVariables.values().iterator().next();
   }
 
   @Override
   public WorkflowModelCollection createCollection() {
     throw new UnsupportedOperationException();
   }
 
+  // TODO: all these methods can use Cognisphere as long as we have access to the `outputName`
+
   @Override
   public WorkflowModel from(boolean value) {
     return new JavaModel(value);
@@ -75,7 +91,9 @@ public WorkflowModel from(OffsetDateTime value) {
 
   @Override
   public WorkflowModel from(Map<String, Object> map) {
-    return new JavaModel(map);
+    final Cognisphere cognisphere = new CognisphereRegistryAssessor().getCognisphere();
+    cognisphere.writeStates(map);
+    return new AgenticModel(cognisphere);
   }
 
   @Override

experimental/agentic/src/main/java/io/serverlessworkflow/impl/expressions/agentic/langchain4j/CognisphereRegistryAssessor.java

@@ -0,0 +1,70 @@
+/*
+ * Copyright 2020-Present The Serverless Workflow Specification Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.serverlessworkflow.impl.expressions.agentic.langchain4j;
+
+import dev.langchain4j.agentic.cognisphere.CognisphereRegistry;
+import dev.langchain4j.agentic.cognisphere.DefaultCognisphere;
+import dev.langchain4j.agentic.internal.CognisphereOwner;
+import java.util.Objects;
+import java.util.UUID;
+import java.util.concurrent.atomic.AtomicReference;
+
+public class CognisphereRegistryAssessor implements CognisphereOwner {
+
+  private final AtomicReference<CognisphereRegistry> cognisphereRegistry = new AtomicReference<>();
+  private final String agentId;
+  private DefaultCognisphere cognisphere;
+  private Object memoryId;
+
+  public CognisphereRegistryAssessor(String agentId) {
+    Objects.requireNonNull(agentId, "Agent id cannot be null");
+    this.agentId = agentId;
+  }
+
+  // TODO: have access to the workflow definition and assign its name instead
+  public CognisphereRegistryAssessor() {
+    this.agentId = UUID.randomUUID().toString();
+  }
+
+  public void setMemoryId(Object memoryId) {
+    this.memoryId = memoryId;
+  }
+
+  public DefaultCognisphere getCognisphere() {
+    if (cognisphere != null) {
+      return cognisphere;
+    }
+
+    if (memoryId != null) {
+      this.cognisphere = registry().getOrCreate(memoryId);
+    } else {
+      this.cognisphere = registry().createEphemeralCognisphere();
+    }
+    return this.cognisphere;
+  }
+
+  @Override
+  public CognisphereOwner withCognisphere(DefaultCognisphere cognisphere) {
+    this.cognisphere = cognisphere;
+    return this;
+  }
+
+  @Override
+  public CognisphereRegistry registry() {
+    cognisphereRegistry.compareAndSet(null, new CognisphereRegistry(agentId));
+    return cognisphereRegistry.get();
+  }
+}

fluent/agentic-langchain4j/src/main/java/io/serverlessworkflow/fluent/agentic/langchain4j/AbstractAgentService.java

@@ -36,9 +36,7 @@ public abstract class AbstractAgentService<T, S> implements WorkflowDefinitionBu
   protected final Class<T> agentServiceClass;
 
   protected AbstractAgentService(Class<T> agentServiceClass) {
-    this.workflowBuilder =
-        AgentWorkflowBuilder.workflow()
-            .outputAs(DEFAULT_OUTPUT_FUNCTION);
+    this.workflowBuilder = AgentWorkflowBuilder.workflow().outputAs(DEFAULT_OUTPUT_FUNCTION);
     this.agentServiceClass = agentServiceClass;
     this.workflowExecBuilder = WorkflowApplication.builder();
   }
@@ -49,7 +47,8 @@ public T build() {
         Proxy.newProxyInstance(
             this.agentServiceClass.getClassLoader(),
             new Class<?>[] {agentServiceClass, AgentSpecification.class, CognisphereOwner.class},
-            new WorkflowInvocationHandler(this.workflowBuilder.build(), this.workflowExecBuilder, this.agentServiceClass));
+            new WorkflowInvocationHandler(
+                this.workflowBuilder.build(), this.workflowExecBuilder, this.agentServiceClass));
   }
 
   @SuppressWarnings("unchecked")

fluent/agentic-langchain4j/src/main/java/io/serverlessworkflow/fluent/agentic/langchain4j/WorkflowInvocationHandler.java

@@ -18,7 +18,6 @@
 import dev.langchain4j.agentic.UntypedAgent;
 import dev.langchain4j.agentic.cognisphere.Cognisphere;
 import dev.langchain4j.agentic.cognisphere.CognisphereAccess;
-import dev.langchain4j.agentic.cognisphere.CognisphereKey;
 import dev.langchain4j.agentic.cognisphere.CognisphereRegistry;
 import dev.langchain4j.agentic.cognisphere.DefaultCognisphere;
 import dev.langchain4j.agentic.cognisphere.ResultWithCognisphere;
@@ -28,29 +27,26 @@
 import dev.langchain4j.service.MemoryId;
 import io.serverlessworkflow.api.types.Workflow;
 import io.serverlessworkflow.impl.WorkflowApplication;
-import io.serverlessworkflow.impl.WorkflowModel;
-
+import io.serverlessworkflow.impl.expressions.agentic.langchain4j.CognisphereRegistryAssessor;
 import java.lang.reflect.InvocationHandler;
 import java.lang.reflect.Method;
 import java.lang.reflect.Parameter;
 import java.util.Map;
-import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
-import java.util.concurrent.atomic.AtomicReference;
 
 public class WorkflowInvocationHandler implements InvocationHandler, CognisphereOwner {
 
   private final Workflow workflow;
   private final WorkflowApplication.Builder workflowApplicationBuilder;
-  private DefaultCognisphere cognisphere;
-  private Class<?> agentServiceClass;
-  private final AtomicReference<CognisphereRegistry> cognisphereRegistry = new AtomicReference<>();
+  private final CognisphereRegistryAssessor cognisphereRegistryAssessor;
 
   WorkflowInvocationHandler(
-      Workflow workflow, WorkflowApplication.Builder workflowApplicationBuilder, Class<?> agentServiceClass) {
+      Workflow workflow,
+      WorkflowApplication.Builder workflowApplicationBuilder,
+      Class<?> agentServiceClass) {
     this.workflow = workflow;
     this.workflowApplicationBuilder = workflowApplicationBuilder;
-    this.agentServiceClass = agentServiceClass;
+    this.cognisphereRegistryAssessor = new CognisphereRegistryAssessor(agentServiceClass.getName());
   }
 
   @SuppressWarnings("unchecked")
@@ -67,22 +63,26 @@ private static void writeCognisphereState(Cognisphere cognisphere, Method method
     }
   }
 
-  private String agentId() {
-    return workflow.getDocument().getName();
+  private String outputName() {
+    Object outputName =
+        this.workflow
+            .getDocument()
+            .getMetadata()
+            .getAdditionalProperties()
+            .get(WorkflowDefinitionBuilder.META_KEY_OUTPUTNAME);
+    if (outputName != null) {
+      return outputName.toString();
+    }
+    return null;
   }
 
   @Override
   public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
-    CognisphereRegistry registry = cognisphereRegistry();
+    CognisphereRegistry registry = registry();
     // outputName
     if (method.getDeclaringClass() == AgentSpecification.class) {
       return switch (method.getName()) {
-        case "outputName" ->
-            this.workflow
-                .getDocument()
-                .getMetadata()
-                .getAdditionalProperties()
-                .get(WorkflowDefinitionBuilder.META_KEY_OUTPUTNAME);
+        case "outputName" -> outputName();
         default ->
             throw new UnsupportedOperationException(
                 "Unknown method on AgentInstance class : " + method.getName());
@@ -113,39 +113,44 @@ public Object invoke(Object proxy, Method method, Object[] args) throws Throwabl
     }
 
     // invoke
-    return executeWorkflow(currentCognisphere(cognisphereRegistry(), method, args), method, args);
+    return executeWorkflow(currentCognisphere(method, args), method, args);
   }
 
   private Object executeWorkflow(DefaultCognisphere cognisphere, Method method, Object[] args) {
     writeCognisphereState(cognisphere, method, args);
 
     try (WorkflowApplication app = workflowApplicationBuilder.build()) {
-      CompletableFuture<WorkflowModel> workflowInstance = app.workflowDefinition(workflow).instance(cognisphere).start();
-      Object result = workflowInstance.get().asJavaObject();
-      return method.getReturnType().equals(ResultWithCognisphere.class) ?
-              new ResultWithCognisphere<>(cognisphere, result) :
-              result;
+      // TODO improve result handling
+      DefaultCognisphere output =
+          app.workflowDefinition(workflow)
+              .instance(cognisphere)
+              .start()
+              .get()
+              .as(DefaultCognisphere.class)
+              .orElseThrow(
+                  () ->
+                      new IllegalArgumentException(
+                          "Workflow hasn't returned a Cognisphere object."));
+      Object result = output.readState(outputName());
+
+      return method.getReturnType().equals(ResultWithCognisphere.class)
+          ? new ResultWithCognisphere<>(output, result)
+          : result;
 
     } catch (ExecutionException | InterruptedException e) {
       throw new RuntimeException(
-          "Failed to execute workflow: " + agentId() + " - Cognisphere: " + cognisphere, e);
+          "Failed to execute workflow: "
+              + workflow.getDocument().getName()
+              + " - Cognisphere: "
+              + cognisphere,
+          e);
     }
   }
 
-  private CognisphereRegistry cognisphereRegistry() {
-    cognisphereRegistry.compareAndSet(null, new CognisphereRegistry(this.agentServiceClass.getName()));
-    return cognisphereRegistry.get();
-  }
-
-  private DefaultCognisphere currentCognisphere(CognisphereRegistry registry, Method method, Object[] args) {
-    if (cognisphere != null) {
-      return cognisphere;
-    }
-
+  private DefaultCognisphere currentCognisphere(Method method, Object[] args) {
     Object memoryId = memoryId(method, args);
-    return memoryId != null
-            ? registry.getOrCreate(new CognisphereKey(this.agentId(), memoryId))
-            : registry.createEphemeralCognisphere();
+    this.cognisphereRegistryAssessor.setMemoryId(memoryId);
+    return this.cognisphereRegistryAssessor.getCognisphere();
   }
 
   private Object memoryId(Method method, Object[] args) {
@@ -160,12 +165,12 @@ private Object memoryId(Method method, Object[] args) {
 
   @Override
   public CognisphereOwner withCognisphere(DefaultCognisphere cognisphere) {
-    this.cognisphere = cognisphere;
+    this.cognisphereRegistryAssessor.withCognisphere(cognisphere);
     return this;
   }
 
   @Override
   public CognisphereRegistry registry() {
-    return this.cognisphereRegistry();
+    return this.cognisphereRegistryAssessor.registry();
   }
 }