Align with latest changes from langchain4j-agentic

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

Author: ricardozanini

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

@@ -16,7 +16,7 @@
 package io.serverlessworkflow.fluent.agentic.langchain4j;
 
 import dev.langchain4j.agentic.cognisphere.Cognisphere;
-import dev.langchain4j.agentic.internal.AgentInstance;
+import dev.langchain4j.agentic.internal.AgentSpecification;
 import dev.langchain4j.agentic.internal.CognisphereOwner;
 import io.serverlessworkflow.api.types.Workflow;
 import io.serverlessworkflow.fluent.agentic.AgentWorkflowBuilder;
@@ -50,8 +50,8 @@ public T build() {
     return (T)
         Proxy.newProxyInstance(
             this.agentServiceClass.getClassLoader(),
-            new Class<?>[] {agentServiceClass, AgentInstance.class, CognisphereOwner.class},
-            new WorkflowInvocationHandler(this.workflowBuilder.build(), this.workflowExecBuilder));
+            new Class<?>[] {agentServiceClass, AgentSpecification.class, CognisphereOwner.class},
+            new WorkflowInvocationHandler(this.workflowBuilder.build(), this.workflowExecBuilder, this.agentServiceClass));
   }
 
   @SuppressWarnings("unchecked")

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

@@ -20,32 +20,38 @@
 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;
-import dev.langchain4j.agentic.internal.AgentInstance;
+import dev.langchain4j.agentic.internal.AgentInvoker;
 import dev.langchain4j.agentic.internal.AgentSpecification;
 import dev.langchain4j.agentic.internal.CognisphereOwner;
 import dev.langchain4j.service.MemoryId;
 import io.serverlessworkflow.api.types.Workflow;
 import io.serverlessworkflow.impl.WorkflowApplication;
 import io.serverlessworkflow.impl.WorkflowModel;
 
+import java.lang.reflect.InvocationHandler;
 import java.lang.reflect.Method;
 import java.lang.reflect.Parameter;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.atomic.AtomicReference;
 
-public class WorkflowInvocationHandler implements CognisphereOwner {
+public class WorkflowInvocationHandler implements InvocationHandler, CognisphereOwner {
 
   private final Workflow workflow;
   private final WorkflowApplication.Builder workflowApplicationBuilder;
-  private Cognisphere cognisphere;
+  private DefaultCognisphere cognisphere;
+  private Class<?> agentServiceClass;
+  private final AtomicReference<CognisphereRegistry> cognisphereRegistry = new AtomicReference<>();
 
   WorkflowInvocationHandler(
-      Workflow workflow, WorkflowApplication.Builder workflowApplicationBuilder) {
+      Workflow workflow, WorkflowApplication.Builder workflowApplicationBuilder, Class<?> agentServiceClass) {
     this.workflow = workflow;
     this.workflowApplicationBuilder = workflowApplicationBuilder;
+    this.agentServiceClass = agentServiceClass;
   }
 
   @SuppressWarnings("unchecked")
@@ -56,7 +62,7 @@ private static void writeCognisphereState(Cognisphere cognisphere, Method method
       Parameter[] parameters = method.getParameters();
       for (int i = 0; i < parameters.length; i++) {
         int index = i;
-        AgentSpecification.optionalParameterName(parameters[i])
+        AgentInvoker.optionalParameterName(parameters[i])
             .ifPresent(argName -> cognisphere.writeState(argName, args[index]));
       }
     }
@@ -70,7 +76,7 @@ private static void writeWorkflowInputState(final Map<String, Object> input, Met
       Parameter[] parameters = method.getParameters();
       for (int i = 0; i < parameters.length; i++) {
         int index = i;
-          AgentSpecification.optionalParameterName(parameters[i])
+        AgentInvoker.optionalParameterName(parameters[i])
                 .ifPresent(argName -> input.put(argName, args[index]));
       }
     }
@@ -82,8 +88,9 @@ private String agentId() {
 
   @Override
   public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
+    CognisphereRegistry registry = cognisphereRegistry();
     // outputName
-    if (method.getDeclaringClass() == AgentInstance.class) {
+    if (method.getDeclaringClass() == AgentSpecification.class) {
       return switch (method.getName()) {
         case "outputName" ->
             this.workflow
@@ -101,7 +108,8 @@ public Object invoke(Object proxy, Method method, Object[] args) throws Throwabl
       // Ingest the workflow input as a Cognisphere object
       // Later, retrieve it and start the workflow with it as input.
       return switch (method.getName()) {
-        case "withCognisphere" -> this.withCognisphere((Cognisphere) args[0]);
+        case "withCognisphere" -> this.withCognisphere((DefaultCognisphere) args[0]);
+        case "registry" -> registry;
         default ->
             throw new UnsupportedOperationException(
                 "Unknown method on CognisphereOwner class : " + method.getName());
@@ -111,10 +119,8 @@ public Object invoke(Object proxy, Method method, Object[] args) throws Throwabl
     // evictCognisphere
     if (method.getDeclaringClass() == CognisphereAccess.class) {
       return switch (method.getName()) {
-        case "getCognisphere" ->
-            CognisphereRegistry.get(new CognisphereKey(this.agentId(), args[0]));
-        case "evictCognisphere" ->
-            CognisphereRegistry.evict(new CognisphereKey(this.agentId(), args[0]));
+        case "getCognisphere" -> registry().get(args[0]);
+        case "evictCognisphere" -> registry().evict(args[0]);
         default ->
             throw new UnsupportedOperationException(
                 "Unknown method on CognisphereAccess class : " + method.getName());
@@ -126,10 +132,6 @@ public Object invoke(Object proxy, Method method, Object[] args) throws Throwabl
   }
 
   private Object executeWorkflow(Method method, Object[] args) {
-    // TODO: actually, we must own the Cognisphere object creation upon calling the workflow
-
-    //writeCognisphereState(cognisphere, method, args);
-
     Object input;
     if (args == null || args.length == 0) {
       input = new HashMap<>();
@@ -155,21 +157,20 @@ private Object executeWorkflow(Method method, Object[] args) {
     }
   }
 
-  @Override
-  public CognisphereOwner withCognisphere(Cognisphere cognisphere) {
-    this.cognisphere = cognisphere;
-    return this;
+  private CognisphereRegistry cognisphereRegistry() {
+    cognisphereRegistry.compareAndSet(null, new CognisphereRegistry(this.agentServiceClass.getName()));
+    return cognisphereRegistry.get();
   }
 
-  private Cognisphere currentCognisphere(Method method, Object[] args) {
+  private Cognisphere currentCognisphere(CognisphereRegistry registry, Method method, Object[] args) {
     if (cognisphere != null) {
       return cognisphere;
     }
 
     Object memoryId = memoryId(method, args);
     return memoryId != null
-        ? CognisphereRegistry.getOrCreate(new CognisphereKey(this.agentId(), memoryId))
-        : CognisphereRegistry.createEphemeralCognisphere();
+            ? registry.getOrCreate(new CognisphereKey(this.agentId(), memoryId))
+            : registry.createEphemeralCognisphere();
   }
 
   private Object memoryId(Method method, Object[] args) {
@@ -181,4 +182,15 @@ private Object memoryId(Method method, Object[] args) {
     }
     return null;
   }
+
+  @Override
+  public CognisphereOwner withCognisphere(DefaultCognisphere cognisphere) {
+    this.cognisphere = cognisphere;
+    return this;
+  }
+
+  @Override
+  public CognisphereRegistry registry() {
+    return this.cognisphereRegistry();
+  }
 }