refactor: finalize AsyncIterator support

experimental feature
bsorrentino · Mar 25, 2024 · f404e50 · f404e50
1 parent 19b43fd
commit f404e50
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 34 deletions.
diff --git a/src/main/java/org/bsc/langgraph4j/async/AsyncIterator.java b/src/main/java/org/bsc/langgraph4j/async/AsyncIterator.java
@@ -9,31 +9,18 @@
 
 public interface AsyncIterator<T> extends Iterable<T> {
 
-    record Data<T>(T data, boolean done, Throwable error) {
-        public Data(T data, boolean done) {
-            this(data, done, null );
-        }
-        public Data(Throwable error) {
-            this(null, false, error );
-        }
-    }
+    record Data<T>(T data, boolean done) {}
 
     CompletableFuture<Data<T>> next();
 
     default CompletableFuture<Void> forEachAsync(  final AsyncFunction<T,Void> consumer) {
 
         return next().thenCompose(data -> {
-                    if (data.error != null  ) {
-                        var error = new CompletableFuture<Void>();
-                        error.completeExceptionally(data.error);
-                        return error;
-                    }
                     if (data.done) {
                         return completedFuture(null);
                     }
                     return consumer.apply(data.data)
                             .thenCompose( v -> forEachAsync(consumer) );
-
                 });
     }
 
@@ -46,27 +33,19 @@ public boolean hasNext() {
                 if (currentFetchedData.get() != null) {
                     return false;
                 }
-
-                var next =  currentFetchedData.updateAndGet( (v) -> AsyncIterator.this.next().join() );
-
-                return !next.done();
+                return !currentFetchedData.updateAndGet( (v) -> AsyncIterator.this.next().join() ).done();
             }
 
             @Override
             public T next() {
-
                 if (currentFetchedData.get() == null) {
                     if( !hasNext() ) {
                         throw new NoSuchElementException("no more elements into iterator");
                     }
                 }
-                if (currentFetchedData.get().error() != null ) {
-                    throw new IllegalStateException(currentFetchedData.get().error());
-                }
                 if (currentFetchedData.get().done()) {
                     throw new NoSuchElementException("no more elements into iterator");
                 }
-
                 return currentFetchedData.getAndUpdate((v) -> null).data();
             }
         };

diff --git a/src/main/java/org/bsc/langgraph4j/async/AsyncQueue.java b/src/main/java/org/bsc/langgraph4j/async/AsyncQueue.java
@@ -26,24 +26,20 @@ public AsyncQueue(Executor executor) {
      * @throws InterruptedException if interrupted while waiting for space to become available
      */
     public void put(E e) throws InterruptedException {
-        if( exception.get() != null ) {
-            throw new IllegalStateException("Queue has been closed with exception!");
-        }
+        Objects.requireNonNull(queue);
         queue.put(new Data<>(e, false));
     }
 
     public void closeExceptionally(Throwable ex) {
+        Objects.requireNonNull(queue);
         exception.set(ex);
     }
 
     @Override
     public void close() throws Exception {
-        if (exception.get() != null) {
-            queue.put(new Data<>(exception.get()));
-        }
-        else {
-            queue.put(new Data<>(null, true));
-        }
+        Objects.requireNonNull(queue);
+        queue.put(new Data<>(null, true));
+        queue = null;
     }
 
 
@@ -52,8 +48,8 @@ public CompletableFuture<Data<E>> next() {
         return CompletableFuture.supplyAsync( () -> {
             try {
                 var result = queue.take();
-                if( result.error() != null ) {
-                    queue = null;
+                if( exception.get()!=null ) {
+                    throw new RuntimeException(exception.get());
                 }
                 return result;
             } catch (InterruptedException e) {