cellaflow 0.1.0__tar.gz

cellaflow-0.1.0/PKG-INFO

@@ -0,0 +1,58 @@
+Metadata-Version: 2.4
+Name: cellaflow
+Version: 0.1.0
+Summary: Python SDK for Cellaflow
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: grpcio>=1.60.0
+Requires-Dist: protobuf>=4.25.0
+Requires-Dist: msgpack>=1.0.7
+Requires-Dist: rfc8785>=0.1.4
+Provides-Extra: dev
+Requires-Dist: pytest>=8.0.0; extra == "dev"
+Requires-Dist: pytest-asyncio>=0.23.0; extra == "dev"
+Requires-Dist: black>=24.0.0; extra == "dev"
+Requires-Dist: flake8>=7.0.0; extra == "dev"
+Requires-Dist: mypy>=1.8.0; extra == "dev"
+Requires-Dist: grpcio-tools>=1.60.0; extra == "dev"
+Requires-Dist: mypy-protobuf>=3.5.0; extra == "dev"
+
+# Cellaflow Python SDK
+
+The official Python SDK for the Cellaflow Engine.
+## Development Setup
+
+To get up and running with the Python SDK for development, follow these steps:
+
+1. **Create and activate a virtual environment**:
+   ```bash
+   cd python
+   python3 -m venv venv
+   source venv/bin/activate
+   ```
+
+2. **Install the SDK in editable mode with development dependencies**:
+   ```bash
+   pip install -e ".[dev]"
+   ```
+
+## Generating Protobufs and Typing Stubs (`.pyi`)
+
+We use `grpcio-tools` and `mypy-protobuf` to generate Python code and typing stubs from the `.proto` files. This ensures your IDE (e.g. VS Code with Pylance or mypy) will properly resolve gRPC stubs and message attributes.
+
+To regenerate the protobuf definitions and their corresponding `.pyi` typing stubs, run the generation script from the `python` directory:
+
+```bash
+./scripts/generate_protos.sh
+```
+
+This will automatically compile the protos from `../proto/cellaflow/v1/*.proto` and output both the Python code (`*_pb2.py`, `*_pb2_grpc.py`) and typing stubs (`*.pyi`) into `src/cellaflow/v1/`.
+
+## Running Tests and Linters
+
+You can validate your changes using the following tools:
+
+- **Tests**: `pytest`
+- **Formatting**: `black src tests`
+- **Linting**: `flake8 src tests`
+- **Type Checking**: `mypy src tests`

cellaflow-0.1.0/README.md

@@ -0,0 +1,39 @@
+# Cellaflow Python SDK
+
+The official Python SDK for the Cellaflow Engine.
+## Development Setup
+
+To get up and running with the Python SDK for development, follow these steps:
+
+1. **Create and activate a virtual environment**:
+   ```bash
+   cd python
+   python3 -m venv venv
+   source venv/bin/activate
+   ```
+
+2. **Install the SDK in editable mode with development dependencies**:
+   ```bash
+   pip install -e ".[dev]"
+   ```
+
+## Generating Protobufs and Typing Stubs (`.pyi`)
+
+We use `grpcio-tools` and `mypy-protobuf` to generate Python code and typing stubs from the `.proto` files. This ensures your IDE (e.g. VS Code with Pylance or mypy) will properly resolve gRPC stubs and message attributes.
+
+To regenerate the protobuf definitions and their corresponding `.pyi` typing stubs, run the generation script from the `python` directory:
+
+```bash
+./scripts/generate_protos.sh
+```
+
+This will automatically compile the protos from `../proto/cellaflow/v1/*.proto` and output both the Python code (`*_pb2.py`, `*_pb2_grpc.py`) and typing stubs (`*.pyi`) into `src/cellaflow/v1/`.
+
+## Running Tests and Linters
+
+You can validate your changes using the following tools:
+
+- **Tests**: `pytest`
+- **Formatting**: `black src tests`
+- **Linting**: `flake8 src tests`
+- **Type Checking**: `mypy src tests`

cellaflow-0.1.0/pyproject.toml

@@ -0,0 +1,50 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "cellaflow"
+version = "0.1.0"
+description = "Python SDK for Cellaflow"
+readme = "README.md"
+requires-python = ">=3.9"
+dependencies = [
+    "grpcio>=1.60.0",
+    "protobuf>=4.25.0",
+    "msgpack>=1.0.7",
+    "rfc8785>=0.1.4",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=8.0.0",
+    "pytest-asyncio>=0.23.0",
+    "black>=24.0.0",
+    "flake8>=7.0.0",
+    "mypy>=1.8.0",
+    "grpcio-tools>=1.60.0",
+    "mypy-protobuf>=3.5.0",
+]
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.pytest.ini_options]
+addopts = "-ra -q"
+testpaths = [
+    "tests",
+]
+
+[tool.black]
+line-length = 88
+force-exclude = ".*_pb2.*\\.py"
+
+[tool.mypy]
+python_version = "3.10"
+strict = true
+ignore_missing_imports = true
+exclude = [
+    ".*_pb2\\.py$",
+    ".*_pb2_grpc\\.py$",
+    ".*\\.pyi$"
+]

cellaflow-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

cellaflow-0.1.0/src/cellaflow/__init__.py

@@ -0,0 +1,5 @@
+from cellaflow.client import CellaflowClient
+from cellaflow.decorators import workflow, step, tool
+from cellaflow.idempotency import IdempotencyScope
+
+__all__ = ["CellaflowClient", "workflow", "step", "tool", "IdempotencyScope"]

cellaflow-0.1.0/src/cellaflow/client.py

@@ -0,0 +1,165 @@
+import grpc
+from typing import Any, Dict, List, Optional, Tuple, cast
+
+from cellaflow.v1 import service_pb2, service_pb2_grpc
+from cellaflow.v1 import common_pb2, idempotency_pb2
+from cellaflow.serialization import serialize, deserialize
+
+
+class CellaflowClient:
+    """
+    gRPC Client for the Cellaflow Engine.
+    Handles communication with the engine and strictly uses
+    MessagePack for state payloads.
+    """
+
+    def __init__(self, target: str = "localhost:50051", secure: bool = False) -> None:
+        if secure:
+            self.channel = grpc.secure_channel(target, grpc.ssl_channel_credentials())
+        else:
+            self.channel = grpc.insecure_channel(target)
+        self.stub = service_pb2_grpc.WorkflowEngineServiceStub(self.channel)
+
+    def start_session(
+        self, workflow_id: str, version: str, session_id: Optional[str] = None
+    ) -> service_pb2.StartSessionResponse:
+        req = service_pb2.StartSessionRequest(
+            workflow_id=workflow_id,
+            version=version,
+        )
+        if session_id:
+            req.session_id = session_id
+
+        return cast(service_pb2.StartSessionResponse, self.stub.StartSession(req))
+
+    def commit_step(
+        self,
+        session_id: str,
+        sequence: int,
+        name: str,
+        status: "common_pb2.StepStatus.ValueType",
+        output_payload: Dict[str, Any],
+        idempotency_key: Optional[str] = None,
+        idempotency_fencing_token: Optional[int] = None,
+    ) -> service_pb2.CommitStepResponse:
+        # Strictly serialize dict to MessagePack
+        serialized_state = serialize(output_payload)
+
+        step_result = common_pb2.StepResult(
+            sequence=sequence,
+            name=name,
+            status=status,
+            output_payload=serialized_state,
+        )
+
+        req = service_pb2.CommitStepRequest(
+            session_id=session_id,
+            step_result=step_result,
+        )
+        if idempotency_key is not None:
+            req.idempotency_key = idempotency_key
+            if idempotency_fencing_token is None:
+                raise ValueError(
+                    "idempotency_fencing_token required if idempotency_key is set"
+                )
+            req.idempotency_fencing_token = idempotency_fencing_token
+
+        return cast(service_pb2.CommitStepResponse, self.stub.CommitStep(req))
+
+    def get_graph(
+        self, session_id: str, limit: Optional[int] = None, cursor: Optional[str] = None
+    ) -> Tuple[List[Dict[str, Any]], Optional[str]]:
+        """
+        Returns a tuple of (list of deserialized step results, next_cursor).
+        Each step result is a dictionary representation of the StepResult proto,
+        with the output_payload fully deserialized into a Python dictionary.
+        """
+        req = service_pb2.GetGraphRequest(session_id=session_id)
+        if limit is not None:
+            req.limit = limit
+        if cursor is not None:
+            req.cursor = cursor
+
+        resp: service_pb2.GetGraphResponse = self.stub.GetGraph(req)
+
+        results = []
+        for step in resp.steps:
+            results.append(
+                {
+                    "sequence": step.sequence,
+                    "name": step.name,
+                    "status": step.status,
+                    "output_payload": deserialize(step.output_payload),
+                    "idempotency_key": (
+                        step.idempotency_key
+                        if step.HasField("idempotency_key")
+                        else None
+                    ),
+                }
+            )
+
+        next_cursor = resp.next_cursor if resp.next_cursor else None
+        return results, next_cursor
+
+    def check_idempotency_cache(
+        self,
+        agent_id: str,
+        idempotency_key: str,
+        wait_timeout_ms: Optional[int] = None,
+        lease_ttl_ms: Optional[int] = None,
+    ) -> idempotency_pb2.CheckCacheResponse:
+        req = idempotency_pb2.CheckCacheRequest(
+            agent_id=agent_id,
+            idempotency_key=idempotency_key,
+        )
+        if wait_timeout_ms is not None:
+            req.wait_timeout_ms = wait_timeout_ms
+        if lease_ttl_ms is not None:
+            req.lease_ttl_ms = lease_ttl_ms
+
+        return cast(
+            idempotency_pb2.CheckCacheResponse,
+            self.stub.CheckIdempotencyCache(req),
+        )
+
+    def renew_lease(
+        self,
+        agent_id: str,
+        idempotency_key: str,
+        fencing_token: int,
+        extend_ms: int,
+    ) -> idempotency_pb2.RenewLeaseResponse:
+        req = idempotency_pb2.RenewLeaseRequest(
+            agent_id=agent_id,
+            idempotency_key=idempotency_key,
+            fencing_token=fencing_token,
+            extend_ms=extend_ms,
+        )
+
+        return cast(
+            idempotency_pb2.RenewLeaseResponse,
+            self.stub.RenewLease(req),
+        )
+
+    def release_lease(
+        self,
+        agent_id: str,
+        idempotency_key: str,
+        fencing_token: int,
+        reason: Optional[str] = None,
+    ) -> idempotency_pb2.ReleaseLeaseResponse:
+        req = idempotency_pb2.ReleaseLeaseRequest(
+            agent_id=agent_id,
+            idempotency_key=idempotency_key,
+            fencing_token=fencing_token,
+        )
+        if reason:
+            req.reason = reason
+
+        return cast(
+            idempotency_pb2.ReleaseLeaseResponse,
+            self.stub.ReleaseLease(req),
+        )
+
+    def close(self) -> None:
+        self.channel.close()

cellaflow-0.1.0/src/cellaflow/context.py

@@ -0,0 +1,38 @@
+from dataclasses import dataclass, field
+from typing import Any, Dict
+import contextvars
+
+from cellaflow.client import CellaflowClient
+
+
+@dataclass
+class WorkflowContext:
+    client: CellaflowClient
+    session_id: str
+    workflow_version: str
+    sequence: int = 0
+    # Map sequence number to a deserialized step payload for fast replay lookups
+    replayed_steps: Dict[int, Dict[str, Any]] = field(default_factory=dict)
+
+
+_current_context: contextvars.ContextVar[WorkflowContext] = contextvars.ContextVar(
+    "workflow_context"
+)
+
+
+def get_context() -> WorkflowContext:
+    try:
+        return _current_context.get()
+    except LookupError:
+        raise RuntimeError(
+            "No active workflow context found. "
+            "Are you calling a @step inside a @workflow?"
+        )
+
+
+def set_context(context: WorkflowContext) -> contextvars.Token[WorkflowContext]:
+    return _current_context.set(context)
+
+
+def reset_context(token: contextvars.Token[WorkflowContext]) -> None:
+    _current_context.reset(token)

cellaflow-0.1.0/src/cellaflow/decorators.py

@@ -0,0 +1,304 @@
+import functools
+import inspect
+import uuid
+import asyncio
+import time
+import logging
+from typing import Any, Callable, TypeVar, cast, Optional
+
+from cellaflow.client import CellaflowClient
+from cellaflow.context import WorkflowContext, set_context, reset_context, get_context
+from cellaflow.idempotency import derive_idempotency_key, IdempotencyScope
+from cellaflow.lease import LeaseHeartbeat
+from cellaflow.serialization import deserialize
+from cellaflow.v1.common_pb2 import STEP_STATUS_SUCCESS
+from cellaflow.v1.idempotency_pb2 import (
+    CACHE_STATUS_HIT,
+    CACHE_STATUS_ACQUIRED,
+    CACHE_STATUS_IN_PROGRESS,
+)
+
+logger = logging.getLogger(__name__)
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def workflow(
+    version: str = "1.0.0", target: str = "localhost:50051", secure: bool = False
+) -> Callable[[F], F]:
+    def decorator(func: F) -> F:
+        @functools.wraps(func)
+        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+            # Pop _session_id if provided for recovery, else generate new
+            session_id = kwargs.pop("_session_id", str(uuid.uuid4()))
+
+            client = CellaflowClient(target=target, secure=secure)
+            resp = client.start_session(
+                workflow_id=func.__name__, version=version, session_id=session_id
+            )
+
+            replayed_steps = {}
+            if resp.is_recovered:
+                steps, _ = client.get_graph(resp.session_id)
+                for step_info in steps:
+                    replayed_steps[step_info["sequence"]] = step_info["output_payload"]
+
+            ctx = WorkflowContext(
+                client=client,
+                session_id=resp.session_id,
+                workflow_version=resp.version,
+                sequence=0,
+                replayed_steps=replayed_steps,
+            )
+            token = set_context(ctx)
+            try:
+                return await func(*args, **kwargs)
+            finally:
+                reset_context(token)
+
+        @functools.wraps(func)
+        def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+            session_id = kwargs.pop("_session_id", str(uuid.uuid4()))
+
+            client = CellaflowClient(target=target, secure=secure)
+            resp = client.start_session(
+                workflow_id=func.__name__, version=version, session_id=session_id
+            )
+
+            replayed_steps = {}
+            if resp.is_recovered:
+                steps, _ = client.get_graph(resp.session_id)
+                for step_info in steps:
+                    replayed_steps[step_info["sequence"]] = step_info["output_payload"]
+
+            ctx = WorkflowContext(
+                client=client,
+                session_id=resp.session_id,
+                workflow_version=resp.version,
+                sequence=0,
+                replayed_steps=replayed_steps,
+            )
+            token = set_context(ctx)
+            try:
+                return func(*args, **kwargs)
+            finally:
+                reset_context(token)
+
+        if inspect.iscoroutinefunction(func):
+            return cast(F, async_wrapper)
+        return cast(F, sync_wrapper)
+
+    return decorator
+
+
+def step(
+    func: Optional[F] = None,
+    *,
+    idempotency_key: Optional[str] = None,
+    agent_id: str = "default",
+    tool_name: Optional[str] = None,
+    scope: IdempotencyScope = IdempotencyScope.SCOPE_SESSION_WIDE,
+) -> Any:
+    if func is None:
+        return functools.partial(
+            step,
+            idempotency_key=idempotency_key,
+            agent_id=agent_id,
+            tool_name=tool_name,
+            scope=scope,
+        )
+
+    actual_tool_name = tool_name or func.__name__
+
+    @functools.wraps(func)
+    async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+        ctx = get_context()
+        ctx.sequence += 1
+        seq = ctx.sequence
+
+        if seq in ctx.replayed_steps:
+            return ctx.replayed_steps[seq].get("result")
+
+        # Derive idempotency key
+        ikey = idempotency_key
+        if not ikey:
+            ikey = derive_idempotency_key(
+                ctx.session_id,
+                ctx.workflow_version,
+                seq,
+                agent_id,
+                actual_tool_name,
+                scope,
+                *args,
+                **kwargs,
+            )
+
+        fencing_token = 0
+        hb: Optional[LeaseHeartbeat] = None
+
+        while True:
+            resp = ctx.client.check_idempotency_cache(
+                agent_id=agent_id,
+                idempotency_key=ikey,
+            )
+            if resp.status == CACHE_STATUS_HIT:
+                if resp.cached_result and resp.cached_result.output_payload:
+                    deserialized = deserialize(resp.cached_result.output_payload)
+                    return deserialized.get("result")
+                return None
+            elif resp.status == CACHE_STATUS_IN_PROGRESS:
+                retry_ms = resp.retry_after_ms or 1000
+                logger.info(
+                    "Step %s is in progress by another worker. Sleeping %d ms...",
+                    actual_tool_name,
+                    retry_ms,
+                )
+                await asyncio.sleep(retry_ms / 1000.0)
+            elif resp.status == CACHE_STATUS_ACQUIRED:
+                fencing_token = resp.fencing_token or 0
+                interval_ms = resp.heartbeat_interval_ms or 5000
+                hb = LeaseHeartbeat(
+                    client=ctx.client,
+                    agent_id=agent_id,
+                    idempotency_key=ikey,
+                    fencing_token=fencing_token,
+                    heartbeat_interval_ms=interval_ms,
+                )
+                hb.start_async()
+                break
+            else:
+                break
+
+        try:
+            result = await func(*args, **kwargs)
+            ctx.client.commit_step(
+                session_id=ctx.session_id,
+                sequence=seq,
+                name=actual_tool_name,
+                status=STEP_STATUS_SUCCESS,
+                output_payload={"result": result},
+                idempotency_key=ikey,
+                idempotency_fencing_token=fencing_token,
+            )
+            return result
+        except Exception as e:
+            if fencing_token > 0:
+                ctx.client.release_lease(
+                    agent_id=agent_id,
+                    idempotency_key=ikey,
+                    fencing_token=fencing_token,
+                    reason="TOOL_ERROR",
+                )
+            raise e
+        finally:
+            if hb:
+                await hb.stop_async()
+
+    @functools.wraps(func)
+    def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+        ctx = get_context()
+        ctx.sequence += 1
+        seq = ctx.sequence
+
+        if seq in ctx.replayed_steps:
+            return ctx.replayed_steps[seq].get("result")
+
+        # Derive idempotency key
+        ikey = idempotency_key
+        if not ikey:
+            ikey = derive_idempotency_key(
+                ctx.session_id,
+                ctx.workflow_version,
+                seq,
+                agent_id,
+                actual_tool_name,
+                scope,
+                *args,
+                **kwargs,
+            )
+
+        fencing_token = 0
+        hb: Optional[LeaseHeartbeat] = None
+
+        while True:
+            resp = ctx.client.check_idempotency_cache(
+                agent_id=agent_id,
+                idempotency_key=ikey,
+            )
+            if resp.status == CACHE_STATUS_HIT:
+                if resp.cached_result and resp.cached_result.output_payload:
+                    deserialized = deserialize(resp.cached_result.output_payload)
+                    return deserialized.get("result")
+                return None
+            elif resp.status == CACHE_STATUS_IN_PROGRESS:
+                retry_ms = resp.retry_after_ms or 1000
+                logger.info(
+                    "Step %s is in progress by another worker. Sleeping %d ms...",
+                    actual_tool_name,
+                    retry_ms,
+                )
+                time.sleep(retry_ms / 1000.0)
+            elif resp.status == CACHE_STATUS_ACQUIRED:
+                fencing_token = resp.fencing_token or 0
+                interval_ms = resp.heartbeat_interval_ms or 5000
+                hb = LeaseHeartbeat(
+                    client=ctx.client,
+                    agent_id=agent_id,
+                    idempotency_key=ikey,
+                    fencing_token=fencing_token,
+                    heartbeat_interval_ms=interval_ms,
+                )
+                hb.start_sync()
+                break
+            else:
+                break
+
+        try:
+            result = func(*args, **kwargs)
+            ctx.client.commit_step(
+                session_id=ctx.session_id,
+                sequence=seq,
+                name=actual_tool_name,
+                status=STEP_STATUS_SUCCESS,
+                output_payload={"result": result},
+                idempotency_key=ikey,
+                idempotency_fencing_token=fencing_token,
+            )
+            return result
+        except Exception as e:
+            if fencing_token > 0:
+                ctx.client.release_lease(
+                    agent_id=agent_id,
+                    idempotency_key=ikey,
+                    fencing_token=fencing_token,
+                    reason="TOOL_ERROR",
+                )
+            raise e
+        finally:
+            if hb:
+                hb.stop_sync()
+
+    if inspect.iscoroutinefunction(func):
+        return cast(F, async_wrapper)
+    return cast(F, sync_wrapper)
+
+
+def tool(
+    func: Optional[F] = None,
+    *,
+    idempotency_key: Optional[str] = None,
+    agent_id: str = "default",
+    tool_name: Optional[str] = None,
+    scope: IdempotencyScope = IdempotencyScope.SCOPE_SESSION_WIDE,
+) -> Any:
+    """
+    Decorator for tool steps that require idempotency tracking.
+    Functionally identical to @step for now, but semantically distinct.
+    """
+    return step(
+        func=func,
+        idempotency_key=idempotency_key,
+        agent_id=agent_id,
+        tool_name=tool_name,
+        scope=scope,
+    )