mageflow 0.0.1__py3-none-any.whl

mageflow/__init__.py

@@ -0,0 +1,30 @@
+from mageflow.callbacks import register_task, handle_task_callback
+from mageflow.chain.creator import chain
+from mageflow.client import Mageflow
+from mageflow.init import init_mageflow_hatchet_tasks
+from mageflow.signature.creator import (
+    sign,
+    load_signature,
+    resume_task,
+    lock_task,
+    resume,
+    pause,
+)
+from mageflow.signature.status import TaskStatus
+from mageflow.swarm.creator import swarm
+
+
+__all__ = [
+    "load_signature",
+    "resume_task",
+    "lock_task",
+    "resume",
+    "pause",
+    "sign",
+    "init_mageflow_hatchet_tasks",
+    "register_task",
+    "handle_task_callback",
+    "Mageflow",
+    "chain",
+    "swarm",
+]

mageflow/callbacks.py

@@ -0,0 +1,72 @@
+import asyncio
+import functools
+import inspect
+from enum import Enum
+from typing import Any
+
+from hatchet_sdk import Context
+from hatchet_sdk.runnables.types import EmptyModel
+from pydantic import BaseModel
+
+from mageflow.invokers.hatchet import HatchetInvoker
+from mageflow.utils.pythonic import flexible_call
+
+
+class AcceptParams(Enum):
+    JUST_MESSAGE = 1
+    NO_CTX = 2
+    ALL = 3
+
+
+class HatchetResult(BaseModel):
+    hatchet_results: Any
+
+
+def handle_task_callback(
+    expected_params: AcceptParams = AcceptParams.NO_CTX, wrap_res: bool = True
+):
+    def task_decorator(func):
+        @functools.wraps(func)
+        async def wrapper(message: EmptyModel, ctx: Context, *args, **kwargs):
+            invoker = HatchetInvoker(message, ctx)
+            if not await invoker.should_run_task():
+                await ctx.aio_cancel()
+                await asyncio.sleep(10)
+                # NOTE: This should not run, the task should cancel, but just in case
+                return {"Error": "Task should have been canceled"}
+            try:
+                await invoker.start_task()
+                if expected_params == AcceptParams.JUST_MESSAGE:
+                    result = await flexible_call(func, message)
+                elif expected_params == AcceptParams.NO_CTX:
+                    result = await flexible_call(func, message, *args, **kwargs)
+                else:
+                    result = await flexible_call(func, message, ctx, *args, **kwargs)
+            except (Exception, asyncio.CancelledError) as e:
+                await invoker.run_error()
+                await invoker.remove_task(with_error=False)
+                raise
+            else:
+                task_results = HatchetResult(hatchet_results=result)
+                dumped_results = task_results.model_dump(mode="json")
+                await invoker.run_success(dumped_results["hatchet_results"])
+                await invoker.remove_task(with_success=False)
+                if wrap_res:
+                    return task_results
+                else:
+                    return result
+
+        wrapper.__signature__ = inspect.signature(func)
+        return wrapper
+
+    return task_decorator
+
+
+def register_task(register_name: str):
+    from mageflow.startup import REGISTERED_TASKS
+
+    def decorator(func):
+        REGISTERED_TASKS.append((func, register_name))
+        return func
+
+    return decorator

mageflow/chain/consts.py

@@ -0,0 +1,8 @@
+# Params
+from mageflow.signature.consts import MAGEFLOW_TASK_INITIALS
+
+CHAIN_TASK_ID_NAME = "chain_task_id"
+
+# Task names
+ON_CHAIN_ERROR = f"{MAGEFLOW_TASK_INITIALS}on_chain_error"
+ON_CHAIN_END = f"{MAGEFLOW_TASK_INITIALS}on_chain_done"

mageflow/chain/creator.py

@@ -0,0 +1,73 @@
+import asyncio
+
+from mageflow.chain.consts import ON_CHAIN_END, ON_CHAIN_ERROR
+from mageflow.chain.messages import ChainSuccessTaskCommandMessage
+from mageflow.chain.model import ChainTaskSignature
+from mageflow.signature.creator import (
+    TaskSignatureConvertible,
+    resolve_signature_key,
+)
+from mageflow.signature.model import (
+    TaskIdentifierType,
+    TaskSignature,
+    TaskInputType,
+)
+
+
+async def chain(
+    tasks: list[TaskSignatureConvertible],
+    name: str = None,
+    error: TaskInputType = None,
+    success: TaskInputType = None,
+) -> ChainTaskSignature:
+    tasks = [await resolve_signature_key(task) for task in 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,
+    )
+    await chain_task_signature.save()
+
+    callback_kwargs = dict(chain_task_id=chain_task_signature.key)
+    on_chain_error = TaskSignature(
+        task_name=ON_CHAIN_ERROR,
+        task_identifiers=callback_kwargs,
+        model_validators=ChainSuccessTaskCommandMessage,
+    )
+    on_chain_success = TaskSignature(
+        task_name=ON_CHAIN_END,
+        task_identifiers=callback_kwargs,
+        model_validators=ChainSuccessTaskCommandMessage,
+    )
+    await _chain_task_to_previous_success(tasks, on_chain_error, on_chain_success)
+    return chain_task_signature
+
+
+async def _chain_task_to_previous_success(
+    tasks: list[TaskSignature], error: TaskSignature, success: TaskSignature
+) -> TaskIdentifierType:
+    """
+    Take a list of tasks and connect each one to the previous one.
+    """
+    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."
+        )
+
+    total_tasks = tasks + [success]
+    error_tasks = await error.duplicate_many(len(tasks))
+    store_errors = [error.save() for error in error_tasks]
+
+    # Store tasks
+    await asyncio.gather(success.save(), *store_errors)
+    update_tasks = [
+        task.add_callbacks(success=[total_tasks[i + 1]], errors=[error_tasks[i]])
+        for i, task in enumerate(tasks)
+    ]
+    chained_tasks = await asyncio.gather(*update_tasks)
+    return chained_tasks[0]

mageflow/chain/messages.py

@@ -0,0 +1,9 @@
+from typing import Any, Annotated
+
+from pydantic import BaseModel
+
+from mageflow.models.message import ReturnValue
+
+
+class ChainSuccessTaskCommandMessage(BaseModel):
+    chain_results: Annotated[Any, ReturnValue()]

mageflow/chain/model.py

@@ -0,0 +1,61 @@
+import asyncio
+
+from pydantic import field_validator, Field
+
+from mageflow.errors import MissingSignatureError
+from mageflow.signature.model import TaskSignature, TaskIdentifierType
+from mageflow.signature.status import SignatureStatus
+
+
+class ChainTaskSignature(TaskSignature):
+    tasks: list[TaskIdentifierType] = 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]
+
+    async def workflow(self, **task_additional_params):
+        first_task = await TaskSignature.get_safe(self.tasks[0])
+        if first_task is None:
+            raise MissingSignatureError(f"First task from chain {self.key} not found")
+        return await first_task.workflow(**task_additional_params)
+
+    async def delete_chain_tasks(self, with_errors=True, with_success=True):
+        signatures = await asyncio.gather(
+            *[TaskSignature.get_safe(signature_id) for signature_id in self.tasks],
+            return_exceptions=True,
+        )
+        signatures = [sign for sign in signatures if isinstance(sign, TaskSignature)]
+        delete_tasks = [
+            signature.remove(with_errors, with_success) for signature in signatures
+        ]
+        await asyncio.gather(*delete_tasks)
+
+    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)

mageflow/chain/workflows.py

@@ -0,0 +1,65 @@
+import asyncio
+
+from hatchet_sdk import Context
+from hatchet_sdk.runnables.types import EmptyModel
+
+from mageflow.chain.consts import CHAIN_TASK_ID_NAME
+from mageflow.chain.messages import ChainSuccessTaskCommandMessage
+from mageflow.chain.model import ChainTaskSignature
+from mageflow.invokers.hatchet import HatchetInvoker
+from mageflow.signature.consts import TASK_ID_PARAM_NAME
+from mageflow.signature.model import TaskSignature
+
+
+async def chain_end_task(msg: ChainSuccessTaskCommandMessage, ctx: Context) -> None:
+    try:
+        task_data = HatchetInvoker(msg, ctx).task_ctx
+        chain_task_id = task_data[CHAIN_TASK_ID_NAME]
+        current_task_id = task_data[TASK_ID_PARAM_NAME]
+
+        chain_task_signature, current_task = await asyncio.gather(
+            ChainTaskSignature.get_safe(chain_task_id),
+            TaskSignature.get_safe(current_task_id),
+        )
+        ctx.log(f"Chain task done {chain_task_signature.task_name}")
+
+        # Calling error callback from a chain task - This is done before deletion because a deletion error should not disturb the workflow
+        await chain_task_signature.activate_success(msg.chain_results)
+        ctx.log(f"Chain task success {chain_task_signature.task_name}")
+
+        # Remove tasks
+        await asyncio.gather(
+            chain_task_signature.remove(with_success=False), current_task.remove()
+        )
+    except Exception as e:
+        ctx.log(f"MAJOR - infrastructure error in chain end task: {e}")
+        raise
+
+
+# This task needs to be added as a workflow
+async def chain_error_task(msg: EmptyModel, ctx: Context) -> None:
+    try:
+        task_data = HatchetInvoker(msg, ctx).task_ctx
+        chain_task_id = task_data[CHAIN_TASK_ID_NAME]
+        current_task_id = task_data[TASK_ID_PARAM_NAME]
+        chain_packed_task, current_task = await asyncio.gather(
+            ChainTaskSignature.get_safe(chain_task_id),
+            TaskSignature.get_safe(current_task_id),
+        )
+        ctx.log(
+            f"Chain task failed {chain_packed_task.task_name} on task id - {current_task_id}"
+        )
+
+        # Calling error callback from chain task
+        await chain_packed_task.activate_error(msg)
+        ctx.log(f"Chain task error {chain_packed_task.task_name}")
+
+        # Remove tasks
+        await chain_packed_task.delete_chain_tasks()
+        await asyncio.gather(
+            chain_packed_task.remove(with_error=False), current_task.remove()
+        )
+        ctx.log(f"Clean redis from chain tasks {chain_packed_task.task_name}")
+    except Exception as e:
+        ctx.log(f"MAJOR - infrastructure error in chain error task: {e}")
+        raise

mageflow/client.py

@@ -0,0 +1,140 @@
+import functools
+import os
+from typing import TypeVar, Any, overload, Unpack
+
+import redis
+from hatchet_sdk import Hatchet, Worker
+from hatchet_sdk.runnables.workflow import BaseWorkflow
+from hatchet_sdk.worker.worker import LifespanFn
+from redis.asyncio import Redis
+
+from mageflow.callbacks import AcceptParams, register_task, handle_task_callback
+from mageflow.chain.creator import chain
+from mageflow.init import init_mageflow_hatchet_tasks
+from mageflow.signature.creator import sign, TaskSignatureConvertible
+from mageflow.signature.model import TaskSignature, TaskInputType
+from mageflow.signature.types import HatchetTaskType
+from mageflow.startup import (
+    lifespan_initialize,
+    mageflow_config,
+    init_mageflow,
+    teardown_mageflow,
+)
+from mageflow.swarm.creator import swarm, SignatureOptions
+
+
+async def merge_lifespan(original_lifespan: LifespanFn):
+    await init_mageflow()
+    async for res in original_lifespan():
+        yield res
+    await teardown_mageflow()
+
+
+class HatchetMageflow(Hatchet):
+    def __init__(
+        self,
+        hatchet: Hatchet,
+        redis_client: Redis,
+        param_config: AcceptParams = AcceptParams.NO_CTX,
+    ):
+        super().__init__(client=hatchet._client)
+        self.hatchet = hatchet
+        self.redis = redis_client
+        self.param_config = param_config
+
+    def task(self, *, name: str | None = None, **kwargs):
+        hatchet_task = super().task(name=name, **kwargs)
+
+        def decorator(func):
+            handler_dec = handle_task_callback(self.param_config)
+            func = handler_dec(func)
+            wf = hatchet_task(func)
+
+            nonlocal name
+            task_name = name or func.__name__
+            register = register_task(task_name)
+            return register(wf)
+
+        return decorator
+
+    def durable_task(self, *, name: str | None = None, **kwargs):
+        hatchet_task = super().durable_task(name=name, **kwargs)
+
+        def decorator(func):
+            handler_dec = handle_task_callback(self.param_config)
+            func = handler_dec(func)
+            wf = hatchet_task(func)
+            nonlocal name
+            task_name = name or func.__name__
+            register = register_task(task_name)
+            return register(wf)
+
+        return decorator
+
+    def worker(
+        self,
+        *args,
+        workflows: list[BaseWorkflow[Any]] | None = None,
+        lifespan: LifespanFn | None = None,
+        **kwargs,
+    ) -> Worker:
+        mageflow_flows = init_mageflow_hatchet_tasks(self.hatchet)
+        workflows += mageflow_flows
+        if lifespan is None:
+            lifespan = lifespan_initialize
+        else:
+            lifespan = functools.partial(merge_lifespan, lifespan)
+
+        return super().worker(*args, workflows=workflows, lifespan=lifespan, **kwargs)
+
+    async def sign(self, task: str | HatchetTaskType, **options: Any) -> TaskSignature:
+        return await sign(task, **options)
+
+    async def chain(
+        self,
+        tasks: list[TaskSignatureConvertible],
+        name: str = None,
+        error: TaskInputType = None,
+        success: TaskInputType = None,
+    ):
+        return await chain(tasks, name, error, success)
+
+    async def swarm(
+        self,
+        tasks: list[TaskSignatureConvertible] = None,
+        task_name: str = None,
+        **kwargs: Unpack[SignatureOptions],
+    ):
+        return await swarm(tasks, task_name, **kwargs)
+
+
+T = TypeVar("T")
+
+
+@overload
+def Mageflow(
+    hatchet_client: Hatchet, redis_client: Redis | str = None
+) -> HatchetMageflow: ...
+
+
+def Mageflow(
+    hatchet_client: T = None,
+    redis_client: Redis | str = None,
+    param_config: AcceptParams = AcceptParams.NO_CTX,
+) -> T:
+    if hatchet_client is None:
+        hatchet_client = Hatchet()
+
+    # Create a hatchet client with empty namespace for creating wf
+    config = hatchet_client._client.config.model_copy(deep=True)
+    config.namespace = ""
+    hatchet_caller = Hatchet(config=config, debug=hatchet_client._client.debug)
+    mageflow_config.hatchet_client = hatchet_caller
+
+    if redis_client is None:
+        redis_url = os.getenv("REDIS_URL")
+        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)
+    mageflow_config.redis_client = redis_client
+    return HatchetMageflow(hatchet_client, redis_client, param_config)

mageflow/errors.py

@@ -0,0 +1,22 @@
+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

mageflow/init.py

@@ -0,0 +1,53 @@
+from hatchet_sdk import Hatchet
+
+from mageflow.callbacks import register_task
+from mageflow.chain.consts import ON_CHAIN_END, ON_CHAIN_ERROR
+from mageflow.chain.messages import ChainSuccessTaskCommandMessage
+from mageflow.chain.workflows import chain_end_task, chain_error_task
+from mageflow.swarm.consts import ON_SWARM_ERROR, ON_SWARM_END, ON_SWARM_START
+from mageflow.swarm.messages import SwarmResultsMessage
+from mageflow.swarm.workflows import (
+    swarm_item_failed,
+    swarm_item_done,
+    swarm_start_tasks,
+)
+
+
+def init_mageflow_hatchet_tasks(hatchet: Hatchet):
+    # Chain tasks
+    hatchet_chain_done = hatchet.task(
+        name=ON_CHAIN_END,
+        input_validator=ChainSuccessTaskCommandMessage,
+    )
+    hatchet_chain_error = hatchet.task(name=ON_CHAIN_ERROR)
+    chain_done_task = hatchet_chain_done(chain_end_task)
+    on_chain_error_task = hatchet_chain_error(chain_error_task)
+    register_chain_done = register_task(ON_CHAIN_END)
+    register_chain_error = register_task(ON_CHAIN_ERROR)
+    chain_done_task = register_chain_done(chain_done_task)
+    on_chain_error_task = register_chain_error(on_chain_error_task)
+
+    # Swarm tasks
+    swarm_start = hatchet.task(name=ON_SWARM_START)
+    swarm_done = hatchet.task(
+        name=ON_SWARM_END,
+        input_validator=SwarmResultsMessage,
+    )
+    swarm_error = hatchet.task(name=ON_SWARM_ERROR)
+    swarm_start = swarm_start(swarm_start_tasks)
+    swarm_done = swarm_done(swarm_item_done)
+    swarm_error = swarm_error(swarm_item_failed)
+    register_swarm_start = register_task(ON_SWARM_START)
+    register_swarm_done = register_task(ON_SWARM_END)
+    register_swarm_error = register_task(ON_SWARM_ERROR)
+    swarm_start = register_swarm_start(swarm_start)
+    swarm_done = register_swarm_done(swarm_done)
+    swarm_error = register_swarm_error(swarm_error)
+
+    return [
+        on_chain_error_task,
+        chain_done_task,
+        swarm_start,
+        swarm_done,
+        swarm_error,
+    ]

mageflow/invokers/base.py

@@ -0,0 +1,34 @@
+import abc
+from abc import ABC
+from typing import Any
+
+from mageflow.signature.model import TaskSignature
+
+
+class BaseInvoker(ABC):
+    @property
+    @abc.abstractmethod
+    def task_ctx(self) -> dict:
+        pass
+
+    @abc.abstractmethod
+    async def start_task(self):
+        pass
+
+    @abc.abstractmethod
+    async def run_success(self, result: Any) -> bool:
+        pass
+
+    @abc.abstractmethod
+    async def run_error(self) -> bool:
+        pass
+
+    @abc.abstractmethod
+    async def remove_task(
+        self, with_success: bool = True, with_error: bool = True
+    ) -> TaskSignature | None:
+        pass
+
+    @abc.abstractmethod
+    async def should_run_task(self) -> bool:
+        pass

mageflow/invokers/hatchet.py

@@ -0,0 +1,82 @@
+import asyncio
+from typing import Any
+
+from hatchet_sdk import Context
+from hatchet_sdk.runnables.contextvars import ctx_additional_metadata
+from pydantic import BaseModel
+
+from mageflow.invokers.base import BaseInvoker
+from mageflow.signature.consts import TASK_ID_PARAM_NAME
+from mageflow.signature.model import TaskSignature
+from mageflow.signature.status import SignatureStatus
+from mageflow.workflows import TASK_DATA_PARAM_NAME
+
+
+class HatchetInvoker(BaseInvoker):
+    def __init__(self, message: BaseModel, ctx: Context):
+        self.message = message
+        self.task_data = ctx.additional_metadata.get(TASK_DATA_PARAM_NAME, {})
+        self.workflow_id = ctx.workflow_id
+        hatchet_ctx_metadata = ctx_additional_metadata.get() or {}
+        hatchet_ctx_metadata.pop(TASK_DATA_PARAM_NAME, None)
+        ctx_additional_metadata.set(hatchet_ctx_metadata)
+
+    @property
+    def task_ctx(self) -> dict:
+        return self.task_data
+
+    async def start_task(self):
+        task_id = self.task_data.get(TASK_ID_PARAM_NAME, None)
+        if task_id:
+            async with TaskSignature.lock_from_key(task_id) as signature:
+                await signature.change_status(SignatureStatus.ACTIVE)
+                await signature.task_status.aupdate(worker_task_id=self.workflow_id)
+
+    async def run_success(self, result: Any) -> bool:
+        success_publish_tasks = []
+        task_id = self.task_data.get(TASK_ID_PARAM_NAME, None)
+        if task_id:
+            current_task = await TaskSignature.get_safe(task_id)
+            task_success_workflows = current_task.activate_success(result)
+            success_publish_tasks.append(asyncio.create_task(task_success_workflows))
+
+        if success_publish_tasks:
+            await asyncio.gather(*success_publish_tasks)
+            return True
+        return False
+
+    async def run_error(self) -> bool:
+        error_publish_tasks = []
+        task_id = self.task_data.get(TASK_ID_PARAM_NAME, None)
+        if task_id:
+            current_task = await TaskSignature.get_safe(task_id)
+            task_error_workflows = current_task.activate_error(self.message)
+            error_publish_tasks.append(asyncio.create_task(task_error_workflows))
+
+        if error_publish_tasks:
+            await asyncio.gather(*error_publish_tasks)
+            return True
+        return False
+
+    async def remove_task(
+        self, with_success: bool = True, with_error: bool = True
+    ) -> TaskSignature | None:
+        task_id = self.task_data.get(TASK_ID_PARAM_NAME, None)
+        if task_id:
+            signature = await TaskSignature.get_safe(task_id)
+            if signature:
+                await signature.remove(with_error, with_success)
+
+    async def should_run_task(self) -> bool:
+        task_id = self.task_data.get(TASK_ID_PARAM_NAME, None)
+        if task_id:
+            signature = await TaskSignature.get_safe(task_id)
+            if signature is None:
+                return False
+            should_task_run = await signature.should_run()
+            if should_task_run:
+                return True
+            await signature.task_status.aupdate(last_status=SignatureStatus.ACTIVE)
+            await signature.handle_inactive_task(self.message)
+            return False
+        return True

mageflow/models/message.py

@@ -0,0 +1,6 @@
+import dataclasses
+
+
+@dataclasses.dataclass
+class ReturnValue:
+    pass

mageflow/signature/consts.py

@@ -0,0 +1,3 @@
+TASK_ID_PARAM_NAME = "task_id"
+
+MAGEFLOW_TASK_INITIALS = "mageflow_"