Complete python gRPC requests when closing

LangStream · Nov 6, 2023 · d7f39f6 · d7f39f6
1 parent eec9581
commit d7f39f6
Show file tree

Hide file tree

Showing 8 changed files with 41 additions and 9 deletions.
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
@@ -98,4 +98,12 @@ public void onCompleted() {
             }
         };
     }
+
+    @Override
+    public synchronized void close() throws Exception {
+        if (request != null) {
+            request.onCompleted();
+        }
+        super.close();
+    }
 }
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
@@ -133,4 +133,12 @@ public void onCompleted() {
             }
         };
     }
+
+    @Override
+    public synchronized void close() throws Exception {
+        if (request != null) {
+            request.onCompleted();
+        }
+        super.close();
+    }
 }
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
@@ -39,12 +39,16 @@ public void start() throws Exception {
 
     @Override
     public synchronized void close() throws Exception {
-        if (server != null) server.close(false);
+        if (server != null) {
+            server.close(false);
+        }
     }
 
     @Override
     protected synchronized void stop() throws Exception {
-        if (server != null) server.close(true);
+        if (server != null) {
+            server.close(true);
+        }
     }
 
     @Override

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
@@ -39,13 +39,17 @@ public void start() throws Exception {
 
     @Override
     public synchronized void close() throws Exception {
-        if (server != null) server.close(false);
+        if (server != null) {
+            server.close(false);
+        }
         super.close();
     }
 
     @Override
     protected synchronized void stop() throws Exception {
-        if (server != null) server.close(true);
+        if (server != null) {
+            server.close(true);
+        }
     }
 
     @Override

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
@@ -39,13 +39,17 @@ public void start() throws Exception {
 
     @Override
     public synchronized void close() throws Exception {
-        if (server != null) server.close(false);
+        if (server != null) {
+            server.close(false);
+        }
         super.close();
     }
 
     @Override
     protected synchronized void stop() throws Exception {
-        if (server != null) server.close(true);
+        if (server != null) {
+            server.close(true);
+        }
     }
 
     @Override

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 {
         log.info("Restart is not supported for agent type {}", agentType());