mageflow 0.3.0__tar.gz → 0.3.2__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mageflow
-Version: 0.3.0
+Version: 0.3.2
 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
+Requires-Dist: thirdmagic<0.1.0,>=0.0.3
 Provides-Extra: dev
 Requires-Dist: black>=26.1.0; extra == 'dev'
 Requires-Dist: coverage[toml]<8.0.0,>=7.0.0; extra == 'dev'
@@ -36,5 +36,6 @@ Requires-Dist: psutil<8.0.0,>=7.1.3; extra == 'dev'
 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'
 Provides-Extra: hatchet
 Requires-Dist: hatchet-sdk<1.24.0,>=1.22.5; extra == 'hatchet'

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

@@ -1,14 +1,16 @@
 import rapyer
 from rapyer.fields import RapyerKey
-from thirdmagic.chain.creator import chain as achain
-from thirdmagic.signature import Signature
-from thirdmagic.task import sign as asign, TaskSignature
-from thirdmagic.swarm.creator import swarm as aswarm
 
 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.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
 
-resume_task = TaskSignature.resume_from_key
 lock_task = TaskSignature.alock_from_key
 resume = TaskSignature.resume_from_key
 pause = TaskSignature.pause_from_key
@@ -20,11 +22,11 @@ async def load_sign(key: RapyerKey) -> Signature:
 
 
 load_signature = rapyer.afind_one
+abounded_field = rapyer.apipeline
 
 
 __all__ = [
     "load_signature",
-    "resume_task",
     "lock_task",
     "resume",
     "remove",
@@ -32,6 +34,11 @@ __all__ = [
     "asign",
     "handle_task_callback",
     "Mageflow",
+    "MageflowConfig",
+    "TTLConfig",
+    "SignatureTTLConfig",
     "achain",
     "aswarm",
+    "start_mageflow",
+    "abounded_field",
 ]

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

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

{mageflow-0.3.0 → mageflow-0.3.2}/mageflow/client.py

@@ -1,29 +1,47 @@
 import os
+import warnings
 from typing import TypeVar, overload
 
 import redis
 from hatchet_sdk import Hatchet
 from redis.asyncio import Redis
-from thirdmagic.signature import Signature
 
 from mageflow.callbacks import AcceptParams
 from mageflow.clients.hatchet.adapter import HatchetClientAdapter
 from mageflow.clients.hatchet.mageflow import HatchetMageflow
+from mageflow.config import MageflowConfig
+from thirdmagic.signature import Signature
 
 T = TypeVar("T")
 
 
 @overload
 def Mageflow(
-    hatchet_client: Hatchet, redis_client: Redis | str = None
+    hatchet_client: Hatchet,
+    redis_client: Redis | str = None,
+    param_config: AcceptParams = None,
+    config: MageflowConfig = None,
 ) -> HatchetMageflow: ...
 
 
 def Mageflow(
     hatchet_client: T = None,
     redis_client: Redis | str = None,
-    param_config: AcceptParams = AcceptParams.NO_CTX,
+    param_config: AcceptParams = None,
+    config: MageflowConfig = None,
 ) -> T:
+    if config is None:
+        config = MageflowConfig()
+
+    if param_config is not None:
+        warnings.warn(
+            "Passing 'param_config' directly to Mageflow() is deprecated. "
+            "Set it via MageflowConfig(param_config=...) instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        config.param_config = param_config
+
     if hatchet_client is None:
         hatchet_client = Hatchet()
 
@@ -35,4 +53,4 @@ def Mageflow(
         redis_client = redis.asyncio.from_url(redis_url, decode_responses=True)
     if isinstance(redis_client, str):
         redis_client = redis.asyncio.from_url(redis_client, decode_responses=True)
-    return HatchetMageflow(hatchet_client, redis_client, param_config)
+    return HatchetMageflow(hatchet_client, redis_client, config)

mageflow-0.3.2/mageflow/clients/__init__.py

@@ -0,0 +1,3 @@
+from thirdmagic.clients.base import BaseClientAdapter, DefaultClientAdapter
+
+__all__ = ["BaseClientAdapter", "DefaultClientAdapter"]

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

@@ -1,20 +1,13 @@
 from typing import Any, cast
 
 import rapyer
-from hatchet_sdk import Hatchet, NonRetryableException, Context
+from hatchet_sdk import Context, Hatchet, NonRetryableException
 from hatchet_sdk.clients.admin import TriggerWorkflowOptions
 from hatchet_sdk.runnables.contextvars import ctx_additional_metadata
 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
@@ -23,16 +16,22 @@ from mageflow.clients.inner_task_names import (
     ON_CHAIN_ERROR,
     ON_SWARM_ITEM_DONE,
     ON_SWARM_ITEM_ERROR,
+    SWARM_FILL_TASK,
 )
-from mageflow.clients.inner_task_names import SWARM_FILL_TASK
 from mageflow.lifecycle.signature import SignatureLifecycle
 from mageflow.lifecycle.task import TaskLifecycle
 from mageflow.swarm.messages import (
-    SwarmMessage,
-    SwarmResultsMessage,
-    SwarmErrorMessage,
     FillSwarmMessage,
+    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):

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

@@ -3,43 +3,36 @@ import functools
 import inspect
 import random
 from datetime import timedelta
-from typing import Any, Unpack, Callable, TypedDict
+from typing import Any, Callable, TypedDict, Unpack
 
-from hatchet_sdk import Hatchet, Worker, Context
+from hatchet_sdk import Context, Hatchet, Worker
 from hatchet_sdk.labels import DesiredWorkerLabel
 from hatchet_sdk.rate_limit import RateLimit
 from hatchet_sdk.runnables.types import (
-    StickyStrategy,
     ConcurrencyExpression,
-    DefaultFilter,
     ConcurrencyLimitStrategy,
+    DefaultFilter,
+    StickyStrategy,
 )
 from hatchet_sdk.runnables.workflow import BaseWorkflow, Standalone
 from hatchet_sdk.worker.worker import LifespanFn
 from redis.asyncio import Redis
-from thirdmagic import sign, chain
-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 TaskSignatureConvertible, TaskSignature, TaskInputType
-from thirdmagic.task_def import MageflowTaskDefinition
-from thirdmagic.utils import HatchetTaskType
 from typing_extensions import override
 
 from mageflow.callbacks import AcceptParams, handle_task_callback
 from mageflow.chain.messages import ChainCallbackMessage, ChainErrorMessage
 from mageflow.chain.workflows import chain_end_task, chain_error_task
 from mageflow.clients.inner_task_names import (
-    ON_CHAIN_ERROR,
     ON_CHAIN_END,
-    SWARM_FILL_TASK,
-    ON_SWARM_ITEM_ERROR,
+    ON_CHAIN_ERROR,
     ON_SWARM_ITEM_DONE,
+    ON_SWARM_ITEM_ERROR,
+    SWARM_FILL_TASK,
 )
+from mageflow.config import MageflowConfig
 from mageflow.startup import (
-    lifespan_initialize,
     init_mageflow,
+    lifespan_initialize,
     teardown_mageflow,
 )
 from mageflow.swarm.consts import SWARM_TASK_ID_PARAM_NAME
@@ -50,10 +43,18 @@ from mageflow.swarm.messages import (
 )
 from mageflow.swarm.workflows import (
     fill_swarm_running_tasks,
-    swarm_item_failed,
     swarm_item_done,
+    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
 
@@ -86,9 +87,12 @@ class WorkerOptions(TypedDict, total=False):
 
 
 async def merge_lifespan(
-    redis: Redis, tasks: list[MageflowTaskDefinition], original_lifespan: LifespanFn
+    redis: Redis,
+    tasks: list[MageflowTaskDefinition],
+    config: MageflowConfig,
+    original_lifespan: LifespanFn,
 ):
-    await init_mageflow(redis, tasks)
+    await init_mageflow(redis, tasks, config)
     async for res in original_lifespan():
         yield res
     await teardown_mageflow()
@@ -99,12 +103,12 @@ class HatchetMageflow(Hatchet):
         self,
         hatchet: Hatchet,
         redis_client: Redis,
-        param_config: AcceptParams = AcceptParams.NO_CTX,
+        config: MageflowConfig = None,
     ):
         super().__init__(client=hatchet._client)
         self.hatchet = hatchet
         self.redis = redis_client
-        self.param_config = param_config
+        self.mageflow_config = config or MageflowConfig()
         self._task_defs: list[MageflowTaskDefinition] = []
 
     @property
@@ -123,7 +127,9 @@ class HatchetMageflow(Hatchet):
 
     def task_decorator(self, func: Callable, hatchet_task):
         param_config = (
-            AcceptParams.ALL if does_task_wants_ctx(func) else self.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)
@@ -227,11 +233,15 @@ class HatchetMageflow(Hatchet):
         workflows += mageflow_flows
         if lifespan is None:
             lifespan = functools.partial(
-                lifespan_initialize, self.redis, self._task_defs
+                lifespan_initialize, self.redis, self._task_defs, self.mageflow_config
             )
         else:
             lifespan = functools.partial(
-                merge_lifespan, self.redis, self._task_defs, lifespan
+                merge_lifespan,
+                self.redis,
+                self._task_defs,
+                self.mageflow_config,
+                lifespan,
             )
 
         return super().worker(name, workflows=workflows, lifespan=lifespan, **kwargs)

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

@@ -3,6 +3,7 @@ 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/config.py

@@ -0,0 +1,47 @@
+import dataclasses
+from dataclasses import dataclass, field
+from typing import Optional
+
+from mageflow.callbacks import AcceptParams
+from thirdmagic.chain import ChainTaskSignature
+from thirdmagic.consts import REMOVED_TASK_TTL
+from thirdmagic.signature import SignatureConfig
+from thirdmagic.swarm import SwarmTaskSignature
+from thirdmagic.swarm.state import PublishState
+from thirdmagic.task import TaskSignature
+
+
+@dataclass
+class SignatureTTLConfig:
+    active_ttl: Optional[int] = None  # seconds, None = use general
+    ttl_when_sign_done: Optional[int] = None
+
+
+@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)
+    task: SignatureTTLConfig = field(default_factory=SignatureTTLConfig)
+    chain: SignatureTTLConfig = field(default_factory=SignatureTTLConfig)
+    swarm: SignatureTTLConfig = field(default_factory=SignatureTTLConfig)
+
+
+@dataclass
+class MageflowConfig:
+    ttl: TTLConfig = field(default_factory=TTLConfig)
+    param_config: AcceptParams = AcceptParams.NO_CTX
+
+
+def apply_ttl_config(ttl_config: TTLConfig):
+    config_mapping = {
+        TaskSignature: ttl_config.task,
+        ChainTaskSignature: ttl_config.chain,
+        SwarmTaskSignature: ttl_config.swarm,
+        PublishState: ttl_config.swarm,
+    }
+    for sig_type, sig_config in config_mapping.items():
+        active_ttl = sig_config.active_ttl or ttl_config.active_ttl
+        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)

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

@@ -1,7 +1,8 @@
 import asyncio
-from typing import Optional, Any, cast
+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.0 → mageflow-0.3.2}/mageflow/lifecycle/task.py

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

{mageflow-0.3.0 → mageflow-0.3.2}/mageflow/startup.py

@@ -1,9 +1,21 @@
 import rapyer
 from redis.asyncio import Redis
+
+from mageflow.config import MageflowConfig, apply_ttl_config
 from thirdmagic.task_def import MageflowTaskDefinition
 
 
-async def init_mageflow(redis: Redis, tasks: list[MageflowTaskDefinition]):
+async def start_mageflow(redis: Redis, config: MageflowConfig = None):
+    await init_mageflow(redis, [], config)
+
+
+async def init_mageflow(
+    redis: Redis,
+    tasks: list[MageflowTaskDefinition],
+    config: MageflowConfig = None,
+):
+    if config is not None:
+        apply_ttl_config(config.ttl)
     # Init redis in local async loop
     redis = Redis(**redis.connection_pool.connection_kwargs)
     await rapyer.init_rapyer(redis, prefer_normal_json_dump=True)
@@ -18,10 +30,14 @@ async def register_workflows(tasks: list[MageflowTaskDefinition]):
     await MageflowTaskDefinition.ainsert(*tasks)
 
 
-async def lifespan_initialize(redis: Redis, tasks: list[MageflowTaskDefinition]):
+async def lifespan_initialize(
+    redis: Redis,
+    tasks: list[MageflowTaskDefinition],
+    config: MageflowConfig = None,
+):
     # Init redis in local async loop
     redis = Redis(**redis.connection_pool.connection_kwargs)
-    await init_mageflow(redis, tasks)
+    await init_mageflow(redis, tasks, config)
     # yield makes the function usable as a Hatchet lifespan context manager (can also be used for FastAPI):
     # - code before yield runs at startup (init config, register workers, etc.)
     # - code after yield would run at shutdown

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

@@ -1,8 +1,9 @@
 from logging import Logger
-from typing import Optional, cast, Any
+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.2/mageflow/utils/mageflow.py

@@ -0,0 +1,2 @@
+def does_task_wants_ctx(func) -> bool:
+    return getattr(func, "__user_ctx__", False)

{mageflow-0.3.0 → mageflow-0.3.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mageflow"
-version = "0.3.0"
+version = "0.3.2"
 description = "Manage Graph Execution Flow - A unified interface for task orchestration across different task managers"
 authors = [
     {name = "imaginary-cherry", email = "yedidyakfir@gmail.com"}
@@ -41,7 +41,7 @@ classifiers = [
     "Operating System :: OS Independent",
 ]
 dependencies = [
-    "thirdmagic",
+    "thirdmagic>=0.0.3, <0.1.0",
 ]
 
 [project.optional-dependencies]
@@ -57,6 +57,7 @@ dev = [
     "coverage[toml]>=7.0.0,<8.0.0",
     "fakeredis[json,lua]>=2.34.0,<3.0.0",
     "black>=26.1.0",
+    "ruff>=0.15.5",
 ]
 
 [project.urls]

{mageflow-0.3.0 → mageflow-0.3.2}/tests/integration/conftest.py

@@ -2,6 +2,9 @@ import pytest_asyncio
 import redis.asyncio
 from dynaconf import Dynaconf
 
+from mageflow import start_mageflow
+from tests.integration.hatchet.worker import TEST_MAGEFLOW_CONFIG
+
 settings = Dynaconf(
     envvar_prefix="DYNACONF",
     settings_files=["settings.toml", ".secrets.toml"],
@@ -24,3 +27,8 @@ async def real_redis(redis_client):
     if delete_keys:
         await redis_client.delete(*delete_keys)
     await redis_client.aclose()
+
+
+@pytest_asyncio.fixture(scope="function", loop_scope="session", autouse=True)
+async def init_settings(real_redis):
+    await start_mageflow(real_redis, TEST_MAGEFLOW_CONFIG)

{mageflow-0.3.0 → mageflow-0.3.2}/tests/integration/hatchet/assertions.py

@@ -5,19 +5,19 @@ from hatchet_sdk import Hatchet
 from hatchet_sdk.clients.rest import V1LogLineList, V1TaskStatus, V1TaskSummary
 from hatchet_sdk.runnables.workflow import TaskRunRef
 from pydantic import BaseModel
+
+from tests.integration.hatchet.conftest import extract_bad_keys_from_redis
+from tests.integration.hatchet.worker import MAX_DONE_TTL
 from thirdmagic.chain.model import ChainTaskSignature
 from thirdmagic.consts import (
     MAGEFLOW_TASK_INITIALS,
     TASK_ID_PARAM_NAME,
-    REMOVED_TASK_TTL,
 )
 from thirdmagic.signature import Signature
 from thirdmagic.swarm.model import SwarmTaskSignature
 from thirdmagic.task import TaskSignature
 from thirdmagic.utils import return_value_field
 
-from tests.integration.hatchet.conftest import extract_bad_keys_from_redis
-
 WF_MAPPING_TYPE = dict[str, V1TaskSummary]
 WF_MAPPING_BY_WF_ID_TYPE = dict[str, V1TaskSummary]
 HatchetRuns = list[V1TaskSummary]
@@ -209,11 +209,11 @@ async def assert_redis_is_clean(redis_client):
     keys_with_invalid_ttl = [
         (key, ttl)
         for key, ttl in zip(non_persistent_keys, ttls)
-        if ttl == -1 or ttl > REMOVED_TASK_TTL
+        if ttl == -1 or ttl > MAX_DONE_TTL
     ]
     assert (
         len(keys_with_invalid_ttl) == 0
-    ), f"Keys without proper TTL (should be <= {REMOVED_TASK_TTL}s): {keys_with_invalid_ttl}"
+    ), f"Keys without proper TTL (should be <= {MAX_DONE_TTL}s): {keys_with_invalid_ttl}"
 
 
 def assert_task_was_paused(runs: HatchetRuns, task: TaskSignature, with_resume=False):

{mageflow-0.3.0 → mageflow-0.3.2}/tests/integration/hatchet/chain/test_stop_resume.py

@@ -1,20 +1,19 @@
 import asyncio
 
 import pytest
-from thirdmagic.task import TaskSignature
 
 import mageflow
 from tests.integration.hatchet.assertions import (
+    assert_chain_done,
+    assert_redis_is_clean,
     assert_signature_not_called,
     assert_task_was_paused,
-    assert_redis_is_clean,
-    assert_chain_done,
     get_runs,
 )
 from tests.integration.hatchet.conftest import HatchetInitData
 from tests.integration.hatchet.models import ContextMessage
-from tests.integration.hatchet.worker import sleep_task
-from tests.integration.hatchet.worker import task2_with_result
+from tests.integration.hatchet.worker import sleep_task, task2_with_result
+from thirdmagic.task import TaskSignature
 
 
 @pytest.mark.asyncio(loop_scope="session")

{mageflow-0.3.0 → mageflow-0.3.2}/tests/integration/hatchet/conftest.py

@@ -11,7 +11,7 @@ from datetime import datetime
 from io import BytesIO
 from pathlib import Path
 from threading import Thread
-from typing import Generator, Callable, AsyncGenerator
+from typing import AsyncGenerator, Callable, Generator
 
 import psutil
 import pytest
@@ -23,21 +23,20 @@ from hatchet_sdk.clients.admin import TriggerWorkflowOptions
 from hatchet_sdk.clients.rest import V1TaskStatus
 from hatchet_sdk.features.runs import BulkCancelReplayOpts, RunFilter
 from redis.asyncio.client import Redis
-from thirdmagic.task_def import MageflowTaskDefinition
 
 import mageflow
 from mageflow import Mageflow
 from mageflow.client import HatchetMageflow
-from mageflow.startup import init_mageflow
 from tests.integration.hatchet.worker import (
+    chain_callback,
     config_obj,
+    fail_task,
     task1,
+    task1_callback,
     task2,
     task3,
-    task1_callback,
-    fail_task,
-    chain_callback,
 )
+from thirdmagic.task_def import MageflowTaskDefinition
 
 # If redis key starts with one of these, it shouldn't be removed
 STATIC_REDIS_PREFIX_KEYS = [MageflowTaskDefinition.__name__]
@@ -79,16 +78,6 @@ async def hatchet_client_init(
     await rapyer.teardown_rapyer()
 
 
-@pytest_asyncio.fixture(scope="function", loop_scope="session", autouse=True)
-async def init_settings(hatchet_client_init: HatchetInitData):
-    redis_client, hatchet = (
-        hatchet_client_init.redis_client,
-        hatchet_client_init.hatchet,
-    )
-    # Load the subclasses of the task signature
-    await init_mageflow(redis_client, [])
-
-
 @pytest_asyncio.fixture(scope="session", loop_scope="session", autouse=True)
 async def hatchet_worker_deploy(
     redis_client,

{mageflow-0.3.0 → mageflow-0.3.2}/tests/integration/hatchet/models.py

@@ -1,6 +1,7 @@
 from typing import Any
 
 from pydantic import BaseModel, Field
+
 from thirdmagic.message import ReturnValue
 
 
@@ -41,5 +42,14 @@ class SleepTaskMessage(ContextMessage):
     result: Any = None
 
 
+class SignatureKeysResult(BaseModel):
+    task_keys: list[str]
+    chain_key: str
+    chain_sub_task_keys: list[str]
+    swarm_key: str
+    swarm_sub_task_keys: list[str]
+    publish_state_key: str
+
+
 class MageflowTestError(Exception):
     pass