Complete python gRPC requests when closing

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

Hide file tree

Showing 11 changed files with 45 additions and 31 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
@@ -95,7 +95,7 @@ public void start() throws Exception {
     }
 
     @Override
-    public void setContext(AgentContext context) throws Exception {
+    public void setContext(AgentContext context) {
         this.agentContext = context;
     }
 
@@ -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
@@ -101,20 +101,13 @@ public void process(List<ai.langstream.api.runner.code.Record> records, RecordSi
                         records.forEach(recordSink::emitEmptyList);
                     } else {
                         records.forEach(
-                                e -> {
-                                    recordSink.emitError(e, stopped);
-                                });
+                                e -> recordSink.emitError(e, stopped));
                     }
                 }
             }
         }
     }
 
-    @Override
-    public synchronized void close() throws Exception {
-        stopBeforeRestart();
-    }
-
     private SourceRecordAndResult fromGrpc(
             ai.langstream.api.runner.code.Record sourceRecord, ProcessorResult result)
             throws IOException {
@@ -202,8 +195,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/...m-agents/langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/GrpcAgentSink.java b/...m-agents/langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/GrpcAgentSink.java
@@ -144,8 +144,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/...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
@@ -72,7 +72,7 @@ public List<Record> read() throws Exception {
     }
 
     @Override
-    public void permanentFailure(Record record, Exception error) throws Exception {
+    public void permanentFailure(Record record, Exception error) {
         if (record instanceof GrpcAgentRecord grpcAgentRecord) {
             request.onNext(
                     SourceRequest.newBuilder()
@@ -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();
+                        }
                     };
                 }
             };
@@ -272,11 +274,12 @@ void testAvroAndSchema() throws Exception {
 
     @Test
     void testInfo() throws Exception {
-        GrpcAgentProcessor processor = new GrpcAgentProcessor(channel);
-        processor.setContext(new TestAgentContext());
-        processor.start();
-        Map<String, Object> info = processor.buildAdditionalInfo();
-        assertEquals("test-info-value", info.get("test-info-key"));
+        try (GrpcAgentProcessor processor = new GrpcAgentProcessor(channel)) {
+            processor.setContext(new TestAgentContext());
+            processor.start();
+            Map<String, Object> info = processor.buildAdditionalInfo();
+            assertEquals("test-info-value", info.get("test-info-key"));
+        }
     }
 
     private static void assertProcessSuccessful(GrpcAgentProcessor processor, Record inputRecord)

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 {}
 }