Set gRPC message limit to 2GB (LangStream#746)

vectorize-io · Dec 12, 2023 · 1c4e368 · 1c4e368
1 parent aba5a69
commit 1c4e368
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 3 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
@@ -100,7 +100,11 @@ public void start() throws Exception {
         }
         blockingStub =
                 AgentServiceGrpc.newBlockingStub(channel).withDeadlineAfter(30, TimeUnit.SECONDS);
-        asyncStub = AgentServiceGrpc.newStub(channel).withWaitForReady();
+        asyncStub =
+                AgentServiceGrpc.newStub(channel)
+                        .withWaitForReady()
+                        .withMaxInboundMessageSize(Integer.MAX_VALUE)
+                        .withMaxOutboundMessageSize(Integer.MAX_VALUE);
 
         topicProducerWriteResults = new CompletableFuture<>();
         topicProducerWriteResults.complete(

diff --git a/langstream-runtime/langstream-runtime-impl/src/main/python/langstream_grpc/grpc_service.py b/langstream-runtime/langstream-runtime-impl/src/main/python/langstream_grpc/grpc_service.py
@@ -415,7 +415,12 @@ def get_topic_producer(self) -> TopicProducer:
 class AgentServer(object):
     def __init__(self, target: str):
         self.target = target
-        self.grpc_server = grpc.aio.server()
+        self.grpc_server = grpc.aio.server(
+            options=[
+                ("grpc.max_send_message_length", 0x7FFFFFFF),
+                ("grpc.max_receive_message_length", 0x7FFFFFFF),
+            ]
+        )
         self.port = self.grpc_server.add_insecure_port(target)
         self.agent = None
 

diff --git a/...-runtime/langstream-runtime-impl/src/main/python/langstream_grpc/tests/server_and_stub.py b/...-runtime/langstream-runtime-impl/src/main/python/langstream_grpc/tests/server_and_stub.py
@@ -36,7 +36,13 @@ async def __aenter__(self):
         self.server = AgentServer("[::]:0")
         await self.server.init(json.dumps(self.config), json.dumps(self.context))
         await self.server.start()
-        self.channel = grpc.aio.insecure_channel("localhost:%d" % self.server.port)
+        self.channel = grpc.aio.insecure_channel(
+            "localhost:%d" % self.server.port,
+            options=[
+                ("grpc.max_send_message_length", 0x7FFFFFFF),
+                ("grpc.max_receive_message_length", 0x7FFFFFFF),
+            ],
+        )
         self.stub = AgentServiceStub(channel=self.channel)
         return self
 

diff --git a/...time/langstream-runtime-impl/src/main/python/langstream_grpc/tests/test_grpc_processor.py b/...time/langstream-runtime-impl/src/main/python/langstream_grpc/tests/test_grpc_processor.py
@@ -201,6 +201,25 @@ async def test_future_record(klass):
             assert response.results[0].records[0].value.string_value == "test"
 
 
+async def test_big_record():
+    async with ServerAndStub(
+        "langstream_grpc.tests.test_grpc_processor.MyProcessor"
+    ) as server_and_stub:
+        long_string = "a" * 10_000_000
+        response: ProcessorResponse
+        async for response in server_and_stub.stub.process(
+            [
+                ProcessorRequest(
+                    records=[GrpcRecord(value=Value(string_value=long_string))]
+                )
+            ]
+        ):
+            assert len(response.results) == 1
+            assert response.results[0].HasField("error") is False
+            assert len(response.results[0].records) == 1
+            assert response.results[0].records[0].value.string_value == long_string
+
+
 async def test_info():
     async with ServerAndStub(
         "langstream_grpc.tests.test_grpc_processor.MyProcessor"