mageflow 0.3.2__tar.gz → 0.3.4__tar.gz

{mageflow-0.3.2 → mageflow-0.3.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mageflow
-Version: 0.3.2
+Version: 0.3.4
 Summary: Manage Graph Execution Flow - A unified interface for task orchestration across different task managers
 Project-URL: Homepage, https://imaginary-cherry.github.io/mageflow/
 Project-URL: Documentation, https://imaginary-cherry.github.io/mageflow/
@@ -26,7 +26,7 @@ Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Classifier: Topic :: System :: Distributed Computing
 Classifier: Typing :: Typed
 Requires-Python: <3.14,>=3.10
-Requires-Dist: thirdmagic<0.1.0,>=0.0.3
+Requires-Dist: thirdmagic<0.1.0,>=0.0.5
 Provides-Extra: dev
 Requires-Dist: black>=26.1.0; extra == 'dev'
 Requires-Dist: coverage[toml]<8.0.0,>=7.0.0; extra == 'dev'
@@ -37,5 +37,11 @@ Requires-Dist: pytest-asyncio<2.0.0,>=1.2.0; extra == 'dev'
 Requires-Dist: pytest<10.0.0,>=9.0.2; extra == 'dev'
 Requires-Dist: requests<3.0.0,>=2.32.5; extra == 'dev'
 Requires-Dist: ruff>=0.15.5; extra == 'dev'
+Requires-Dist: testcontainers[redis]<5.0.0,>=4.14.0; extra == 'dev'
 Provides-Extra: hatchet
 Requires-Dist: hatchet-sdk<1.24.0,>=1.22.5; extra == 'hatchet'
+Provides-Extra: testing
+Requires-Dist: fakeredis[json,lua]<3.0.0,>=2.34.0; extra == 'testing'
+Requires-Dist: pytest-asyncio<2.0.0,>=1.2.0; extra == 'testing'
+Requires-Dist: pytest<10.0.0,>=9.0.0; extra == 'testing'
+Requires-Dist: testcontainers[redis]<5.0.0,>=4.14.0; extra == 'testing'

{mageflow-0.3.2 → mageflow-0.3.4}/mageflow/__init__.py

@@ -1,16 +1,17 @@
 import rapyer
 from rapyer.fields import RapyerKey
-
-from mageflow.callbacks import handle_task_callback
-from mageflow.client import Mageflow
-from mageflow.config import MageflowConfig, SignatureTTLConfig, TTLConfig
-from mageflow.startup import start_mageflow
+from thirdmagic import abounded_field
 from thirdmagic.chain.creator import chain as achain
 from thirdmagic.signature import Signature
 from thirdmagic.swarm.creator import swarm as aswarm
 from thirdmagic.task import TaskSignature
 from thirdmagic.task import sign as asign
 
+from mageflow.callbacks import handle_task_callback
+from mageflow.client import Mageflow
+from mageflow.config import MageflowConfig, SignatureTTLConfig, TTLConfig
+from mageflow.startup import start_mageflow
+
 lock_task = TaskSignature.alock_from_key
 resume = TaskSignature.resume_from_key
 pause = TaskSignature.pause_from_key
@@ -22,7 +23,6 @@ async def load_sign(key: RapyerKey) -> Signature:
 
 
 load_signature = rapyer.afind_one
-abounded_field = rapyer.apipeline
 
 
 __all__ = [

{mageflow-0.3.2 → mageflow-0.3.4}/mageflow/callbacks.py

@@ -7,11 +7,16 @@ from typing import Any
 from hatchet_sdk import Context
 from hatchet_sdk.runnables.types import EmptyModel
 from pydantic import BaseModel
-
-from mageflow.utils.pythonic import flexible_call
+from thirdmagic.signature.retry_cache import (
+    retry_cache_ctx,
+    setup_retry_cache,
+    teardown_retry_cache,
+)
 from thirdmagic.task.model import TaskSignature
 from thirdmagic.task_def import MageflowTaskDefinition
 
+from mageflow.utils.pythonic import flexible_call
+
 
 class AcceptParams(Enum):
     JUST_MESSAGE = 1
@@ -27,6 +32,7 @@ def handle_task_callback(
     expected_params: AcceptParams = AcceptParams.NO_CTX,
     wrap_res: bool = True,
     send_signature: bool = False,
+    is_idempotent: bool = False,
 ):
     def task_decorator(func):
         @functools.wraps(func)
@@ -40,8 +46,18 @@ def handle_task_callback(
                 # NOTE: This should not run, the task should cancel, but just in case
                 return {"Error": "Task should have been canceled"}
             is_normal_run = lifecycle.is_vanilla_run()
+            is_task_finish = False
             signature = await lifecycle.start_task()
 
+            # Setup retry cache for signature idempotency on retries (durable tasks only)
+            cache_token = None
+            cache_state = None
+            if is_idempotent:
+                cache_state = await setup_retry_cache(
+                    ctx.workflow_id, ctx.attempt_number
+                )
+                cache_token = retry_cache_ctx.set(cache_state)
+
             # Add params if user requires
             if send_signature:
                 kwargs["signature"] = signature
@@ -53,20 +69,18 @@ def handle_task_callback(
                     result = await flexible_call(func, message, *args, **kwargs)
                 else:
                     result = await flexible_call(func, message, ctx, *args, **kwargs)
-            except asyncio.CancelledError as e:
-                if not is_normal_run:
-                    await lifecycle.task_failed(msg_data, e)
-                raise
-            except Exception as e:
-                if is_normal_run:
-                    raise
-                if not TaskSignature.ClientAdapter.should_task_retry(
+            except (Exception, asyncio.CancelledError) as e:
+                will_retry = TaskSignature.ClientAdapter.should_task_retry(
                     task_model, ctx.attempt_number, e
-                ):
-                    await lifecycle.task_failed(msg_data, e)
+                )
+                if not will_retry:
+                    is_task_finish = True
+                    if not is_normal_run:
+                        await lifecycle.task_failed(msg_data, e)
                 raise
             else:
                 # If this is a simple task, no signature, then we dont do any manipulation
+                is_task_finish = True
                 if is_normal_run:
                     return result
                 task_results = HatchetResult(hatchet_results=result)
@@ -76,6 +90,11 @@ def handle_task_callback(
                     return task_results
                 else:
                     return result
+            finally:
+                if cache_token is not None:
+                    retry_cache_ctx.reset(cache_token)
+                if is_task_finish and cache_state:
+                    await teardown_retry_cache(cache_state)
 
         wrapper.__signature__ = inspect.signature(func)
         return wrapper

{mageflow-0.3.2 → mageflow-0.3.4}/mageflow/chain/workflows.py

@@ -2,7 +2,6 @@ from logging import Logger
 from typing import Any
 
 from rapyer.fields import RapyerKey
-
 from thirdmagic.clients.lifecycle import BaseLifecycle
 
 

{mageflow-0.3.2 → mageflow-0.3.4}/mageflow/clients/hatchet/adapter.py

@@ -8,6 +8,13 @@ from hatchet_sdk.runnables.types import EmptyModel
 from hatchet_sdk.runnables.workflow import BaseWorkflow
 from pydantic import BaseModel, TypeAdapter
 from rapyer.fields import RapyerKey
+from thirdmagic.chain import ChainTaskSignature
+from thirdmagic.clients.base import BaseClientAdapter
+from thirdmagic.consts import TASK_ID_PARAM_NAME
+from thirdmagic.signature import Signature
+from thirdmagic.swarm import SwarmTaskSignature
+from thirdmagic.task import TaskSignature
+from thirdmagic.task_def import MageflowTaskDefinition
 
 from mageflow.chain.messages import ChainCallbackMessage, ChainErrorMessage
 from mageflow.clients.hatchet.workflow import MageflowWorkflow
@@ -25,13 +32,6 @@ from mageflow.swarm.messages import (
     SwarmErrorMessage,
     SwarmResultsMessage,
 )
-from thirdmagic.chain import ChainTaskSignature
-from thirdmagic.clients.base import BaseClientAdapter
-from thirdmagic.consts import TASK_ID_PARAM_NAME
-from thirdmagic.signature import Signature
-from thirdmagic.swarm import SwarmTaskSignature
-from thirdmagic.task import TaskSignature
-from thirdmagic.task_def import MageflowTaskDefinition
 
 
 class HatchetClientAdapter(BaseClientAdapter):
@@ -123,6 +123,15 @@ class HatchetClientAdapter(BaseClientAdapter):
     def extract_retries(self, client_task: BaseWorkflow) -> int:
         return client_task.tasks[0].retries
 
+    def _prepare_wf(self, signature: TaskSignature, set_return_field: bool, **kwargs):
+        total_kwargs = signature.kwargs | kwargs
+        workflow = self.hatchet.workflow(
+            name=signature.task_name, input_validator=signature.model_validators
+        )
+        return_field_name = signature.return_field_name if set_return_field else None
+        mageflow_wf = MageflowWorkflow(workflow, total_kwargs, return_field_name)
+        return mageflow_wf
+
     async def acall_signature(
         self,
         signature: TaskSignature,
@@ -134,17 +143,23 @@ class HatchetClientAdapter(BaseClientAdapter):
         if msg is None:
             msg = EmptyModel()
         options = self._update_options(signature, options)
-        total_kwargs = signature.kwargs | kwargs
-        workflow = self.hatchet.workflow(
-            name=signature.task_name, input_validator=signature.model_validators
-        )
-        mageflow_wf = MageflowWorkflow(
-            workflow,
-            total_kwargs,
-            signature.return_field_name if set_return_field else None,
-        )
+        mageflow_wf = self._prepare_wf(signature, set_return_field, **kwargs)
         return await mageflow_wf.aio_run_no_wait(msg, options)
 
+    async def await_signature(
+        self,
+        signature: "TaskSignature",
+        msg: Any,
+        set_return_field: bool,
+        options: TriggerWorkflowOptions = None,
+        **kwargs,
+    ):
+        if msg is None:
+            msg = EmptyModel()
+        options = self._update_options(signature, options)
+        mageflow_wf = self._prepare_wf(signature, set_return_field, **kwargs)
+        return await mageflow_wf.aio_run(msg, options)
+
     def should_task_retry(
         self,
         task_definition: MageflowTaskDefinition,
@@ -153,7 +168,7 @@ class HatchetClientAdapter(BaseClientAdapter):
     ) -> bool:
         finish_retry = (
             task_definition.retries is not None
-            and attempt_num < task_definition.retries
+            and attempt_num <= task_definition.retries
         )
         return finish_retry and not isinstance(e, NonRetryableException)
 

{mageflow-0.3.2 → mageflow-0.3.4}/mageflow/clients/hatchet/mageflow.py

@@ -17,6 +17,14 @@ from hatchet_sdk.runnables.types import (
 from hatchet_sdk.runnables.workflow import BaseWorkflow, Standalone
 from hatchet_sdk.worker.worker import LifespanFn
 from redis.asyncio import Redis
+from thirdmagic import chain, sign
+from thirdmagic.chain import ChainTaskSignature
+from thirdmagic.signature import Signature
+from thirdmagic.swarm import SwarmTaskSignature
+from thirdmagic.swarm.creator import SignatureOptions, swarm
+from thirdmagic.task import TaskInputType, TaskSignature, TaskSignatureConvertible
+from thirdmagic.task_def import MageflowTaskDefinition
+from thirdmagic.utils import HatchetTaskType
 from typing_extensions import override
 
 from mageflow.callbacks import AcceptParams, handle_task_callback
@@ -47,14 +55,6 @@ from mageflow.swarm.workflows import (
     swarm_item_failed,
 )
 from mageflow.utils.mageflow import does_task_wants_ctx
-from thirdmagic import chain, sign
-from thirdmagic.chain import ChainTaskSignature
-from thirdmagic.signature import Signature
-from thirdmagic.swarm import SwarmTaskSignature
-from thirdmagic.swarm.creator import SignatureOptions, swarm
-from thirdmagic.task import TaskInputType, TaskSignature, TaskSignatureConvertible
-from thirdmagic.task_def import MageflowTaskDefinition
-from thirdmagic.utils import HatchetTaskType
 
 Duration = timedelta | str
 
@@ -125,14 +125,16 @@ class HatchetMageflow(Hatchet):
             )
         )
 
-    def task_decorator(self, func: Callable, hatchet_task):
+    def task_decorator(self, func: Callable, hatchet_task, is_idempotent: bool = False):
         param_config = (
             AcceptParams.ALL
             if does_task_wants_ctx(func)
             else self.mageflow_config.param_config
         )
         send_signature = getattr(func, "__send_signature__", False)
-        handler_dec = handle_task_callback(param_config, send_signature=send_signature)
+        handler_dec = handle_task_callback(
+            param_config, send_signature=send_signature, is_idempotent=is_idempotent
+        )
         func = handler_dec(func)
         wf = hatchet_task(func)
         self._add_task_def(wf)
@@ -146,7 +148,9 @@ class HatchetMageflow(Hatchet):
         """
         hatchet_task = super().task(name=name, **kwargs)
 
-        decorator = functools.partial(self.task_decorator, hatchet_task=hatchet_task)
+        decorator = functools.partial(
+            self.task_decorator, hatchet_task=hatchet_task, is_idempotent=False
+        )
         return decorator
 
     @override
@@ -156,7 +160,9 @@ class HatchetMageflow(Hatchet):
         """
         hatchet_task = super().durable_task(name=name, **kwargs)
 
-        decorator = functools.partial(self.task_decorator, hatchet_task=hatchet_task)
+        decorator = functools.partial(
+            self.task_decorator, hatchet_task=hatchet_task, is_idempotent=True
+        )
 
         return decorator
 

{mageflow-0.3.2 → mageflow-0.3.4}/mageflow/clients/hatchet/workflow.py

@@ -3,7 +3,6 @@ from typing import Any
 from hatchet_sdk.runnables.workflow import Workflow
 from hatchet_sdk.utils.typing import JSONSerializableMapping
 from pydantic import BaseModel
-
 from thirdmagic.utils import deep_merge
 
 

{mageflow-0.3.2 → mageflow-0.3.4}/mageflow/config.py

@@ -1,26 +1,29 @@
 import dataclasses
-from dataclasses import dataclass, field
+from dataclasses import field
 from typing import Optional
 
-from mageflow.callbacks import AcceptParams
+from pydantic import Field
+from pydantic.dataclasses import dataclass
 from thirdmagic.chain import ChainTaskSignature
 from thirdmagic.consts import REMOVED_TASK_TTL
-from thirdmagic.signature import SignatureConfig
+from thirdmagic.signature import Signature, SignatureConfig
 from thirdmagic.swarm import SwarmTaskSignature
 from thirdmagic.swarm.state import PublishState
 from thirdmagic.task import TaskSignature
 
+from mageflow.callbacks import AcceptParams
+
 
 @dataclass
 class SignatureTTLConfig:
     active_ttl: Optional[int] = None  # seconds, None = use general
-    ttl_when_sign_done: Optional[int] = None
+    ttl_when_sign_done: Optional[int] = Field(default=None, ge=REMOVED_TASK_TTL)
 
 
 @dataclass
 class TTLConfig:
     active_ttl: int = 24 * 60 * 60  # general active TTL (default 24h)
-    ttl_when_sign_done: int = REMOVED_TASK_TTL  # general done TTL (default 5min)
+    ttl_when_sign_done: int = Field(default=REMOVED_TASK_TTL, ge=REMOVED_TASK_TTL)
     task: SignatureTTLConfig = field(default_factory=SignatureTTLConfig)
     chain: SignatureTTLConfig = field(default_factory=SignatureTTLConfig)
     swarm: SignatureTTLConfig = field(default_factory=SignatureTTLConfig)
@@ -44,4 +47,5 @@ def apply_ttl_config(ttl_config: TTLConfig):
         done_ttl = sig_config.ttl_when_sign_done or ttl_config.ttl_when_sign_done
 
         sig_type.Meta = dataclasses.replace(sig_type.Meta, ttl=active_ttl)
-        sig_type.SignatureSettings = SignatureConfig(ttl_when_sign_done=done_ttl)
+        if issubclass(sig_type, Signature):
+            sig_type.SignatureSettings = SignatureConfig(ttl_when_sign_done=done_ttl)

{mageflow-0.3.2 → mageflow-0.3.4}/mageflow/lifecycle/signature.py

@@ -2,7 +2,6 @@ import asyncio
 from typing import Any, Optional, cast
 
 import rapyer
-
 from thirdmagic.clients.lifecycle import BaseLifecycle
 from thirdmagic.container import ContainerTaskSignature
 from thirdmagic.signature import Signature

{mageflow-0.3.2 → mageflow-0.3.4}/mageflow/lifecycle/task.py

@@ -1,7 +1,6 @@
 from typing import Any
 
 from pydantic import BaseModel
-
 from thirdmagic.clients.lifecycle import BaseLifecycle
 
 

{mageflow-0.3.2 → mageflow-0.3.4}/mageflow/swarm/workflows.py

@@ -3,7 +3,6 @@ from typing import Any, Optional, cast
 
 import rapyer
 from rapyer.fields import RapyerKey
-
 from thirdmagic.clients.lifecycle import BaseLifecycle
 from thirdmagic.signature import Signature
 from thirdmagic.swarm import PublishState

mageflow-0.3.4/mageflow/testing/__init__.py

@@ -0,0 +1,15 @@
+from mageflow.testing._adapter import (
+    ChainDispatchRecord,
+    RecordedDispatch,
+    SwarmDispatchRecord,
+    TaskDispatchRecord,
+    TestClientAdapter,
+)
+
+__all__ = [
+    "TestClientAdapter",
+    "RecordedDispatch",
+    "TaskDispatchRecord",
+    "SwarmDispatchRecord",
+    "ChainDispatchRecord",
+]