GRPC agents: relax more limits

LangStream · Dec 13, 2023 · c05e0a6 · c05e0a6
1 parent 651eec5
commit c05e0a6
Show file tree

Hide file tree

Showing 5 changed files with 26 additions and 5 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
@@ -99,7 +99,10 @@ public void start() throws Exception {
             throw new IllegalStateException("Channel not initialized");
         }
         blockingStub =
-                AgentServiceGrpc.newBlockingStub(channel).withDeadlineAfter(30, TimeUnit.SECONDS);
+                AgentServiceGrpc.newBlockingStub(channel)
+                        .withMaxInboundMessageSize(Integer.MAX_VALUE)
+                        .withMaxOutboundMessageSize(Integer.MAX_VALUE)
+                        .withDeadlineAfter(30, TimeUnit.SECONDS);
         asyncStub =
                 AgentServiceGrpc.newStub(channel)
                         .withWaitForReady()

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
@@ -59,7 +59,12 @@ public synchronized void onNewSchemaToSend(Schema schema) {
     @Override
     public void start() throws Exception {
         super.start();
-        request = AgentServiceGrpc.newStub(channel).withWaitForReady().process(responseObserver);
+        request =
+                AgentServiceGrpc.newStub(channel)
+                        .withMaxInboundMessageSize(Integer.MAX_VALUE)
+                        .withMaxOutboundMessageSize(Integer.MAX_VALUE)
+                        .withWaitForReady()
+                        .process(responseObserver);
         restarting.set(false);
         startFailedButDevelopmentMode = false;
     }

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
@@ -53,7 +53,12 @@ public void onNewSchemaToSend(Schema schema) {
     @Override
     public void start() throws Exception {
         super.start();
-        request = AgentServiceGrpc.newStub(channel).withWaitForReady().write(responseObserver);
+        request =
+                AgentServiceGrpc.newStub(channel)
+                        .withMaxInboundMessageSize(Integer.MAX_VALUE)
+                        .withMaxOutboundMessageSize(Integer.MAX_VALUE)
+                        .withWaitForReady()
+                        .write(responseObserver);
         restarting.set(false);
         startFailedButDevelopmentMode = false;
     }

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
@@ -53,7 +53,12 @@ public void onNewSchemaToSend(Schema schema) {
     @Override
     public void start() throws Exception {
         super.start();
-        request = AgentServiceGrpc.newStub(channel).withWaitForReady().read(responseObserver);
+        request =
+                AgentServiceGrpc.newStub(channel)
+                        .withMaxInboundMessageSize(Integer.MAX_VALUE)
+                        .withMaxOutboundMessageSize(Integer.MAX_VALUE)
+                        .withWaitForReady()
+                        .read(responseObserver);
         restarting.set(false);
         startFailedButDevelopmentMode = false;
     }

diff --git a/...gents/langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/PythonGrpcServer.java b/...gents/langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/PythonGrpcServer.java
@@ -90,7 +90,10 @@ public ManagedChannel start() throws Exception {
         ManagedChannel channel =
                 ManagedChannelBuilder.forAddress("localhost", port).usePlaintext().build();
         AgentServiceGrpc.AgentServiceBlockingStub stub =
-                AgentServiceGrpc.newBlockingStub(channel).withDeadlineAfter(30, TimeUnit.SECONDS);
+                AgentServiceGrpc.newBlockingStub(channel)
+                        .withMaxInboundMessageSize(Integer.MAX_VALUE)
+                        .withMaxOutboundMessageSize(Integer.MAX_VALUE)
+                        .withDeadlineAfter(30, TimeUnit.SECONDS);
         for (int i = 0; ; i++) {
             try {
                 stub.agentInfo(Empty.getDefaultInstance());