Add workflow streams, closes #461

src/python/txtai/app/base.py

@@ -137,6 +137,10 @@ class Application:
                 # Resolve callable functions
                 config["tasks"] = [self.resolve(task) for task in config["tasks"]]
 
+                # Resolve stream functions
+                if "stream" in config:
+                    config["stream"] = self.resolve(config["stream"])
+
                 # Get scheduler config
                 schedule = config.pop("schedule", None)
 
@@ -640,8 +644,12 @@ class Application:
             processed elements
         """
 
-        # Convert lists to tuples
-        elements = [tuple(element) if isinstance(element, list) else element for element in elements]
+        if hasattr(elements, "__len__") and hasattr(elements, "__getitem__"):
+            # Convert to tuples and return as a list since input is sized
+            elements = [tuple(element) if isinstance(element, list) else element for element in elements]
+        else:
+            # Convert to tuples and return as a generator since input is not sized
+            elements = (tuple(element) if isinstance(element, list) else element for element in elements)
 
         # Execute workflow
         return self.workflows[name](elements)

src/python/txtai/workflow/base.py

@@ -27,7 +27,7 @@ class Workflow:
     Base class for all workflows.
     """
 
-    def __init__(self, tasks, batch=100, workers=None, name=None):
+    def __init__(self, tasks, batch=100, workers=None, name=None, stream=None):
         """
         Creates a new workflow. Workflows are lists of tasks to execute.
 
@@ -36,12 +36,14 @@ class Workflow:
             batch: how many items to process at a time, defaults to 100
             workers: number of concurrent workers
             name: workflow name
+            stream: workflow stream processor
         """
 
         self.tasks = tasks
         self.batch = batch
         self.workers = workers
         self.name = name
+        self.stream = stream
 
         # Set default number of executor workers to max number of actions in a task
         self.workers = max(len(task.action) for task in self.tasks) if not self.workers else self.workers
@@ -63,6 +65,9 @@ class Workflow:
             # Run task initializers
             self.initialize()
 
+            # Process elements with stream processor, if available
+            elements = self.stream(elements) if self.stream else elements
+
             # Process elements in batches
             for batch in self.chunk(elements):
                 yield from self.process(batch, executor)

src/python/txtai/workflow/factory.py

@@ -32,5 +32,11 @@ class WorkflowFactory:
 
         config["tasks"] = tasks
 
+        if "stream" in config:
+            sconfig = config["stream"]
+            task = sconfig.pop("task") if "task" in sconfig else "stream"
+
+            config["stream"] = TaskFactory.create(sconfig, task)
+
         # Create workflow
         return Workflow(**config, name=name)

src/python/txtai/workflow/task/__init__.py

@@ -11,6 +11,7 @@ from .image import ImageTask
 from .retrieve import RetrieveTask
 from .service import ServiceTask
 from .storage import StorageTask
+from .stream import StreamTask
 from .template import ExtractorTask, TemplateTask
 from .url import UrlTask
 from .workflow import WorkflowTask

src/python/txtai/workflow/task/stream.py

@@ -0,0 +1,33 @@
+"""
+StreamTask module
+"""
+
+from .base import Task
+
+
+class StreamTask(Task):
+    """
+    Task that calls a task action and yields results.
+    """
+
+    def register(self, batch=False):
+        """
+        Adds stream parameters to task.
+
+        Args:
+            batch: all elements are passed to a single action call if True, otherwise an action call is executed per element, defaults to False
+        """
+
+        # pylint: disable=W0201
+        # All elements are passed to a single action call if True, otherwise an action call is executed per element, defaults to False
+        self.batch = batch
+
+    def __call__(self, elements, executor=None):
+        for action in self.action:
+            if self.batch:
+                # Single batch call
+                yield from action(elements)
+            else:
+                # Call action for each element
+                for x in elements:
+                    yield from action(x)

test/python/testapp.py

@@ -3,6 +3,7 @@ Application module tests
 """
 
 import unittest
+import types
 
 from txtai.app import Application
 from txtai.pipeline import Pipeline
@@ -36,6 +37,35 @@ class TestApp(unittest.TestCase):
         # Check that application instance is not None
         self.assertIsNotNone(app.pipelines["testapp.TestPipeline"].application)
 
+    def testStream(self):
+        """
+        Test workflow streams
+        """
+
+        app = Application(
+            """
+            workflow:
+                stream:
+                    stream:
+                        action: testapp.TestStream
+                    tasks:
+                        - nop
+                batchstream:
+                    stream:
+                        action: testapp.TestStream
+                        batch: True
+                    tasks:
+                        - nop
+        """
+        )
+
+        def generator():
+            yield 10
+
+        self.assertEqual(list(app.workflow("stream", [10])), list(range(10)))
+
+        self.assertEqual(list(app.workflow("batchstream", generator())), list(range(10)))
+
 
 class TestPipeline(Pipeline):
     """
@@ -44,3 +74,16 @@ class TestPipeline(Pipeline):
 
     def __init__(self, application):
         self.application = application
+
+
+class TestStream:
+    """
+    Test workflow stream
+    """
+
+    def __call__(self, arg):
+        if isinstance(arg, types.GeneratorType):
+            for x in arg:
+                yield from range(int(x))
+        else:
+            yield from range(int(arg))