chalk-remote-call-python 1.2.0__tar.gz → 1.3.0__tar.gz

{chalk_remote_call_python-1.2.0 → chalk_remote_call_python-1.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chalk-remote-call-python
-Version: 1.2.0
+Version: 1.3.0
 Summary: Chalk remote call Python runtime interface client
 Author: Chalk AI, Inc.
 Project-URL: Homepage, https://chalk.ai

chalk_remote_call_python-1.3.0/chalk_remote_call/_version.py

@@ -0,0 +1 @@
+__version__ = "1.3.0"

{chalk_remote_call_python-1.2.0 → chalk_remote_call_python-1.3.0}/chalk_remote_call/servicer.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
+import asyncio
 import dataclasses
+import inspect
 from collections.abc import Callable
 from typing import Any
 
@@ -86,6 +88,41 @@ def _coerce_to_record_batch(result: Any) -> pa.RecordBatch:
         raise TypeError(f"Cannot convert handler result of type {type(result).__name__} to RecordBatch: {e}") from e
 
 
+def _get_or_create_event_loop() -> asyncio.AbstractEventLoop:
+    """Get the running event loop or create a new one."""
+    try:
+        return asyncio.get_running_loop()
+    except RuntimeError:
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        return loop
+
+
+def _run_async(coro):
+    """Run a coroutine synchronously."""
+    loop = _get_or_create_event_loop()
+    if loop.is_running():
+        # Already inside an async context — can't use run_until_complete.
+        # Create a new loop in a thread.
+        import concurrent.futures
+
+        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
+            return pool.submit(asyncio.run, coro).result()
+    return loop.run_until_complete(coro)
+
+
+def _collect_async_gen(agen) -> list:
+    """Collect all items from an async generator synchronously."""
+
+    async def _collect():
+        items = []
+        async for item in agen:
+            items.append(item)
+        return items
+
+    return _run_async(_collect())
+
+
 def process_batches(
     ipc_bytes: bytes,
     handler: Callable[..., Any],
@@ -109,12 +146,26 @@ def process_batches(
 
         try:
             result = handler(event, context_metadata)
+            # If the handler is async, await the coroutine.
+            if inspect.iscoroutine(result):
+                result = _run_async(result)
         except Exception as e:
             raise RuntimeError(f"Exception raised during handler execution: {e}") from e
 
         try:
-            result_batch = _coerce_to_record_batch(result)
-            results.append(encode_record_batch(result_batch))
+            if inspect.isgenerator(result):
+                # Sync generator: encode each yielded value as a separate chunk.
+                for item in result:
+                    result_batch = _coerce_to_record_batch(item)
+                    results.append(encode_record_batch(result_batch))
+            elif inspect.isasyncgen(result):
+                # Async generator: collect items and encode each as a separate chunk.
+                for item in _collect_async_gen(result):
+                    result_batch = _coerce_to_record_batch(item)
+                    results.append(encode_record_batch(result_batch))
+            else:
+                result_batch = _coerce_to_record_batch(result)
+                results.append(encode_record_batch(result_batch))
         except Exception as e:
             raise RuntimeError(f"Failed to encode response: {e}") from e
 

{chalk_remote_call_python-1.2.0 → chalk_remote_call_python-1.3.0}/chalk_remote_call_python.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chalk-remote-call-python
-Version: 1.2.0
+Version: 1.3.0
 Summary: Chalk remote call Python runtime interface client
 Author: Chalk AI, Inc.
 Project-URL: Homepage, https://chalk.ai

{chalk_remote_call_python-1.2.0 → chalk_remote_call_python-1.3.0}/tests/test_servicer.py

@@ -12,7 +12,7 @@ from chalk_remote_call._gen.chalk.runtime.v1.remote_python_call_pb2_grpc import
     RemoteCallServiceStub,
 )
 from chalk_remote_call.arrow_utils import decode_ipc_stream, encode_record_batch
-from chalk_remote_call.servicer import _coerce_to_record_batch
+from chalk_remote_call.servicer import _coerce_to_record_batch, process_batches
 
 
 class TestCoerceToRecordBatch:
@@ -222,3 +222,187 @@ class TestStructuredTypeCoercion:
         # Goes through the existing list path → struct array in a "result" column
         assert result.num_columns == 1
         assert result.schema.field(0).name == "result"
+
+
+# ---------------------------------------------------------------------------
+# Generator support tests
+# ---------------------------------------------------------------------------
+
+
+def _generator_handler(event: dict[str, pa.Array], context: Any) -> Any:
+    """Handler that yields each row value individually."""
+    n_col = list(event.values())[0]
+    n = n_col[0].as_py()
+    for i in range(n):
+        yield i
+
+
+def _single_handler(event: dict[str, pa.Array], context: Any) -> Any:
+    """Normal (non-generator) handler for comparison."""
+    a = event["a"]
+    b = event["b"]
+    return [a[i].as_py() + b[i].as_py() for i in range(len(a))]
+
+
+class TestProcessBatchesGenerator:
+    def test_generator_produces_multiple_chunks(self):
+        batch = pa.record_batch([pa.array([5])], names=["n"])
+        ipc_bytes = encode_record_batch(batch)
+
+        results = process_batches(ipc_bytes, _generator_handler, None, {})
+
+        assert len(results) == 5
+        for i, chunk_bytes in enumerate(results):
+            decoded = decode_ipc_stream(chunk_bytes)
+            assert len(decoded) == 1
+            assert decoded[0].column("result").to_pylist() == [i]
+
+    def test_non_generator_produces_single_chunk(self):
+        batch = pa.record_batch(
+            [pa.array([1, 2, 3], type=pa.int64()), pa.array([10, 20, 30], type=pa.int64())],
+            names=["a", "b"],
+        )
+        ipc_bytes = encode_record_batch(batch)
+
+        results = process_batches(ipc_bytes, _single_handler, ["a", "b"], {})
+
+        assert len(results) == 1
+        decoded = decode_ipc_stream(results[0])
+        assert decoded[0].column("result").to_pylist() == [11, 22, 33]
+
+    def test_generator_empty_yields_no_chunks(self):
+        batch = pa.record_batch([pa.array([0])], names=["n"])
+        ipc_bytes = encode_record_batch(batch)
+
+        results = process_batches(ipc_bytes, _generator_handler, None, {})
+
+        assert len(results) == 0
+
+
+@pytest.mark.parametrize(
+    "grpc_server_and_channel",
+    [{"handler": _generator_handler, "arg_names": None}],
+    indirect=True,
+)
+def test_generator_e2e_through_grpc(grpc_server_and_channel):
+    """Generator handler sends multiple response chunks over gRPC."""
+    _thread, channel, address = grpc_server_and_channel
+    stub = RemoteCallServiceStub(channel)
+
+    batch = pa.record_batch([pa.array([3])], names=["n"])
+    request_bytes = _make_request_bytes(batch)
+
+    def requests():
+        yield pb2.CallFunctionRequest(name="handler", feather_stream=request_bytes)
+
+    responses = list(stub.CallFunction(requests()))
+    # Generator yields 3 values → 3 response chunks
+    assert len(responses) == 3
+    for i, resp in enumerate(responses):
+        decoded = decode_ipc_stream(resp.feather_stream)
+        assert decoded[0].column("result").to_pylist() == [i]
+
+
+# ---------------------------------------------------------------------------
+# Async support tests
+# ---------------------------------------------------------------------------
+
+
+async def _async_handler(event: dict[str, pa.Array], context: Any) -> Any:
+    """Async coroutine handler that adds two columns."""
+    a = event["a"]
+    b = event["b"]
+    return [a[i].as_py() + b[i].as_py() for i in range(len(a))]
+
+
+async def _async_generator_handler(event: dict[str, pa.Array], context: Any) -> Any:
+    """Async generator handler that yields each integer 0..n-1."""
+    n_col = list(event.values())[0]
+    n = n_col[0].as_py()
+    for i in range(n):
+        yield i
+
+
+class TestProcessBatchesAsync:
+    def test_async_coroutine_handler(self):
+        """async def handler is awaited and result is encoded."""
+        batch = pa.record_batch(
+            [pa.array([1, 2], type=pa.int64()), pa.array([10, 20], type=pa.int64())],
+            names=["a", "b"],
+        )
+        ipc_bytes = encode_record_batch(batch)
+
+        results = process_batches(ipc_bytes, _async_handler, ["a", "b"], {})
+
+        assert len(results) == 1
+        decoded = decode_ipc_stream(results[0])
+        assert decoded[0].column("result").to_pylist() == [11, 22]
+
+    def test_async_generator_produces_multiple_chunks(self):
+        """async def with yield produces separate chunks."""
+        batch = pa.record_batch([pa.array([4])], names=["n"])
+        ipc_bytes = encode_record_batch(batch)
+
+        results = process_batches(ipc_bytes, _async_generator_handler, None, {})
+
+        assert len(results) == 4
+        for i, chunk_bytes in enumerate(results):
+            decoded = decode_ipc_stream(chunk_bytes)
+            assert decoded[0].column("result").to_pylist() == [i]
+
+    def test_async_generator_empty_yields_no_chunks(self):
+        """async generator that yields nothing produces 0 chunks."""
+        batch = pa.record_batch([pa.array([0])], names=["n"])
+        ipc_bytes = encode_record_batch(batch)
+
+        results = process_batches(ipc_bytes, _async_generator_handler, None, {})
+
+        assert len(results) == 0
+
+
+@pytest.mark.parametrize(
+    "grpc_server_and_channel",
+    [{"handler": _async_generator_handler, "arg_names": None}],
+    indirect=True,
+)
+def test_async_generator_e2e_through_grpc(grpc_server_and_channel):
+    """Async generator handler sends multiple response chunks over gRPC."""
+    _thread, channel, address = grpc_server_and_channel
+    stub = RemoteCallServiceStub(channel)
+
+    batch = pa.record_batch([pa.array([3])], names=["n"])
+    request_bytes = _make_request_bytes(batch)
+
+    def requests():
+        yield pb2.CallFunctionRequest(name="handler", feather_stream=request_bytes)
+
+    responses = list(stub.CallFunction(requests()))
+    assert len(responses) == 3
+    for i, resp in enumerate(responses):
+        decoded = decode_ipc_stream(resp.feather_stream)
+        assert decoded[0].column("result").to_pylist() == [i]
+
+
+@pytest.mark.parametrize(
+    "grpc_server_and_channel",
+    [{"handler": _async_handler, "arg_names": ["a", "b"]}],
+    indirect=True,
+)
+def test_async_coroutine_e2e_through_grpc(grpc_server_and_channel):
+    """Async coroutine handler works end-to-end through gRPC."""
+    _thread, channel, address = grpc_server_and_channel
+    stub = RemoteCallServiceStub(channel)
+
+    batch = pa.record_batch(
+        [pa.array([5, 10], type=pa.int64()), pa.array([3, 7], type=pa.int64())],
+        names=["a", "b"],
+    )
+    request_bytes = _make_request_bytes(batch)
+
+    def requests():
+        yield pb2.CallFunctionRequest(name="handler", feather_stream=request_bytes)
+
+    responses = list(stub.CallFunction(requests()))
+    assert len(responses) == 1
+    decoded = decode_ipc_stream(responses[0].feather_stream)
+    assert decoded[0].column("result").to_pylist() == [8, 17]

chalk_remote_call_python-1.2.0/chalk_remote_call/_version.py

@@ -1 +0,0 @@
-__version__ = "1.2.0"