mageflow 0.0.1__tar.gz

mageflow-0.0.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 yedidyakfir
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

mageflow-0.0.1/PKG-INFO

@@ -0,0 +1,164 @@
+Metadata-Version: 2.4
+Name: mageflow
+Version: 0.0.1
+Summary: Manage Graph Execution Flow - A unified interface for task orchestration across different task managers
+License: MIT
+License-File: LICENSE
+Keywords: task-orchestration,workflow,task-manager,hatchet,async,redis,task-chain,task-swarm,distributed,microservices,pipeline,graph-execution,workflow-engine,task-queue,python
+Author: yedidyakfir
+Author-email: yedidyakfir@gmail.com
+Requires-Python: >=3.10,<3.14
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Framework :: AsyncIO
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Classifier: Topic :: System :: Distributed Computing
+Classifier: Topic :: Software Development :: Libraries :: Application Frameworks
+Classifier: Typing :: Typed
+Classifier: Operating System :: OS Independent
+Provides-Extra: display
+Provides-Extra: hatchet
+Requires-Dist: dash-bootstrap-components (>=2.0.4,<3.0.0) ; extra == "display"
+Requires-Dist: dash-cytoscape (>=1.0.2,<2.0.0) ; extra == "display"
+Requires-Dist: dash[async] (>=3.3.0,<4.0.0) ; extra == "display"
+Requires-Dist: hatchet-sdk (>=1.21.1,<2.0.0) ; extra == "hatchet"
+Requires-Dist: rapyer (>=1.1.0,<2.0.0)
+Project-URL: Bug Tracker, https://github.com/yedidyakfir/mageflow/issues
+Project-URL: Changelog, https://github.com/yedidyakfir/mageflow/releases
+Project-URL: Documentation, https://yedidyakfir.github.io/mageflow/
+Project-URL: Homepage, https://yedidyakfir.github.io/mageflow/
+Project-URL: Repository, https://github.com/yedidyakfir/mageflow
+Project-URL: Source Code, https://github.com/yedidyakfir/mageflow
+Description-Content-Type: text/markdown
+
+<div align="center">
+  <img src="logo.png" alt="MageFlow Logo" width="200"/>
+
+ [![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/)
+  [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+  [![codecov](https://codecov.io/gh/yedidyakfir/mageflow/branch/main/graph/badge.svg)](https://codecov.io/gh/yedidyakfir/mageflow)
+  [![PyPI version](https://badge.fury.io/py/mageflow.svg)](https://badge.fury.io/py/mageflow)
+  [![Downloads](https://static.pepy.tech/badge/mageflow)](https://pepy.tech/project/mageflow)
+  [![Documentation](https://img.shields.io/badge/docs-github.io-blue)](https://yedidyakfir.github.io/mageflow/)
+
+  
+  📚 **[Full Documentation](https://yedidyakfir.github.io/mageflow/)** | [Installation](https://yedidyakfir.github.io/mageflow/setup/) | [API Reference](https://yedidyakfir.github.io/mageflow/api/)
+
+</div>
+
+# MageFlow
+
+**Ma**nage **G**raph **E**xecution Flow - A unified interface for task orchestration across different task managers.
+
+## Why MageFlow?
+
+Instead of spreading workflow logic throughout your codebase, MageFlow centralizes task orchestration with a clean, unified API. Switch between task managers (Hatchet, Taskiq, etc.) without rewriting your orchestration code.
+
+## Key Features
+
+🔗 **Task Chaining** - Sequential workflows where tasks depend on previous completions  
+🐝 **Task Swarms** - Parallel execution with intelligent coordination  
+📞 **Callback System** - Robust success/error handling  
+🎯 **Task Signatures** - Flexible task definition with validation  
+⏯️ **Lifecycle Control** - Pause, resume, and monitor task execution  
+💾 **Persistent State** - Redis-backed state management with recovery  
+
+## Installation
+
+```bash
+pip install mageflow[hatchet]  # For Hatchet backend
+```
+
+## Quick Setup
+
+```python
+import asyncio
+import redis
+from hatchet_sdk import Hatchet, ClientConfig
+import mageflow
+
+# Configure backend and Redis
+config = ClientConfig(token="your-hatchet-token")
+redis_client = redis.asyncio.from_url("redis://localhost", decode_responses=True)
+hatchet_client = Hatchet(config=config)
+
+# Create MageFlow instance
+mf = mageflow.Mageflow(hatchet_client, redis_client=redis_client)
+```
+
+## Example Usage
+
+### Define Tasks
+
+```python
+from pydantic import BaseModel
+
+class ProcessData(BaseModel):
+    data: str
+
+@mf.task(name="process-data", input_validator=ProcessData)
+async def process_data(msg: ProcessData):
+    return {"processed": msg.data}
+
+@mf.task(name="send-notification") 
+async def send_notification(msg):
+    print(f"Notification sent: {msg}")
+    return {"status": "sent"}
+```
+
+### Chain Tasks
+
+```python
+# Sequential execution
+workflow = await mageflow.chain([
+    process_data_task,
+    send_notification_task
+], name="data-pipeline")
+```
+
+### Parallel Swarms
+
+```python
+# Parallel execution
+swarm = await mageflow.swarm([
+    process_user_task,
+    update_cache_task,
+    send_email_task
+], task_name="user-onboarding")
+```
+
+### Task Signatures with Callbacks
+
+```python
+task_signature = await mageflow.sign(
+    task_name="process-order",
+    task_identifiers={"order_id": "12345"},
+    success_callbacks=[send_confirmation_task],
+    error_callbacks=[handle_error_task]
+)
+```
+
+## Use Cases
+
+- **Data Pipelines** - ETL operations with error handling
+- **Microservice Coordination** - Orchestrate distributed service calls  
+- **Batch Processing** - Parallel processing of large datasets
+- **User Workflows** - Multi-step onboarding and registration
+- **Content Processing** - Media processing with multiple stages
+
+## Documentation
+
+- [Setup Guide](docs/setup.md)
+- [API Reference](docs/api/)
+- [Task Lifecycle](docs/documentation/task-lifecycle.md)
+- [Callbacks](docs/documentation/callbacks.md)
+
+## License
+
+MIT

mageflow-0.0.1/README.md

@@ -0,0 +1,125 @@
+<div align="center">
+  <img src="logo.png" alt="MageFlow Logo" width="200"/>
+
+ [![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/)
+  [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+  [![codecov](https://codecov.io/gh/yedidyakfir/mageflow/branch/main/graph/badge.svg)](https://codecov.io/gh/yedidyakfir/mageflow)
+  [![PyPI version](https://badge.fury.io/py/mageflow.svg)](https://badge.fury.io/py/mageflow)
+  [![Downloads](https://static.pepy.tech/badge/mageflow)](https://pepy.tech/project/mageflow)
+  [![Documentation](https://img.shields.io/badge/docs-github.io-blue)](https://yedidyakfir.github.io/mageflow/)
+
+  
+  📚 **[Full Documentation](https://yedidyakfir.github.io/mageflow/)** | [Installation](https://yedidyakfir.github.io/mageflow/setup/) | [API Reference](https://yedidyakfir.github.io/mageflow/api/)
+
+</div>
+
+# MageFlow
+
+**Ma**nage **G**raph **E**xecution Flow - A unified interface for task orchestration across different task managers.
+
+## Why MageFlow?
+
+Instead of spreading workflow logic throughout your codebase, MageFlow centralizes task orchestration with a clean, unified API. Switch between task managers (Hatchet, Taskiq, etc.) without rewriting your orchestration code.
+
+## Key Features
+
+🔗 **Task Chaining** - Sequential workflows where tasks depend on previous completions  
+🐝 **Task Swarms** - Parallel execution with intelligent coordination  
+📞 **Callback System** - Robust success/error handling  
+🎯 **Task Signatures** - Flexible task definition with validation  
+⏯️ **Lifecycle Control** - Pause, resume, and monitor task execution  
+💾 **Persistent State** - Redis-backed state management with recovery  
+
+## Installation
+
+```bash
+pip install mageflow[hatchet]  # For Hatchet backend
+```
+
+## Quick Setup
+
+```python
+import asyncio
+import redis
+from hatchet_sdk import Hatchet, ClientConfig
+import mageflow
+
+# Configure backend and Redis
+config = ClientConfig(token="your-hatchet-token")
+redis_client = redis.asyncio.from_url("redis://localhost", decode_responses=True)
+hatchet_client = Hatchet(config=config)
+
+# Create MageFlow instance
+mf = mageflow.Mageflow(hatchet_client, redis_client=redis_client)
+```
+
+## Example Usage
+
+### Define Tasks
+
+```python
+from pydantic import BaseModel
+
+class ProcessData(BaseModel):
+    data: str
+
+@mf.task(name="process-data", input_validator=ProcessData)
+async def process_data(msg: ProcessData):
+    return {"processed": msg.data}
+
+@mf.task(name="send-notification") 
+async def send_notification(msg):
+    print(f"Notification sent: {msg}")
+    return {"status": "sent"}
+```
+
+### Chain Tasks
+
+```python
+# Sequential execution
+workflow = await mageflow.chain([
+    process_data_task,
+    send_notification_task
+], name="data-pipeline")
+```
+
+### Parallel Swarms
+
+```python
+# Parallel execution
+swarm = await mageflow.swarm([
+    process_user_task,
+    update_cache_task,
+    send_email_task
+], task_name="user-onboarding")
+```
+
+### Task Signatures with Callbacks
+
+```python
+task_signature = await mageflow.sign(
+    task_name="process-order",
+    task_identifiers={"order_id": "12345"},
+    success_callbacks=[send_confirmation_task],
+    error_callbacks=[handle_error_task]
+)
+```
+
+## Use Cases
+
+- **Data Pipelines** - ETL operations with error handling
+- **Microservice Coordination** - Orchestrate distributed service calls  
+- **Batch Processing** - Parallel processing of large datasets
+- **User Workflows** - Multi-step onboarding and registration
+- **Content Processing** - Media processing with multiple stages
+
+## Documentation
+
+- [Setup Guide](docs/setup.md)
+- [API Reference](docs/api/)
+- [Task Lifecycle](docs/documentation/task-lifecycle.md)
+- [Callbacks](docs/documentation/callbacks.md)
+
+## License
+
+MIT

mageflow-0.0.1/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-0.0.1/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-0.0.1/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-0.0.1/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-0.0.1/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-0.0.1/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-0.0.1/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