Merge branch 'dev'

Author: Maplemx

agently/core/ModelRequest.py

@@ -432,6 +432,9 @@ def __init__(
         self.get_data = FunctionShifter.syncify(self.async_get_data)
         self.get_data_object = FunctionShifter.syncify(self.async_get_data_object)
 
+        self.start = self.get_data
+        self.async_start = self.async_get_data
+
     def set_prompt(
         self,
         key: "PromptStandardSlot | str",

agently/core/TriggerFlow/BluePrint.py

@@ -132,6 +132,7 @@ def create_execution(
         *,
         execution_id: str | None = None,
         skip_exceptions: bool = False,
+        concurrency: int | None = None,
     ):
         handlers_snapshot: TriggerFlowAllHandlers = {
             "event": {k: v.copy() for k, v in self._handlers["event"].items()},
@@ -143,6 +144,7 @@ def create_execution(
             trigger_flow=trigger_flow,
             id=execution_id,
             skip_exceptions=skip_exceptions,
+            concurrency=concurrency,
         )
 
     def copy(self, *, name: str | None = None):

agently/core/TriggerFlow/Execution.py

@@ -16,6 +16,7 @@
 import uuid
 import asyncio
 import warnings
+from contextvars import ContextVar
 
 from typing import Any, Literal, TYPE_CHECKING
 
@@ -37,6 +38,7 @@ def __init__(
         trigger_flow: "TriggerFlow",
         id: str | None = None,
         skip_exceptions: bool = False,
+        concurrency: int | None = None,
     ):
         # Basic Attributions
         self.id = id if id is not None else uuid.uuid4().hex
@@ -45,6 +47,11 @@ def __init__(
         self._runtime_data = RuntimeData()
         self._system_runtime_data = RuntimeData()
         self._skip_exceptions = skip_exceptions
+        self._concurrency_semaphore = asyncio.Semaphore(concurrency) if concurrency and concurrency > 0 else None
+        self._concurrency_depth = ContextVar(
+            f"trigger_flow_execution_concurrency_depth_{ self.id }",
+            default=0,
+        )
 
         # Settings
         self.settings = Settings(
@@ -126,19 +133,29 @@ async def async_emit(
                     },
                     self.settings,
                 )
-                tasks.append(
-                    asyncio.ensure_future(
-                        FunctionShifter.asyncify(handler)(
-                            TriggerFlowEventData(
-                                trigger_event=trigger_event,
-                                trigger_type=trigger_type,
-                                value=value,
-                                execution=self,
-                                _layer_marks=_layer_marks,
-                            )
-                        )
+                async def run_handler(handler_func):
+                    if self._concurrency_semaphore is None:
+                        return await handler_func
+                    depth = self._concurrency_depth.get()
+                    token = self._concurrency_depth.set(depth + 1)
+                    try:
+                        if depth > 0:
+                            return await handler_func
+                        async with self._concurrency_semaphore:
+                            return await handler_func
+                    finally:
+                        self._concurrency_depth.reset(token)
+
+                handler_task = FunctionShifter.asyncify(handler)(
+                    TriggerFlowEventData(
+                        trigger_event=trigger_event,
+                        trigger_type=trigger_type,
+                        value=value,
+                        execution=self,
+                        _layer_marks=_layer_marks,
                     )
                 )
+                tasks.append(asyncio.ensure_future(run_handler(handler_task)))
 
         if tasks:
             await asyncio.gather(*tasks, return_exceptions=self._skip_exceptions)

agently/core/TriggerFlow/TriggerFlow.py

@@ -99,13 +99,19 @@ def wrapper(func: "TriggerFlowHandler"):
             self._blue_print.chunks[handler_or_name.__name__] = chunk
             return chunk
 
-    def create_execution(self, *, skip_exceptions: bool | None = None):
+    def create_execution(
+        self,
+        *,
+        skip_exceptions: bool | None = None,
+        concurrency: int | None = None,
+    ):
         execution_id = uuid.uuid4().hex
         skip_exceptions = skip_exceptions if skip_exceptions is not None else self._skip_exceptions
         execution = self._blue_print.create_execution(
             self,
             execution_id=execution_id,
             skip_exceptions=skip_exceptions,
+            concurrency=concurrency,
         )
         self._executions[execution_id] = execution
         return execution
@@ -118,8 +124,14 @@ def remove_execution(self, execution: "TriggerFlowExecution | str"):
             if execution.id in self._executions:
                 del self._executions[execution.id]
 
-    async def async_start_execution(self, initial_value: Any, *, wait_for_result: bool = False):
-        execution = self.create_execution()
+    async def async_start_execution(
+        self,
+        initial_value: Any,
+        *,
+        wait_for_result: bool = False,
+        concurrency: int | None = None,
+    ):
+        execution = self.create_execution(concurrency=concurrency)
         await execution.async_start(initial_value, wait_for_result=wait_for_result)
         return execution
 
@@ -192,8 +204,9 @@ async def async_start(
         *,
         wait_for_result: bool = True,
         timeout: int | None = 10,
+        concurrency: int | None = None,
     ):
-        execution = await self.async_start_execution(initial_value)
+        execution = await self.async_start_execution(initial_value, concurrency=concurrency)
         if wait_for_result:
             return await execution.async_get_result(timeout=timeout)
 
@@ -202,8 +215,9 @@ def get_async_runtime_stream(
         initial_value: Any = None,
         *,
         timeout: int | None = 10,
+        concurrency: int | None = None,
     ):
-        execution = self.create_execution()
+        execution = self.create_execution(concurrency=concurrency)
         return execution.get_async_runtime_stream(
             initial_value,
             timeout=timeout,
@@ -214,8 +228,9 @@ def get_runtime_stream(
         initial_value: Any = None,
         *,
         timeout: int | None = 10,
+        concurrency: int | None = None,
     ):
-        execution = self.create_execution()
+        execution = self.create_execution(concurrency=concurrency)
         return execution.get_runtime_stream(
             initial_value,
             timeout=timeout,

agently/core/TriggerFlow/process/BaseProcess.py

@@ -14,7 +14,7 @@
 
 
 import uuid
-from asyncio import Event
+from asyncio import Event, Semaphore
 from threading import Lock
 
 from typing import Callable, Any, Literal, TYPE_CHECKING, overload, cast
@@ -266,11 +266,13 @@ def batch(
         self,
         *chunks: "TriggerFlowChunk | TriggerFlowHandler | tuple[str, TriggerFlowHandler]",
         side_branch: bool = False,
+        concurrency: int | None = None,
     ):
         batch_trigger = f"Batch-{ uuid.uuid4().hex }"
         results = {}
         triggers_to_wait = {}
         trigger_to_chunk_name = {}
+        semaphore = Semaphore(concurrency) if concurrency and concurrency > 0 else None
 
         async def wait_all_chunks(data: "TriggerFlowEventData"):
             if data.event in triggers_to_wait:
@@ -295,10 +297,18 @@ async def wait_all_chunks(data: "TriggerFlowEventData"):
             triggers_to_wait[chunk.trigger] = False
             trigger_to_chunk_name[chunk.trigger] = chunk.name
             results[chunk.name] = None
+
+            if semaphore is None:
+                handler = chunk.async_call
+            else:
+                async def handler(data: "TriggerFlowEventData", _chunk=chunk):
+                    async with semaphore:
+                        return await _chunk.async_call(data)
+
             self._blue_print.add_handler(
                 self.trigger_type,
                 self.trigger_event,
-                chunk.async_call,
+                handler,
             )
             self._blue_print.add_event_handler(chunk.trigger, wait_all_chunks)
 

agently/core/TriggerFlow/process/ForEachProcess.py

@@ -22,7 +22,7 @@
 
 
 class TriggerFlowForEachProcess(TriggerFlowBaseProcess):
-    def for_each(self):
+    def for_each(self, *, concurrency: int | None = None):
         for_each_id = uuid.uuid4().hex
         for_each_block_data = TriggerFlowBlockData(
             outer_block=self._block_data,
@@ -31,40 +31,46 @@ def for_each(self):
             },
         )
         send_item_trigger = f"ForEach-{ for_each_id }-Send"
+        semaphore = asyncio.Semaphore(concurrency) if concurrency and concurrency > 0 else None
 
         async def send_items(data: "TriggerFlowEventData"):
             data.layer_in()
             for_each_instance_id = data.layer_mark
             assert for_each_instance_id is not None
 
             send_tasks = []
-            if not isinstance(data.value, str) and isinstance(data.value, Sequence):
-                items = list(data.value)
-                for item in items:
-                    data.layer_in()
-                    item_id = data.layer_mark
-                    assert item_id is not None
-                    data._system_runtime_data.set(f"for_each_results.{ for_each_instance_id }.{ item_id }", EMPTY)
-                    send_tasks.append(
-                        data.async_emit(
-                            send_item_trigger,
-                            item,
-                            data._layer_marks.copy(),
-                        )
-                    )
-                    data.layer_out()
-                await asyncio.gather(*send_tasks)
-            else:
+            def prepare_item(item):
                 data.layer_in()
                 item_id = data.layer_mark
                 assert item_id is not None
+                layer_marks = data._layer_marks.copy()
                 data._system_runtime_data.set(f"for_each_results.{ for_each_instance_id }.{ item_id }", EMPTY)
-                await data.async_emit(
-                    send_item_trigger,
-                    data.value,
-                    data._layer_marks.copy(),
-                )
                 data.layer_out()
+                return item_id, layer_marks, item
+
+            async def emit_item(item, layer_marks):
+                if semaphore is None:
+                    await data.async_emit(
+                        send_item_trigger,
+                        item,
+                        layer_marks,
+                    )
+                else:
+                    async with semaphore:
+                        await data.async_emit(
+                            send_item_trigger,
+                            item,
+                            layer_marks,
+                        )
+            if not isinstance(data.value, str) and isinstance(data.value, Sequence):
+                items = list(data.value)
+                for item in items:
+                    _, layer_marks, item_value = prepare_item(item)
+                    send_tasks.append(emit_item(item_value, layer_marks))
+                await asyncio.gather(*send_tasks)
+            else:
+                _, layer_marks, item_value = prepare_item(data.value)
+                await emit_item(item_value, layer_marks)
 
         self.to(send_items)
 

examples/trigger_flow/batch.py

@@ -17,6 +17,7 @@ async def echo(data: TriggerFlowEventData):
     ("echo_2", echo),
     ("echo_3", echo),
     ("echo_4", echo),
+    concurrency=2,
 ).end()
 execution = flow.create_execution()
 result = execution.start("Agently")

examples/trigger_flow/for_each.py

@@ -12,7 +12,7 @@ async def handle(data: TriggerFlowEventData):
 
 flow_1 = TriggerFlow()
 
-flow_1.for_each().to(handle).end_for_each().to(lambda data: data.value).end()
+flow_1.for_each(concurrency=2).to(handle).end_for_each().to(lambda data: data.value).end()
 
 execution_1 = flow_1.create_execution()
 result = execution_1.start(

examples/trigger_flow/global_concurrency.py

@@ -0,0 +1,21 @@
+import asyncio
+from agently import TriggerFlow, TriggerFlowEventData
+
+
+async def handle(data: TriggerFlowEventData):
+    print(f"Hi, { data.value }")
+    await asyncio.sleep(2)
+    return f"handled: { data.value }"
+
+
+flow = TriggerFlow()
+flow.batch(
+    ("a", handle),
+    ("b", handle),
+    ("c", handle),
+    ("d", handle),
+).end()
+
+execution = flow.create_execution(concurrency=2)
+result = execution.start("hello")
+print(result)

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "agently"
-version = "4.0.6.11"
+version = "4.0.7"
 description = ""
 authors = [
     {name = "Agently Team",email = "developer@agently.tech"},