Complete python gRPC requests when closing

LangStream · Nov 6, 2023 · 230c202 · 230c202
1 parent 6ca69e9
commit 230c202
Show file tree

Hide file tree

Showing 10 changed files with 32 additions and 17 deletions.
diff --git a/...ents/langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/AbstractGrpcAgent.java b/...ents/langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/AbstractGrpcAgent.java
@@ -128,6 +128,7 @@ public void stopChannel(boolean wait) throws Exception {
 
     @Override
     public synchronized void close() throws Exception {
+        stopBeforeRestart();
         stopChannel(true);
     }
 

diff --git a/...nts/langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/GrpcAgentProcessor.java b/...nts/langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/GrpcAgentProcessor.java
@@ -110,11 +110,6 @@ public void process(List<ai.langstream.api.runner.code.Record> records, RecordSi
         }
     }
 
-    @Override
-    public synchronized void close() throws Exception {
-        stopBeforeRestart();
-    }
-
     private SourceRecordAndResult fromGrpc(
             ai.langstream.api.runner.code.Record sourceRecord, ProcessorResult result)
             throws IOException {

diff --git a/...agents/langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/GrpcAgentSource.java b/...agents/langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/GrpcAgentSource.java
@@ -156,8 +156,8 @@ protected void stopBeforeRestart() throws Exception {
             if (request != null) {
                 try {
                     request.onCompleted();
-                } catch (IllegalStateException ignored) {
-                    log.info("Ignoring error while stopping {}", ignored + "");
+                } catch (IllegalStateException e) {
+                    log.info("Ignoring error while stopping {}", e + "");
                 }
             }
         }

diff --git a/...ngstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/PythonGrpcAgentProcessor.java b/...ngstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/PythonGrpcAgentProcessor.java
@@ -47,12 +47,17 @@ public void start() throws Exception {
 
     @Override
     public synchronized void close() throws Exception {
-        if (server != null) server.close(false);
+        super.close();
+        if (server != null) {
+            server.close(false);
+        }
     }
 
     @Override
     protected synchronized void stopBeforeRestart() throws Exception {
         super.stopBeforeRestart();
-        if (server != null) server.close(true);
+        if (server != null) {
+            server.close(true);
+        }
     }
 }
diff --git a/...ts/langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/PythonGrpcAgentSink.java b/...ts/langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/PythonGrpcAgentSink.java
@@ -45,13 +45,17 @@ public void start() throws Exception {
 
     @Override
     public synchronized void close() throws Exception {
-        if (server != null) server.close(false);
         super.close();
+        if (server != null) {
+            server.close(false);
+        }
     }
 
     @Override
     protected synchronized void stopBeforeRestart() throws Exception {
         super.stopBeforeRestart();
-        if (server != null) server.close(true);
+        if (server != null) {
+            server.close(true);
+        }
     }
 }
diff --git a/.../langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/PythonGrpcAgentSource.java b/.../langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/PythonGrpcAgentSource.java
@@ -45,13 +45,17 @@ public void start() throws Exception {
 
     @Override
     public synchronized void close() throws Exception {
-        if (server != null) server.close(false);
         super.close();
+        if (server != null) {
+            server.close(false);
+        }
     }
 
     @Override
     protected synchronized void stopBeforeRestart() throws Exception {
         super.stopBeforeRestart();
-        if (server != null) server.close(true);
+        if (server != null) {
+            server.close(true);
+        }
     }
 }
diff --git a/...langstream-agent-grpc/src/test/java/ai/langstream/agents/grpc/GrpcAgentProcessorTest.java b/...langstream-agent-grpc/src/test/java/ai/langstream/agents/grpc/GrpcAgentProcessorTest.java
@@ -128,7 +128,9 @@ public void onNext(ProcessorRequest request) {
                         public void onError(Throwable throwable) {}
 
                         @Override
-                        public void onCompleted() {}
+                        public void onCompleted() {
+                            response.onCompleted();
+                        }
                     };
                 }
             };

diff --git a/...ents/langstream-agent-grpc/src/test/java/ai/langstream/agents/grpc/GrpcAgentSinkTest.java b/...ents/langstream-agent-grpc/src/test/java/ai/langstream/agents/grpc/GrpcAgentSinkTest.java
@@ -176,7 +176,9 @@ public void onNext(SinkRequest request) {
                 public void onError(Throwable throwable) {}
 
                 @Override
-                public void onCompleted() {}
+                public void onCompleted() {
+                    responseObserver.onCompleted();
+                }
             };
         }
     }

diff --git a/...ts/langstream-agent-grpc/src/test/java/ai/langstream/agents/grpc/GrpcAgentSourceTest.java b/...ts/langstream-agent-grpc/src/test/java/ai/langstream/agents/grpc/GrpcAgentSourceTest.java
@@ -227,7 +227,9 @@ public void onNext(SourceRequest request) {
                 public void onError(Throwable throwable) {}
 
                 @Override
-                public void onCompleted() {}
+                public void onCompleted() {
+                    responseObserver.onCompleted();
+                }
             };
         }
     }

diff --git a/langstream-api/src/main/java/ai/langstream/api/runner/code/AgentCode.java b/langstream-api/src/main/java/ai/langstream/api/runner/code/AgentCode.java
@@ -66,7 +66,7 @@ default void close() throws Exception {}
     /**
      * Gracefully restart the agent.
      *
-     * @throws Exception
+     * @throws Exception if an error occurs
      */
     default void restart() throws Exception {}
 }