thirdmagic 0.0.1__py3-none-any.whl

thirdmagic/__init__.py

@@ -0,0 +1,5 @@
+from thirdmagic.chain.creator import chain
+from thirdmagic.task.creator import sign
+from thirdmagic.swarm.creator import swarm
+
+__all__ = ["sign", "chain", "swarm"]

thirdmagic/chain/__init__.py

@@ -0,0 +1,4 @@
+from thirdmagic.chain.creator import chain
+from thirdmagic.chain.model import ChainTaskSignature
+
+__all__ = ["chain", "ChainTaskSignature"]

thirdmagic/chain/creator.py

@@ -0,0 +1,34 @@
+from thirdmagic.chain.model import ChainTaskSignature
+from thirdmagic.signature.model import TaskInputType
+from thirdmagic.task.creator import resolve_signatures, TaskSignatureConvertible
+
+
+async def chain(
+    tasks: list[TaskSignatureConvertible],
+    name: str = None,
+    error: TaskInputType = None,
+    success: TaskInputType = None,
+    **kwargs,
+) -> ChainTaskSignature:
+    if len(tasks) < 2:
+        raise ValueError(
+            "Chained tasks must contain at least two tasks. "
+            "If you want to run a single task, use `create_workflow` instead."
+        )
+    tasks = await resolve_signatures(tasks)
+
+    # Create a chain task that will be deleted only at the end of the chain
+    first_task = tasks[0]
+    chain_task_signature = ChainTaskSignature(
+        task_name=f"chain-task:{name or first_task.task_name}",
+        success_callbacks=[success] if success else [],
+        error_callbacks=[error] if error else [],
+        tasks=tasks,
+        kwargs=kwargs,
+    )
+    async with first_task.apipeline():
+        for task in tasks:
+            task.signature_container_id = chain_task_signature.key
+        await chain_task_signature.asave()
+
+    return chain_task_signature

thirdmagic/chain/model.py

@@ -0,0 +1,91 @@
+import asyncio
+from typing import cast, Any
+
+import rapyer
+from pydantic import field_validator, Field, BaseModel
+from rapyer.fields import RapyerKey
+
+from thirdmagic.container import ContainerTaskSignature
+from thirdmagic.errors import MissingSignatureError
+from thirdmagic.signature.status import SignatureStatus
+from thirdmagic.task.model import TaskSignature
+from thirdmagic.utils import HAS_HATCHET
+
+if HAS_HATCHET:
+    from hatchet_sdk.clients.admin import TriggerWorkflowOptions
+
+
+class ChainTaskSignature(ContainerTaskSignature):
+    tasks: list[RapyerKey] = Field(default_factory=list)
+
+    @field_validator("tasks", mode="before")
+    @classmethod
+    def validate_tasks(cls, v: list[TaskSignature]):
+        return [cls.validate_task_key(item) for item in v]
+
+    @property
+    def task_ids(self):
+        return self.tasks
+
+    async def on_sub_task_done(self, sub_task: TaskSignature, results: Any):
+        sub_task_idx = self.tasks.index(sub_task.key)
+        # If this is the last task, activate chain success callbacks
+        if sub_task_idx == len(self.tasks) - 1:
+            await self.ClientAdapter.acall_chain_done(results, self)
+        else:
+            next_task_key = self.tasks[sub_task_idx + 1]
+            next_task = await rapyer.aget(next_task_key)
+            next_task = cast(TaskSignature, next_task)
+            await next_task.acall(results, set_return_field=True, **self.kwargs)
+
+    async def on_sub_task_error(
+        self, sub_task: TaskSignature, error: BaseException, original_msg: dict
+    ):
+        await self.ClientAdapter.acall_chain_error(original_msg, error, self, sub_task)
+
+    async def sub_tasks(self) -> list[TaskSignature]:
+        sub_tasks = await rapyer.afind(*self.tasks, skip_missing=True)
+        return cast(list[TaskSignature], sub_tasks)
+
+    async def acall(self, msg: Any, set_return_field: bool = True, **kwargs):
+        first_task = await rapyer.afind_one(self.tasks[0])
+        if first_task is None:
+            raise MissingSignatureError(f"First task from chain {self.key} not found")
+
+        full_kwargs = self.kwargs | kwargs
+        return await first_task.acall(msg, set_return_field, **full_kwargs)
+
+    if HAS_HATCHET:
+
+        async def aio_run_no_wait(
+            self, msg: BaseModel, options: TriggerWorkflowOptions = None
+        ):
+            return await self.acall(msg, options=options, set_return_field=False)
+
+    async def change_status(self, status: SignatureStatus):
+        pause_chain_tasks = [
+            TaskSignature.safe_change_status(task, status) for task in self.tasks
+        ]
+        pause_chain = super().change_status(status)
+        await asyncio.gather(pause_chain, *pause_chain_tasks, return_exceptions=True)
+
+    async def suspend(self):
+        await asyncio.gather(
+            *[TaskSignature.suspend_from_key(task_id) for task_id in self.tasks],
+            return_exceptions=True,
+        )
+        await super().change_status(SignatureStatus.SUSPENDED)
+
+    async def interrupt(self):
+        await asyncio.gather(
+            *[TaskSignature.interrupt_from_key(task_id) for task_id in self.tasks],
+            return_exceptions=True,
+        )
+        await super().change_status(SignatureStatus.INTERRUPTED)
+
+    async def resume(self):
+        await asyncio.gather(
+            *[TaskSignature.resume_from_key(task_key) for task_key in self.tasks],
+            return_exceptions=True,
+        )
+        await super().change_status(self.task_status.last_status)

thirdmagic/clients/__init__.py

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

thirdmagic/clients/base.py

@@ -0,0 +1,150 @@
+import abc
+import asyncio
+from abc import ABC
+from typing import TYPE_CHECKING, Any
+
+from pydantic import BaseModel
+
+from thirdmagic.clients.lifecycle import BaseLifecycle
+from thirdmagic.task_def import MageflowTaskDefinition
+
+if TYPE_CHECKING:
+    from thirdmagic.task import TaskSignature
+    from thirdmagic.signature import Signature
+    from thirdmagic.chain.model import ChainTaskSignature
+    from thirdmagic.swarm.model import SwarmTaskSignature
+    from thirdmagic.utils import HatchetTaskType
+
+
+class BaseClientAdapter(ABC):
+    @abc.abstractmethod
+    def extract_validator(self, client_task) -> type[BaseModel]:
+        pass
+
+    @abc.abstractmethod
+    def extract_retries(self, client_task) -> int:
+        pass
+
+    @abc.abstractmethod
+    async def acall_chain_done(self, results: Any, chain: "ChainTaskSignature"):
+        pass
+
+    @abc.abstractmethod
+    async def acall_chain_error(
+        self,
+        original_msg: dict,
+        error: BaseException,
+        chain: "ChainTaskSignature",
+        failed_task: "Signature",
+    ):
+        pass
+
+    @abc.abstractmethod
+    async def afill_swarm(
+        self, swarm: "SwarmTaskSignature", max_tasks: int = None, **kwargs
+    ):
+        pass
+
+    @abc.abstractmethod
+    async def acall_swarm_item_error(
+        self, error: BaseException, swarm: "SwarmTaskSignature", swarm_item: "Signature"
+    ):
+        pass
+
+    @abc.abstractmethod
+    async def acall_swarm_item_done(
+        self, results: Any, swarm: "SwarmTaskSignature", swarm_item: "Signature"
+    ):
+        pass
+
+    @abc.abstractmethod
+    async def acall_signature(
+        self, signature: "TaskSignature", msg: Any, set_return_field: bool, **kwargs
+    ):
+        pass
+
+    async def acall_signatures(
+        self,
+        signatures: list["Signature"],
+        msg: Any,
+        set_return_field: bool,
+        **kwargs,
+    ):
+        return await asyncio.gather(
+            *[
+                signature.acall(msg, set_return_field, **kwargs)
+                for signature in signatures
+            ]
+        )
+
+    @abc.abstractmethod
+    def should_task_retry(
+        self, task_definition: MageflowTaskDefinition, attempt_num: int, e: BaseException
+    ) -> bool:
+        pass
+
+    @abc.abstractmethod
+    def task_name(self, task: "HatchetTaskType") -> str:
+        pass
+
+    @abc.abstractmethod
+    async def create_lifecycle(self, *args) -> BaseLifecycle:
+        pass
+
+    @abc.abstractmethod
+    async def lifecycle_from_signature(self, *args) -> BaseLifecycle:
+        pass
+
+
+class DefaultClientAdapter(BaseClientAdapter):
+    async def acall_chain_done(self, results: Any, chain: "ChainTaskSignature"):
+        raise NotImplementedError("Set a client before we start")
+
+    async def acall_chain_error(
+        self,
+        original_msg: dict,
+        error: BaseException,
+        chain: "ChainTaskSignature",
+        failed_task: "Signature",
+    ):
+        raise NotImplementedError("Set a client before we start")
+
+    async def afill_swarm(
+        self, swarm: "SwarmTaskSignature", max_tasks: int = None, **kwargs
+    ):
+        raise NotImplementedError("Set a client before we start")
+
+    async def acall_swarm_item_error(
+        self, error: BaseException, swarm: "SwarmTaskSignature", swarm_item: "Signature"
+    ):
+        raise NotImplementedError("Set a client before we start")
+
+    async def acall_swarm_item_done(
+        self, results: Any, swarm: "SwarmTaskSignature", swarm_item: "Signature"
+    ):
+        raise NotImplementedError("Set a client before we start")
+
+    def extract_validator(self, client_task) -> type[BaseModel]:
+        raise NotImplementedError("Set a client before we start")
+
+    def extract_retries(self, client_task) -> int:
+        raise NotImplementedError("Set a client before we start")
+
+    async def acall_signature(
+        self, signature: "Signature", set_return_field: bool, **kwargs
+    ):
+        raise NotImplementedError("Set a client before we start")
+
+    def should_task_retry(
+        self, task_definition: MageflowTaskDefinition, attempt_num: int, e: BaseException
+    ) -> bool:
+        raise NotImplementedError("Set a client before we start")
+
+    def task_name(self, task: "HatchetTaskType") -> str:
+        raise NotImplementedError("Set a client before we start")
+
+    def create_lifecycle(self, *args) -> BaseLifecycle:
+        raise NotImplementedError("Set a client before we start")
+
+    async def lifecycle_from_signature(self, *args) -> BaseLifecycle:
+        pass

thirdmagic/clients/lifecycle.py

@@ -0,0 +1,24 @@
+import abc
+from abc import ABC
+from typing import Any
+
+
+class BaseLifecycle(ABC):
+    @abc.abstractmethod
+    async def start_task(self):
+        pass
+
+    @abc.abstractmethod
+    async def task_success(self, result: Any):
+        pass
+
+    @abc.abstractmethod
+    async def task_failed(self, message: dict, error: BaseException):
+        pass
+
+    @abc.abstractmethod
+    async def should_run_task(self, message: dict) -> bool:
+        pass
+
+    def is_vanilla_run(self):
+        return False

thirdmagic/consts.py

@@ -0,0 +1,4 @@
+TASK_ID_PARAM_NAME = "task_id"
+
+MAGEFLOW_TASK_INITIALS = "mageflow_"
+REMOVED_TASK_TTL = 5 * 60

thirdmagic/container.py

@@ -0,0 +1,35 @@
+import abc
+import asyncio
+from abc import ABC
+from typing import Any
+
+from pydantic import BaseModel
+
+from thirdmagic.signature import Signature
+
+
+class ContainerTaskSignature(Signature, ABC):
+    @property
+    @abc.abstractmethod
+    def task_ids(self):
+        pass
+
+    @abc.abstractmethod
+    async def sub_tasks(self) -> list[Signature]:
+        pass
+
+    async def remove_references(self):
+        sub_tasks = await self.sub_tasks()
+        await asyncio.gather(
+            *[task.remove() for task in sub_tasks], return_exceptions=True
+        )
+
+    @abc.abstractmethod
+    async def on_sub_task_error(
+        self, sub_task: Signature, error: BaseException, original_msg: dict
+    ):
+        pass
+
+    @abc.abstractmethod
+    async def on_sub_task_done(self, sub_task: Signature, results: Any):
+        pass

thirdmagic/errors.py

@@ -0,0 +1,26 @@
+class MageflowError(Exception):
+    pass
+
+
+class MissingSignatureError(MageflowError):
+    pass
+
+
+class MissingSwarmItemError(MissingSignatureError):
+    pass
+
+
+class SwarmError(MageflowError):
+    pass
+
+
+class TooManyTasksError(SwarmError, RuntimeError):
+    pass
+
+
+class SwarmIsCanceledError(SwarmError, RuntimeError):
+    pass
+
+
+class UnrecognizedTaskError(MageflowError):
+    pass

thirdmagic/message.py

@@ -0,0 +1,26 @@
+import dataclasses
+from typing import Generic, TypeVar, Any, Annotated, TYPE_CHECKING, TypeAlias
+
+DEFAULT_RESULT_NAME = "mageflow_results"
+T = TypeVar("T")
+
+
+@dataclasses.dataclass(frozen=True)
+class ReturnValueAnnotation:
+    pass
+
+
+class _ReturnValue(Generic[T]):
+    def __new__(cls, typ: Any = None):
+        if typ is None:
+            return ReturnValueAnnotation()
+        return Annotated[typ, ReturnValueAnnotation()]
+
+    def __class_getitem__(cls, item):
+        return Annotated[item, ReturnValueAnnotation()]
+
+
+ReturnValue = _ReturnValue
+
+if TYPE_CHECKING:
+    ReturnValue: TypeAlias = Annotated[T, ReturnValueAnnotation()]

thirdmagic/signature/__init__.py

@@ -0,0 +1,4 @@
+from thirdmagic.signature.model import Signature
+from thirdmagic.signature.status import SignatureStatus, PauseActionTypes, TaskStatus
+
+__all__ = ["Signature", "SignatureStatus", "PauseActionTypes", "TaskStatus"]

thirdmagic/signature/model.py

@@ -0,0 +1,191 @@
+import abc
+import asyncio
+from abc import ABC
+from datetime import datetime
+from typing import Optional, Self, Any, TypeAlias, ClassVar, cast
+
+import rapyer
+from pydantic import BaseModel, field_validator, Field
+from rapyer import AtomicRedisModel
+from rapyer.config import RedisConfig
+from rapyer.errors.base import KeyNotFound
+from rapyer.fields import RapyerKey
+from rapyer.types import RedisDict, RedisList, RedisDatetime
+
+from thirdmagic.clients import BaseClientAdapter, DefaultClientAdapter
+from thirdmagic.consts import REMOVED_TASK_TTL
+from thirdmagic.signature.status import TaskStatus, PauseActionTypes, SignatureStatus
+from thirdmagic.utils import HAS_HATCHET
+
+if HAS_HATCHET:
+    from hatchet_sdk.clients.admin import TriggerWorkflowOptions
+
+
+class Signature(AtomicRedisModel, ABC):
+    task_name: str
+    kwargs: RedisDict[Any] = Field(default_factory=dict)
+    creation_time: RedisDatetime = Field(default_factory=datetime.now)
+    success_callbacks: RedisList[RapyerKey] = Field(default_factory=list)
+    error_callbacks: RedisList[RapyerKey] = Field(default_factory=list)
+    task_status: TaskStatus = Field(default_factory=TaskStatus)
+    signature_container_id: Optional[RapyerKey] = None
+
+    Meta: ClassVar[RedisConfig] = RedisConfig(ttl=24 * 60 * 60, refresh_ttl=False)
+    ClientAdapter: ClassVar[BaseClientAdapter] = DefaultClientAdapter()
+
+    @field_validator("success_callbacks", "error_callbacks", mode="before")
+    @classmethod
+    def validate_tasks_id(cls, v: list) -> list[str]:
+        return [cls.validate_task_key(item) for item in v]
+
+    @classmethod
+    def validate_task_key(cls, v) -> str:
+        if isinstance(v, bytes):
+            return RapyerKey(v.decode())
+        if isinstance(v, str):
+            return v
+        elif isinstance(v, Signature):
+            return v.key
+        else:
+            raise ValueError(f"Expected task ID or Signature, got {type(v).__name__}")
+
+    @abc.abstractmethod
+    async def acall(self, msg: Any, set_return_field: bool = True, **kwargs):
+        pass
+
+    if HAS_HATCHET:
+
+        @abc.abstractmethod
+        async def aio_run_no_wait(
+            self, msg: BaseModel, options: TriggerWorkflowOptions = None
+        ):
+            pass
+
+    async def on_pause_signature(self, msg: dict):
+        await self.kwargs.aupdate(**msg)
+
+    async def on_cancel_signature(self, msg: dict):
+        await self.remove()
+
+    async def activate_success(self, msg):
+        success_signatures = await rapyer.afind(*self.success_callbacks)
+        success_signatures = cast(list[Signature], success_signatures)
+        return await self.ClientAdapter.acall_signatures(success_signatures, msg, True)
+
+    async def activate_error(self, msg):
+        error_signatures = await rapyer.afind(*self.error_callbacks)
+        error_signatures = cast(list[Signature], error_signatures)
+        return await self.ClientAdapter.acall_signatures(error_signatures, msg, False)
+
+    async def remove_task(self):
+        await self.aset_ttl(REMOVED_TASK_TTL)
+
+    async def remove_branches(self, success: bool = True, errors: bool = True):
+        keys_to_remove = []
+        if errors:
+            keys_to_remove.extend([error_id for error_id in self.error_callbacks])
+        if success:
+            keys_to_remove.extend([success_id for success_id in self.success_callbacks])
+
+        signatures = cast(list[Signature], await rapyer.afind(*keys_to_remove))
+        await asyncio.gather(*[signature.remove() for signature in signatures])
+
+    async def remove_references(self):
+        pass
+
+    async def remove(self, with_error: bool = True, with_success: bool = True):
+        return await self._remove(with_error, with_success)
+
+    async def _remove(self, with_error: bool = True, with_success: bool = True):
+        await self.remove_branches(with_success, with_error)
+        await self.remove_references()
+        await self.remove_task()
+
+    @classmethod
+    async def remove_from_key(cls, task_key: RapyerKey):
+        async with rapyer.alock_from_key(task_key) as task:
+            task = cast(Signature, task)
+            return await task.remove()
+
+    async def should_run(self):
+        return self.task_status.should_run()
+
+    async def change_status(self, status: SignatureStatus):
+        await self.task_status.aupdate(
+            last_status=self.task_status.status, status=status
+        )
+
+    # When pausing task from outside the task
+    @classmethod
+    async def safe_change_status(cls, task_id: RapyerKey, status: SignatureStatus):
+        try:
+            async with rapyer.alock_from_key(task_id) as task:
+                task = cast(Signature, task)
+                return await task.change_status(status)
+        except Exception as e:
+            return False
+
+    @classmethod
+    async def resume_from_key(cls, task_key: RapyerKey):
+        async with rapyer.alock_from_key(task_key) as task:
+            task = cast(Signature, task)
+            await task.resume()
+
+    @abc.abstractmethod
+    async def resume(self):
+        pass
+
+    @classmethod
+    async def suspend_from_key(cls, task_key: RapyerKey):
+        async with rapyer.alock_from_key(task_key) as task:
+            task = cast(Signature, task)
+            await task.suspend()
+
+    async def done(self):
+        await self.task_status.aupdate(
+            last_status=self.task_status.status, status=SignatureStatus.DONE
+        )
+
+    async def failed(self):
+        await self.task_status.aupdate(
+            last_status=self.task_status.status, status=SignatureStatus.FAILED
+        )
+
+    async def suspend(self):
+        """
+        Task suspension will try and stop the task at before it starts
+        """
+        await self.change_status(SignatureStatus.SUSPENDED)
+
+    @classmethod
+    async def interrupt_from_key(cls, task_key: RapyerKey):
+        async with rapyer.alock_from_key(task_key) as task:
+            task = cast(Signature, task)
+            return task.interrupt()
+
+    async def interrupt(self):
+        """
+        Task interrupt will try to aggressively take hold of the async loop and stop the task
+        """
+        # TODO - not implemented yet - implement
+        await self.suspend()
+
+    @classmethod
+    async def pause_from_key(
+        cls,
+        task_key: RapyerKey,
+        pause_type: PauseActionTypes = PauseActionTypes.SUSPEND,
+    ):
+        async with rapyer.alock_from_key(task_key) as task:
+            task = cast(Signature, task)
+            await task.pause_task(pause_type)
+
+    async def pause_task(self, pause_type: PauseActionTypes = PauseActionTypes.SUSPEND):
+        if pause_type == PauseActionTypes.SUSPEND:
+            return await self.suspend()
+        elif pause_type == PauseActionTypes.INTERRUPT:
+            return await self.interrupt()
+        raise NotImplementedError(f"Pause type {pause_type} not supported")
+
+
+TaskInputType: TypeAlias = RapyerKey | Signature

thirdmagic/signature/status.py

@@ -0,0 +1,32 @@
+from enum import Enum
+from typing import ClassVar
+
+from rapyer import AtomicRedisModel
+from rapyer.config import RedisConfig
+
+
+class SignatureStatus(str, Enum):
+    PENDING = "pending"
+    ACTIVE = "active"
+    FAILED = "failed"
+    DONE = "done"
+    SUSPENDED = "suspended"
+    INTERRUPTED = "interrupted"
+    CANCELED = "canceled"
+
+
+class PauseActionTypes(str, Enum):
+    SUSPEND = "soft"
+    INTERRUPT = "hard"
+
+
+class TaskStatus(AtomicRedisModel):
+    status: SignatureStatus = SignatureStatus.PENDING
+    last_status: SignatureStatus = SignatureStatus.PENDING
+    Meta: ClassVar[RedisConfig] = RedisConfig(ttl=24 * 60 * 60, refresh_ttl=False)
+
+    def is_canceled(self):
+        return self.status in [SignatureStatus.CANCELED]
+
+    def should_run(self):
+        return self.status in [SignatureStatus.PENDING, SignatureStatus.ACTIVE]

thirdmagic/swarm/__init__.py

@@ -0,0 +1,5 @@
+from thirdmagic.swarm.creator import swarm
+from thirdmagic.swarm.model import SwarmTaskSignature, SwarmConfig
+from thirdmagic.swarm.state import PublishState
+
+__all__ = ["swarm", "SwarmTaskSignature", "SwarmConfig", "PublishState"]

thirdmagic/swarm/consts.py

@@ -0,0 +1 @@
+SWARM_MESSAGE_PARAM_NAME = "__mageflow_swarm_message__"

thirdmagic/swarm/creator.py

@@ -0,0 +1,53 @@
+import uuid
+from typing import overload, Any
+
+import rapyer
+
+from thirdmagic.task.creator import TaskSignatureOptions, TaskSignatureConvertible
+from thirdmagic.swarm.model import SwarmConfig, SwarmTaskSignature
+from thirdmagic.swarm.state import PublishState
+from thirdmagic.typing_support import Unpack
+
+
+class SignatureOptions(TaskSignatureOptions):
+    is_swarm_closed: bool
+    config: SwarmConfig
+    task_kwargs: dict
+
+
+@overload
+async def swarm(
+    tasks: list[TaskSignatureConvertible],
+    task_name: str = None,
+    **options: Unpack[TaskSignatureOptions],
+) -> SwarmTaskSignature: ...
+@overload
+async def swarm(
+    tasks: list[TaskSignatureConvertible], task_name: str = None, **options: Any
+) -> SwarmTaskSignature: ...
+
+
+async def swarm(
+    tasks: list[TaskSignatureConvertible] = None,
+    task_name: str = None,
+    **options: Unpack[SignatureOptions],
+) -> SwarmTaskSignature:
+    tasks = tasks or []
+    task_name = task_name or f"swarm-task-{uuid.uuid4()}"
+    publish_state = PublishState()
+    model_fields = list(SwarmTaskSignature.model_fields.keys())
+    direct_kwargs_param = options.pop("kwargs", {})
+    kwargs = {
+        field_name: options.pop(field_name)
+        for field_name in model_fields
+        if field_name in options
+    }
+    swarm_signature = SwarmTaskSignature(
+        **kwargs,
+        task_name=task_name,
+        publishing_state_id=publish_state.key,
+        kwargs=direct_kwargs_param | options,
+    )
+    await rapyer.ainsert(publish_state, swarm_signature)
+    await swarm_signature.add_tasks(tasks)
+    return swarm_signature

thirdmagic/swarm/model.py

@@ -0,0 +1,222 @@
+import asyncio
+from typing import Self, Any, Optional, cast
+
+import rapyer
+from pydantic import Field, field_validator, BaseModel
+from rapyer import AtomicRedisModel
+from rapyer.fields import RapyerKey
+from rapyer.types import RedisList, RedisInt
+
+from thirdmagic.consts import REMOVED_TASK_TTL
+from thirdmagic.container import ContainerTaskSignature
+from thirdmagic.errors import TooManyTasksError, SwarmIsCanceledError
+from thirdmagic.signature import Signature
+from thirdmagic.signature.status import SignatureStatus
+from thirdmagic.swarm.consts import SWARM_MESSAGE_PARAM_NAME
+from thirdmagic.swarm.state import PublishState
+from thirdmagic.task.creator import TaskSignatureConvertible, resolve_signatures
+from thirdmagic.task.model import TaskSignature
+from thirdmagic.utils import HAS_HATCHET
+
+if HAS_HATCHET:
+    from hatchet_sdk.clients.admin import TriggerWorkflowOptions
+    from hatchet_sdk.runnables.types import EmptyModel
+    from hatchet_sdk.runnables.workflow import TaskRunRef
+
+
+class SwarmConfig(AtomicRedisModel):
+    max_concurrency: int = 30
+    stop_after_n_failures: Optional[int] = Field(default=None, gt=0)
+    max_task_allowed: Optional[int] = Field(default=None, gt=0)
+    send_swarm_message_to_return_field: bool = False
+
+    def can_add_task(self, swarm: "SwarmTaskSignature") -> bool:
+        return self.can_add_n_tasks(swarm, 1)
+
+    def can_add_n_tasks(self, swarm: "SwarmTaskSignature", n: int) -> bool:
+        if self.max_task_allowed is None:
+            return True
+        return len(swarm.tasks) + n <= self.max_task_allowed
+
+
+class SwarmTaskSignature(ContainerTaskSignature):
+    # TODO - TASKS list should be set once we enable this in rapyer
+    tasks: RedisList[RapyerKey] = Field(default_factory=list)
+    tasks_left_to_run: RedisList[RapyerKey] = Field(default_factory=list)
+    finished_tasks: RedisList[RapyerKey] = Field(default_factory=list)
+    failed_tasks: RedisList[RapyerKey] = Field(default_factory=list)
+    tasks_results: RedisList[Any] = Field(default_factory=list)
+    # This flag is raised when no more tasks can be added to the swarm
+    is_swarm_closed: bool = False
+    # How many tasks can be added to the swarm at a time
+    current_running_tasks: RedisInt = 0
+    publishing_state_id: str
+    config: SwarmConfig = Field(default_factory=SwarmConfig)
+
+    @field_validator(
+        "tasks", "tasks_left_to_run", "finished_tasks", "failed_tasks", mode="before"
+    )
+    @classmethod
+    def validate_tasks(cls, v):
+        return [cls.validate_task_key(item) for item in v]
+
+    @property
+    def task_ids(self):
+        return self.tasks
+
+    async def sub_tasks(self) -> list[TaskSignature]:
+        tasks = await rapyer.afind(*self.tasks)
+        return cast(list[TaskSignature], tasks)
+
+    async def on_sub_task_done(self, sub_task: TaskSignature, results: Any):
+        await self.ClientAdapter.acall_swarm_item_done(results, self, sub_task)
+
+    async def on_sub_task_error(
+        self, sub_task: TaskSignature, error: BaseException, original_msg: BaseModel
+    ):
+        await self.ClientAdapter.acall_swarm_item_error(error, self, sub_task)
+
+    async def acall(self, msg: Any, set_return_field: bool = True, **kwargs):
+        # We update the kwargs that everyone are using, we also tell weather we should put this in the Return value or just in the message
+        async with self.apipeline():
+            self.kwargs.update(**{SWARM_MESSAGE_PARAM_NAME: msg})
+            self.config.send_swarm_message_to_return_field = set_return_field
+        return await self.ClientAdapter.afill_swarm(self, **kwargs)
+
+    if HAS_HATCHET:
+
+        async def aio_run_no_wait(
+            self, msg: BaseModel, options: "TriggerWorkflowOptions" = None
+        ):
+            return await self.acall(
+                msg.model_dump(mode="json", exclude_unset=True),
+                set_return_field=False,
+                options=options,
+            )
+
+        async def aio_run_in_swarm(
+            self,
+            task: TaskSignatureConvertible,
+            msg: BaseModel,
+            options: TriggerWorkflowOptions = None,
+            close_on_max_task: bool = True,
+        ) -> Optional["TaskRunRef"]:
+            sub_task = await self.add_task(task, close_on_max_task)
+            await sub_task.kwargs.aupdate(**msg.model_dump(mode="json"))
+            return await self.ClientAdapter.afill_swarm(
+                self, max_tasks=1, options=options
+            )
+
+    async def change_status(self, status: SignatureStatus):
+        paused_chain_tasks = [
+            TaskSignature.safe_change_status(task, status) for task in self.tasks
+        ]
+        pause_chain = super().change_status(status)
+        await asyncio.gather(pause_chain, *paused_chain_tasks, return_exceptions=True)
+
+    async def add_tasks(
+        self, tasks: list[TaskSignatureConvertible], close_on_max_task: bool = True
+    ) -> list[Signature]:
+        """
+        tasks - tasks signature to add to swarm
+        close_on_max_task - if true, and you set max task allowed on swarm, this swarm will close if the task reached maximum capacity
+        """
+        if not self.config.can_add_n_tasks(self, len(tasks)):
+            raise TooManyTasksError(
+                f"Swarm {self.task_name} has reached max tasks limit"
+            )
+        if self.task_status.is_canceled():
+            raise SwarmIsCanceledError(
+                f"Swarm {self.task_name} is {self.task_status} - can't add task"
+            )
+
+        tasks = await resolve_signatures(tasks)
+        task_keys = [task.key for task in tasks]
+
+        async with self.apipeline():
+            for task in tasks:
+                task.signature_container_id = self.key
+            self.tasks.extend(task_keys)
+            self.tasks_left_to_run.extend(task_keys)
+
+        if close_on_max_task and not self.config.can_add_task(self):
+            await self.close_swarm()
+
+        return tasks
+
+    async def add_task(
+        self, task: TaskSignatureConvertible, close_on_max_task: bool = True
+    ) -> Signature:
+        """
+        task - task signature to add to swarm
+        close_on_max_task - if true, and you set max task allowed on swarm, this swarm will close if the task reached maximum capacity
+        """
+        added_tasks = await self.add_tasks([task], close_on_max_task)
+        return added_tasks[0]
+
+    async def is_swarm_done(self):
+        done_tasks = self.finished_tasks + self.failed_tasks
+        finished_all_tasks = set(done_tasks) == set(self.tasks)
+        return self.is_swarm_closed and finished_all_tasks
+
+    def has_published_callback(self):
+        return self.task_status.status == SignatureStatus.DONE
+
+    def has_published_errors(self):
+        return self.task_status.status == SignatureStatus.FAILED
+
+    async def activate_success(self, msg):
+        results = await self.tasks_results.aload()
+        tasks_results = [res for res in results]
+
+        await super().activate_success(tasks_results)
+        await self.remove_branches(success=False)
+        await self.remove_task()
+
+    async def suspend(self):
+        await asyncio.gather(
+            *[TaskSignature.suspend_from_key(swarm_id) for swarm_id in self.tasks],
+            return_exceptions=True,
+        )
+        await super().change_status(SignatureStatus.SUSPENDED)
+
+    async def resume(self):
+        await asyncio.gather(
+            *[TaskSignature.resume_from_key(task_id) for task_id in self.tasks],
+            return_exceptions=True,
+        )
+        await super().change_status(self.task_status.last_status)
+
+    async def close_swarm(self) -> Self:
+        await self.aupdate(is_swarm_closed=True)
+        await self.ClientAdapter.afill_swarm(self, max_tasks=0)
+        return self
+
+    def has_swarm_failed(self):
+        should_stop_after_failures = self.config.stop_after_n_failures is not None
+        stop_after_n_failures = self.config.stop_after_n_failures or 0
+        too_many_errors = len(self.failed_tasks) >= stop_after_n_failures
+        return should_stop_after_failures and too_many_errors
+
+    async def finish_task(self, task_key: str, results: Any):
+        async with self.apipeline() as swarm_task:
+            # In case this was already updated
+            if task_key in swarm_task.finished_tasks:
+                return
+            swarm_task.finished_tasks.append(task_key)
+            swarm_task.tasks_results.append(results)
+            swarm_task.current_running_tasks -= 1
+
+    async def task_failed(self, task_key: str):
+        async with self.apipeline() as swarm_task:
+            if task_key in swarm_task.failed_tasks:
+                return
+            swarm_task.failed_tasks.append(task_key)
+            swarm_task.current_running_tasks -= 1
+
+    async def remove_task(self):
+        publish_state = await PublishState.aget(self.publishing_state_id)
+        async with self.apipeline():
+            # TODO - this should be removed once we use foreign key
+            await publish_state.aset_ttl(REMOVED_TASK_TTL)
+            return await super().remove_task()

thirdmagic/swarm/state.py

@@ -0,0 +1,13 @@
+from typing import ClassVar
+
+from pydantic import Field
+from rapyer import AtomicRedisModel
+from rapyer.config import RedisConfig
+from rapyer.fields import RapyerKey
+from rapyer.types import RedisList
+
+
+class PublishState(AtomicRedisModel):
+    task_ids: RedisList[RapyerKey] = Field(default_factory=list)
+
+    Meta: ClassVar[RedisConfig] = RedisConfig(ttl=24 * 60 * 60, refresh_ttl=False)

thirdmagic/task/__init__.py

@@ -0,0 +1,22 @@
+from thirdmagic.signature.model import TaskInputType
+from thirdmagic.signature.status import PauseActionTypes, TaskStatus, SignatureStatus
+from thirdmagic.task.creator import (
+    sign,
+    resolve_signatures,
+    resolve_signature,
+    TaskSignatureConvertible,
+)
+from thirdmagic.task.model import TaskSignature
+
+
+__all__ = [
+    "sign",
+    "resolve_signatures",
+    "resolve_signature",
+    "TaskSignatureConvertible",
+    "TaskSignature",
+    "SignatureStatus",
+    "TaskStatus",
+    "PauseActionTypes",
+    "TaskInputType",
+]

thirdmagic/task/creator.py

@@ -0,0 +1,87 @@
+from datetime import datetime
+from typing import TypedDict, Any, overload, TypeAlias, Optional
+
+import rapyer
+from rapyer.fields import RapyerKey
+
+from thirdmagic.signature import Signature
+from thirdmagic.signature.status import TaskStatus
+from thirdmagic.task.model import TaskSignature
+from thirdmagic.typing_support import Unpack
+from thirdmagic.utils import HatchetTaskType
+
+TaskSignatureConvertible: TypeAlias = RapyerKey | Signature | HatchetTaskType | str
+
+
+async def resolve_signatures(
+    tasks: list[TaskSignatureConvertible],
+) -> list[Signature]:
+    result: list[Optional[Signature]] = [None] * len(tasks)
+    identifier_entries: list[tuple[int, RapyerKey]] = []
+    hatchet_entries: list[tuple[int, HatchetTaskType]] = []
+    task_names: list[tuple[int, str]] = []
+
+    for i, task in enumerate(tasks):
+        if isinstance(task, Signature):
+            result[i] = task
+        elif isinstance(task, RapyerKey):
+            identifier_entries.append((i, task))
+        elif isinstance(task, str):
+            task_names.append((i, task))
+        else:
+            hatchet_entries.append((i, task))
+
+    if identifier_entries:
+        keys = [key for _, key in identifier_entries]
+        found = await rapyer.afind(*keys, skip_missing=True)
+        found_by_key = {sig.key: sig for sig in found}
+        for i, key in identifier_entries:
+            result[i] = found_by_key.get(key)
+
+    if hatchet_entries:
+        async with rapyer.apipeline():
+            for i, task in hatchet_entries:
+                result[i] = await TaskSignature.from_task(task)
+
+    if task_names:
+        async with rapyer.apipeline():
+            for i, task_name in task_names:
+                result[i] = await TaskSignature.from_task_name(task_name)
+
+    return result
+
+
+async def resolve_signature(task: TaskSignatureConvertible) -> Signature:
+    signatures = await resolve_signatures([task])
+    return signatures[0]
+
+
+class TaskSignatureOptions(TypedDict, total=False):
+    kwargs: dict
+    creation_time: datetime
+    model_validators: Any
+    success_callbacks: list[RapyerKey]
+    error_callbacks: list[RapyerKey]
+    task_status: TaskStatus
+
+
+@overload
+async def sign(
+    task: str | HatchetTaskType, **options: Unpack[TaskSignatureOptions]
+) -> TaskSignature: ...
+@overload
+async def sign(task: str | HatchetTaskType, **options: Any) -> TaskSignature: ...
+
+
+async def sign(task: str | HatchetTaskType, **options: Any) -> TaskSignature:
+    model_fields = list(TaskSignature.model_fields.keys())
+    kwargs = {
+        field_name: options.pop(field_name)
+        for field_name in model_fields
+        if field_name in options
+    }
+
+    if isinstance(task, str):
+        return await TaskSignature.from_task_name(task, kwargs=options, **kwargs)
+    else:
+        return await TaskSignature.from_task(task, kwargs=options, **kwargs)

thirdmagic/task/model.py

@@ -0,0 +1,83 @@
+from typing import Optional, Self, Any
+
+from pydantic import BaseModel
+from rapyer.fields import SafeLoad, RapyerKey
+
+from thirdmagic.errors import UnrecognizedTaskError
+from thirdmagic.message import DEFAULT_RESULT_NAME
+from thirdmagic.signature import Signature
+from thirdmagic.signature.status import SignatureStatus
+from thirdmagic.task_def import MageflowTaskDefinition
+from thirdmagic.utils import return_value_field, HAS_HATCHET, HatchetTaskType
+
+if HAS_HATCHET:
+    from hatchet_sdk.clients.admin import TriggerWorkflowOptions
+
+
+class TaskSignature(Signature):
+    model_validators: SafeLoad[Optional[type[BaseModel]]] = None
+    return_field_name: str = DEFAULT_RESULT_NAME
+    worker_task_id: str = ""
+
+    @classmethod
+    async def from_task(
+        cls,
+        task: HatchetTaskType,
+        success_callbacks: list[RapyerKey | Self] = None,
+        error_callbacks: list[RapyerKey | Self] = None,
+        **kwargs,
+    ) -> Self:
+        validator = cls.ClientAdapter.extract_validator(task)
+        return_field_name = return_value_field(validator)
+        signature = cls(
+            task_name=cls.ClientAdapter.task_name(task),
+            model_validators=validator,
+            return_field_name=return_field_name,
+            success_callbacks=success_callbacks or [],
+            error_callbacks=error_callbacks or [],
+            **kwargs,
+        )
+        await signature.asave()
+        return signature
+
+    @classmethod
+    async def from_task_name(
+        cls, task_name: str, model_validators: type[BaseModel] = None, **kwargs
+    ) -> Self:
+        if not model_validators:
+            task_def = await MageflowTaskDefinition.afind_one(task_name)
+            if not task_def:
+                raise UnrecognizedTaskError(f"Task {task_name} was not initialized")
+            model_validators = task_def.input_validator
+            task_name = task_def.mageflow_task_name if task_def else task_name
+        return_field_name = return_value_field(model_validators)
+
+        signature = cls(
+            task_name=task_name,
+            return_field_name=return_field_name,
+            model_validators=model_validators,
+            **kwargs,
+        )
+        await signature.asave()
+        return signature
+
+    async def acall(self, msg: Any, set_return_field: bool = True, **kwargs):
+        return await self.ClientAdapter.acall_signature(
+            self, msg, set_return_field, **kwargs
+        )
+
+    if HAS_HATCHET:
+
+        async def aio_run_no_wait(
+            self, msg: BaseModel, options: TriggerWorkflowOptions = None
+        ):
+            params = dict(options=options) if options else {}
+            return await self.acall(msg, set_return_field=False, **params)
+
+    async def resume(self):
+        last_status = self.task_status.last_status
+        if last_status == SignatureStatus.ACTIVE:
+            await self.change_status(SignatureStatus.PENDING)
+            await self.ClientAdapter.acall_signature(self, None, set_return_field=False)
+        else:
+            await self.change_status(last_status)

thirdmagic/task_def.py

@@ -0,0 +1,12 @@
+from typing import Optional
+
+from pydantic import BaseModel
+from rapyer import AtomicRedisModel
+from rapyer.fields import Key
+
+
+class MageflowTaskDefinition(AtomicRedisModel):
+    mageflow_task_name: Key[str]
+    task_name: str
+    input_validator: Optional[type[BaseModel]] = None
+    retries: Optional[int] = None

thirdmagic/typing_support.py

@@ -0,0 +1,15 @@
+# For python 310
+try:
+    from typing import Self
+except ImportError:
+    from typing_extensions import Self
+
+
+try:
+    # Python 3.12+
+    from typing import Unpack
+except ImportError:
+    # Older Python versions
+    from typing_extensions import Unpack
+
+__all__ = ["Self", "Unpack"]

thirdmagic/utils.py

@@ -0,0 +1,57 @@
+import dataclasses
+from typing import TypeVar, get_type_hints, Optional, Callable, Any
+
+from pydantic import BaseModel
+
+from thirdmagic.message import ReturnValueAnnotation, DEFAULT_RESULT_NAME
+
+PropType = TypeVar("PropType", bound=dataclasses.dataclass)
+
+
+def get_marked_fields(
+    model: type[BaseModel], mark_type: type[PropType]
+) -> list[tuple[PropType, str]]:
+    hints = get_type_hints(model, include_extras=True)
+    marked = []
+    for field_name, annotated_type in hints.items():
+        if hasattr(annotated_type, "__metadata__"):  # Annotated stores extras here
+            for meta in annotated_type.__metadata__:
+                if isinstance(meta, mark_type):
+                    marked.append((meta, field_name))
+    return marked
+
+
+def return_value_field(model_validators: type[BaseModel]) -> Optional[str]:
+    try:
+        marked_field = get_marked_fields(model_validators, ReturnValueAnnotation)
+        return_field_name = marked_field[0][1]
+    except (IndexError, TypeError):
+        return_field_name = None
+    return return_field_name or DEFAULT_RESULT_NAME
+
+
+def deep_merge(base: dict, updates: dict) -> dict:
+    results = base.copy()
+    for key, value in updates.items():
+        if key in base and isinstance(base[key], dict) and isinstance(value, dict):
+            results[key] = deep_merge(base[key], value)
+        else:
+            results[key] = value
+    return results
+
+
+# Which client is installed
+HatchetTaskType = Callable
+try:
+    HAS_HATCHET = True
+    from hatchet_sdk.runnables.workflow import BaseWorkflow
+
+    HatchetTaskType = HatchetTaskType | Callable
+except ImportError:
+    HAS_HATCHET = False
+
+# try:
+#     HAS_TEMPORAL = True
+#     HatchetTaskType = None
+# except ImportError:
+#     HAS_TEMPORAL = False

thirdmagic-0.0.1.dist-info/METADATA

@@ -0,0 +1,16 @@
+Metadata-Version: 2.4
+Name: thirdmagic
+Version: 0.0.1
+Summary: Core models and signatures for mageflow task orchestration
+Author-email: imaginary-cherry <yedidyakfir@gmail.com>
+License: MIT
+Requires-Python: <3.14,>=3.10
+Requires-Dist: pydantic<3.0.0,>=2.0.0
+Requires-Dist: rapyer<1.3.0,>=1.2.3
+Provides-Extra: dev
+Requires-Dist: black>=26.1.0; extra == 'dev'
+Requires-Dist: coverage[toml]<8.0.0,>=7.0.0; extra == 'dev'
+Requires-Dist: fakeredis[json,lua]<3.0.0,>=2.32.1; extra == 'dev'
+Requires-Dist: hatchet-sdk>=1.23.0; 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'

thirdmagic-0.0.1.dist-info/RECORD

@@ -0,0 +1,28 @@
+thirdmagic/__init__.py,sha256=XcEm5efUwT6XeX85pKBo6mEia0Gc4nw5i7ft7Py8TAM,165
+thirdmagic/consts.py,sha256=DvFDNOmi0LpHpi2IArv19NQvLcf7mTgnpqigLqw2P4c,95
+thirdmagic/container.py,sha256=738Yx5aG_1Z7puDRMKkOg3AYFStE7KagjQXvIdz1Pbc,819
+thirdmagic/errors.py,sha256=DfSAovxzKWiZ88eFga9CVD1zhkTOrLpIbeLIbDHFWBA,385
+thirdmagic/message.py,sha256=PS-kbiWeX-BtGxDyokGD-e5pDoAhw-NrdTBtgOLLrsU,637
+thirdmagic/task_def.py,sha256=Exm85y-FdIcKwehrZvHII3VkX-s-DgJLM4MLr1KitqI,316
+thirdmagic/typing_support.py,sha256=sPcXb_BPNOhYwDclrJMv4avgHIKR-F8_LDKSMJwx2VI,284
+thirdmagic/utils.py,sha256=HZaPmu6VWDUG6zhRKgCfNXIF9MpIFysr3777dKVlse0,1760
+thirdmagic/chain/__init__.py,sha256=mLiGXAJtW8-XD_lPmvgCIonr5EAe_Lops0c-HwqH1LM,140
+thirdmagic/chain/creator.py,sha256=8DobAR5qjyGrO3-8Vo-E3w7rFhQZ-euqsdWq7amvfg8,1202
+thirdmagic/chain/model.py,sha256=7RYTN-jJyutZG4CiyryuLd6O7b5VFcA2oWX_8YuMdQc,3486
+thirdmagic/clients/__init__.py,sha256=CQpt8cqfKp_pLzjpWSzwL5vWZy61cB0fJLaQxFtsrkY,133
+thirdmagic/clients/base.py,sha256=1IctGFIkhWeDkdQlGchT3Lw3jTTQe_FKn4vMsz7OppM,4500
+thirdmagic/clients/lifecycle.py,sha256=Kl-sHbXeg5b9jXRzRyqhhX7etZvihAa1kRzWGmUPmtk,494
+thirdmagic/signature/__init__.py,sha256=klbJzw6XbxHb4DWxwcLW01MvdDex4eo5yYxbNbeduu8,213
+thirdmagic/signature/model.py,sha256=K-Wi3bM_YOkTeVEhippcVGsUykybH-V2spREWvomjCw,6895
+thirdmagic/signature/status.py,sha256=rK4ipp9bBtfWUFBCHhkvISZpBT28e502y8NqudACWz8,840
+thirdmagic/swarm/__init__.py,sha256=G8KA4c0ciedzvV4JlRKQ5oA21oKp4J60Cywfw1AMLUU,232
+thirdmagic/swarm/consts.py,sha256=O-bKrVnzCxw4tETvArmSxevLAfZjs3Mv-h4lnvoq2OY,56
+thirdmagic/swarm/creator.py,sha256=QpggNG1R7UCXLq8SosVqU4FUD_z-qcxEfvemP40SOIo,1581
+thirdmagic/swarm/model.py,sha256=4NncuSTN9BBozEUOtEgViZSNg6-1EsTK1NrOnNVkxkU,9042
+thirdmagic/swarm/state.py,sha256=p2QjltycxQHBgaX7jDlG_10cJjZWlhKpKwgISKWvBY0,390
+thirdmagic/task/__init__.py,sha256=JX74AHbxN3J-IK1Gv1KqV6Z-65OR6LQ98hFzp3xPcKU,532
+thirdmagic/task/creator.py,sha256=d6pr0MJC3dUu3d7xv6pxEObO8QFYylSA8-iWW_XnYiQ,2843
+thirdmagic/task/model.py,sha256=yGehcmFhe8pNa0t6bXd3qcCdPiKlV23XL6Wk_Yond58,3058
+thirdmagic-0.0.1.dist-info/METADATA,sha256=bcT2ya9ug_8wjLUQHIeZk4qxMuCNzpKQTzW8mQXik70,649
+thirdmagic-0.0.1.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
+thirdmagic-0.0.1.dist-info/RECORD,,

thirdmagic-0.0.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.29.0
+Root-Is-Purelib: true
+Tag: py3-none-any