fix: make RunnableRails.atransform return AsyncIterator

nemoguardrails/integrations/langchain/runnable_rails.py

@@ -852,24 +852,30 @@ async def abatch(
 
     def transform(
         self,
-        input: Input,
+        input_stream: Iterator[Input],
         config: Optional[RunnableConfig] = None,
         **kwargs: Optional[Any],
-    ) -> Output:
+    ) -> Iterator[Output]:
         """Transform the input.
 
-        This is just an alias for invoke.
+        Consumes the input iterator and yields the result of invoke for each
+        item.  For RunnableRails the entire input must be collected before
+        processing, so streaming granularity matches invoke.
         """
-        return self.invoke(input, config, **kwargs)
+        for item in input_stream:
+            yield self.invoke(item, config, **kwargs)
 
     async def atransform(
         self,
-        input: Input,
+        input_stream: AsyncIterator[Input],
         config: Optional[RunnableConfig] = None,
         **kwargs: Optional[Any],
-    ) -> Output:
+    ) -> AsyncIterator[Output]:
         """Transform the input asynchronously.
 
-        This is just an alias for ainvoke.
+        Consumes the async input iterator and yields the result of ainvoke for
+        each item.  For RunnableRails the entire input must be collected before
+        processing, so streaming granularity matches ainvoke.
         """
-        return await self.ainvoke(input, config, **kwargs)
+        async for item in input_stream:
+            yield await self.ainvoke(item, config, **kwargs)

tests/runnable_rails/test_atransform_async_iter.py

@@ -0,0 +1,67 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests that RunnableRails.atransform returns an AsyncIterator (fixes #1692)."""
+
+from collections.abc import AsyncIterator, Iterator
+
+import pytest
+from langchain_core.runnables import RunnableLambda
+
+from nemoguardrails import RailsConfig
+from nemoguardrails.integrations.langchain.runnable_rails import RunnableRails
+
+
+@pytest.fixture
+def passthrough_rails():
+    """Create a minimal passthrough RunnableRails instance."""
+    config = RailsConfig.from_content(yaml_content="models: []")
+    return RunnableRails(config, passthrough=True)
+
+
+def test_transform_returns_iterator(passthrough_rails):
+    """transform() should yield results, not return a single value."""
+    result = passthrough_rails.transform(iter(["hello"]))
+    assert isinstance(result, Iterator)
+    items = list(result)
+    assert len(items) == 1
+
+
+@pytest.mark.asyncio
+async def test_atransform_returns_async_iterator(passthrough_rails):
+    """atransform() should yield results as an async iterator."""
+
+    async def _aiter():
+        yield "hello"
+
+    result = passthrough_rails.atransform(_aiter())
+    assert isinstance(result, AsyncIterator)
+    items = []
+    async for item in result:
+        items.append(item)
+    assert len(items) == 1
+
+
+@pytest.mark.asyncio
+async def test_astream_through_pipeline(passthrough_rails):
+    """RunnableRails should work in a pipeline with astream (reproduces #1692)."""
+    pipeline = RunnableLambda(lambda x: x) | passthrough_rails
+
+    items = []
+    async for chunk in pipeline.astream("hello"):
+        items.append(chunk)
+
+    # Should get at least one result without TypeError
+    assert len(items) >= 1