pyworkflow-engine 0.2.15__tar.gz → 0.3.1__tar.gz

{pyworkflow_engine-0.2.15 → pyworkflow_engine-0.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyworkflow-engine
-Version: 0.2.15
+Version: 0.3.1
 Summary: A Python implementation of durable, event-sourced workflows inspired by Vercel Workflow
 Author: PyWorkflow Contributors
 License: MIT

{pyworkflow_engine-0.2.15 → pyworkflow_engine-0.3.1}/pyproject.toml

@@ -7,7 +7,7 @@ packages = [{include = "pyworkflow"}]
 
 [project]
 name = "pyworkflow-engine"
-version = "0.2.15"
+version = "0.3.1"
 description = "A Python implementation of durable, event-sourced workflows inspired by Vercel Workflow"
 readme = "README.md"
 requires-python = ">=3.11"

{pyworkflow_engine-0.2.15 → pyworkflow_engine-0.3.1}/pyworkflow/__init__.py

@@ -29,7 +29,7 @@ Quick Start:
     >>> run_id = await start(my_workflow, "Alice")
 """
 
-__version__ = "0.2.15"
+__version__ = "0.3.1"
 
 # Configuration
 from pyworkflow.config import (

{pyworkflow_engine-0.2.15 → pyworkflow_engine-0.3.1}/pyworkflow/celery/singleton.py

@@ -11,8 +11,9 @@ Based on:
 
 import inspect
 import json
+from collections.abc import Callable
 from hashlib import md5
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, TypeVar
 from uuid import uuid4
 
 from celery import Task
@@ -20,8 +21,11 @@ from celery.exceptions import WorkerLostError
 from loguru import logger
 
 if TYPE_CHECKING:
+    from redis import Redis
     from redis.sentinel import Sentinel
 
+T = TypeVar("T")
+
 
 def generate_lock_key(
     task_name: str,
@@ -74,10 +78,17 @@ class SingletonConfig:
 
 
 class RedisLockBackend:
-    """Redis backend for distributed locking with Sentinel support."""
+    """Redis backend for distributed locking with Sentinel support.
+
+    During Sentinel failover (master rotation), cached connections can raise
+    MasterNotFoundError or socket TimeoutError. On those, we refresh the client
+    via ``sentinel.master_for()`` and retry the operation once. See
+    ``_execute_with_refresh``.
+    """
 
     _sentinel: "Sentinel | None"
     _master_name: str | None
+    redis: "Redis"
 
     def __init__(
         self,
@@ -91,21 +102,62 @@ class RedisLockBackend:
             from redis.sentinel import Sentinel
 
             sentinels = self._parse_sentinel_url(url)
+            # socket_timeout=5 matches FlowHunt's production-tested value.
+            # The previous 0.5s was too tight during Sentinel failover windows.
             self._sentinel = Sentinel(
                 sentinels,
-                socket_timeout=0.5,
+                socket_timeout=5,
                 decode_responses=True,
             )
             self._master_name = sentinel_master or "mymaster"
-            self.redis = self._sentinel.master_for(
-                self._master_name,
-                decode_responses=True,
-            )
+            self.redis = self._build_master_client()
         else:
             self._sentinel = None
             self._master_name = None
             self.redis = redis.from_url(url, decode_responses=True)
 
+    def _build_master_client(self) -> "Redis":
+        """Get a fresh master client from Sentinel. Sentinel mode only."""
+        assert self._sentinel is not None and self._master_name is not None
+        return self._sentinel.master_for(
+            self._master_name,
+            decode_responses=True,
+        )
+
+    def _refresh_client(self) -> None:
+        """Re-resolve master via Sentinel after a failover-related error."""
+        if self._sentinel is None:
+            return
+        self.redis = self._build_master_client()
+        logger.info(
+            "Refreshed Redis master connection for sentinel master '{}'",
+            self._master_name,
+        )
+
+    def _execute_with_refresh(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> T:
+        """Run ``fn``; on Sentinel failover errors, refresh the client and retry once.
+
+        In non-sentinel mode this is a no-op passthrough — refresh has nothing
+        to do, so any exception propagates as before.
+        """
+        from redis.exceptions import ConnectionError as RedisConnectionError
+        from redis.exceptions import TimeoutError as RedisTimeoutError
+        from redis.sentinel import MasterNotFoundError
+
+        try:
+            return fn(*args, **kwargs)
+        except (MasterNotFoundError, RedisConnectionError, RedisTimeoutError) as exc:
+            if self._sentinel is None:
+                raise
+            logger.warning(
+                "Redis lock op hit sentinel failover error ({}: {}). "
+                "Refreshing master and retrying once.",
+                type(exc).__name__,
+                exc,
+            )
+            self._refresh_client()
+            return fn(*args, **kwargs)
+
     @staticmethod
     def _parse_sentinel_url(url: str) -> list[tuple[str, int]]:
         """
@@ -142,15 +194,19 @@ class RedisLockBackend:
 
     def lock(self, lock_key: str, task_id: str, expiry: int | None = None) -> bool:
         """Acquire lock atomically. Returns True if acquired."""
-        return bool(self.redis.set(lock_key, task_id, nx=True, ex=expiry))
+        return bool(
+            self._execute_with_refresh(
+                lambda: self.redis.set(lock_key, task_id, nx=True, ex=expiry)
+            )
+        )
 
     def unlock(self, lock_key: str) -> None:
         """Release the lock."""
-        self.redis.delete(lock_key)
+        self._execute_with_refresh(lambda: self.redis.delete(lock_key))
 
     def get(self, lock_key: str) -> str | None:
         """Get the task ID holding the lock."""
-        return self.redis.get(lock_key)
+        return self._execute_with_refresh(lambda: self.redis.get(lock_key))
 
 
 class DuplicateTaskError(Exception):
@@ -364,6 +420,16 @@ class SingletonWorkflowTask(Task):
 
         # Lock not acquired - check for existing task
         existing_task_id = self.get_existing_task_id(lock_key)
+
+        # Self-retry: when celery's Task.retry() calls apply_async with the
+        # same task_id we already hold, it's the running task re-queueing
+        # itself. Allow the queue write through — on_retry will release the
+        # lock before the retried message is picked up. Without this branch,
+        # apply_async would silently return AsyncResult without enqueueing,
+        # and the retry would be lost.
+        if existing_task_id == task_id:
+            return super().apply_async(args, kwargs, task_id=task_id, **options)
+
         if existing_task_id:
             logger.debug(
                 "Singleton: duplicate task blocked",
@@ -427,13 +493,17 @@ class SingletonWorkflowTask(Task):
                     error=str(exc),
                 )
 
-        # Log appropriately
+        # Log appropriately. Keep str(exc) out of the message string: when
+        # kwargs are passed, Loguru calls .format() on the message, which
+        # parses any `{` from the exception (e.g. JSON in the message) as a
+        # format spec and raises KeyError.
         if isinstance(exc, WorkerLostError):
             logger.warning("Task interrupted due to worker loss", task_id=task_id)
         else:
             logger.error(
-                f"Task {self.name} failed: {exc}",
+                f"Task {self.name} failed",
                 task_id=task_id,
+                error=str(exc),
                 traceback=einfo.traceback if einfo else None,
             )
 

{pyworkflow_engine-0.2.15 → pyworkflow_engine-0.3.1}/pyworkflow/celery/tasks.py

@@ -10,6 +10,7 @@ These tasks enable:
 """
 
 import asyncio
+import hashlib
 import random
 import uuid
 from collections.abc import Callable
@@ -93,6 +94,8 @@ def execute_step_task(
     context_data: dict[str, Any] | None = None,
     context_class_name: str | None = None,
     workflow_name: str | None = None,
+    tracing: dict[str, Any] | None = None,
+    is_generator: bool = False,
 ) -> Any:
     """
     Execute a workflow step on a Celery worker.
@@ -242,6 +245,7 @@ def execute_step_task(
             step_workflow_ctx._runtime = "celery"
             step_workflow_ctx._storage_config = storage_config
             step_workflow_ctx._is_step_worker = True
+            step_workflow_ctx.tracing = tracing
 
             # Set parent_run_id if this workflow is a child
             wf_run = run_async(storage.get_run(run_id))
@@ -284,6 +288,26 @@ def execute_step_task(
             step_id=step_id,
         )
 
+        # Record tracing span for this step on the worker
+        if tracing and result is not None:
+            try:
+                from pyworkflow.tracing import create_tracing_provider
+
+                _tp = create_tracing_provider(tracing)
+                if _tp:
+                    _tp.record_step_span(
+                        trace_id=hashlib.md5(run_id.encode()).hexdigest(),
+                        step_name=step_name,
+                        step_id=step_id,
+                        is_generator=is_generator,
+                        result=result,
+                        step_input=dict(kwargs) if kwargs else None,
+                        trace_name=workflow_name or "workflow",
+                    )
+                    run_async(_tp.shutdown())
+            except Exception:
+                pass  # Never fail step because of tracing
+
         # Record STEP_COMPLETED event and trigger workflow resumption
         run_async(
             _record_step_completion_and_resume(
@@ -418,8 +442,12 @@ def execute_step_task(
         if self.request.retries < max_retries:
             # Use exponential backoff for unexpected errors
             countdown = _calculate_exponential_backoff(self.request.retries)
+            # Don't inline str(e) into the message: Loguru calls .format() on
+            # the message when kwargs are present, and any `{` from the
+            # exception (e.g. a JSON body) is parsed as a format spec and
+            # raises KeyError, which would skip the retry below.
             logger.warning(
-                f"Step failed (unexpected): {step_name}, retrying in {countdown:.1f}s...: {str(e)}",
+                f"Step failed (unexpected): {step_name}, retrying in {countdown:.1f}s...",
                 run_id=run_id,
                 step_id=step_id,
                 error=str(e),
@@ -855,6 +883,7 @@ def start_workflow_task(
     run_id: str,
     storage_config: dict[str, Any] | None = None,
     idempotency_key: str | None = None,
+    tracing: dict[str, Any] | None = None,
 ) -> str:
     """
     Start a workflow execution.
@@ -906,6 +935,7 @@ def start_workflow_task(
             storage_config=storage_config,
             idempotency_key=idempotency_key,
             run_id=run_id,
+            tracing=tracing,
         )
     )
 
@@ -1431,6 +1461,9 @@ async def _recover_workflow_on_worker(
     args = deserialize_args(run.input_args)
     kwargs = deserialize_kwargs(run.input_kwargs)
 
+    # Extract tracing config persisted at start
+    _recover_tracing = kwargs.pop("_tracing_config", None)
+
     # Execute workflow with event replay
     try:
         result = await execute_workflow_with_context(
@@ -1444,6 +1477,7 @@ async def _recover_workflow_on_worker(
             runtime="celery",
             storage_config=storage_config,
             parent_run_id=run.parent_run_id,
+            tracing=workflow_meta.tracing or _recover_tracing,
         )
 
         # Update run status to completed
@@ -1579,6 +1613,7 @@ async def _start_workflow_on_worker(
     storage_config: dict[str, Any] | None = None,
     idempotency_key: str | None = None,
     run_id: str | None = None,
+    tracing: dict[str, Any] | None = None,
 ) -> str:
     """
     Internal function to start workflow on Celery worker.
@@ -1779,10 +1814,10 @@ async def _start_workflow_on_worker(
         created_at=datetime.now(UTC),
         started_at=datetime.now(UTC),
         input_args=serialize_args(*args),
-        input_kwargs=serialize_kwargs(**kwargs),
+        input_kwargs=serialize_kwargs(**kwargs, _tracing_config=tracing),
         idempotency_key=idempotency_key,
         max_duration=workflow_meta.max_duration,
-        context={},  # Step context (not from decorator)
+        context={},  # Step context data
         recovery_attempts=0,
         max_recovery_attempts=max_recovery_attempts,
         recover_on_worker_loss=recover_on_worker_loss,
@@ -1812,6 +1847,7 @@ async def _start_workflow_on_worker(
             kwargs=kwargs,
             runtime="celery",
             storage_config=storage_config,
+            tracing=tracing or workflow_meta.tracing,
         )
 
         # Update run status to completed
@@ -2392,6 +2428,9 @@ async def _resume_workflow_on_worker(
     args = deserialize_args(run.input_args)
     kwargs = deserialize_kwargs(run.input_kwargs)
 
+    # Extract tracing config persisted at start
+    _resume_tracing = kwargs.pop("_tracing_config", None)
+
     # Execute workflow with event replay
     try:
         result = await execute_workflow_with_context(
@@ -2406,6 +2445,7 @@ async def _resume_workflow_on_worker(
             runtime="celery",
             storage_config=storage_config,
             parent_run_id=run.parent_run_id,
+            tracing=workflow_meta.tracing or _resume_tracing,
         )
 
         # Update run status to completed

{pyworkflow_engine-0.2.15 → pyworkflow_engine-0.3.1}/pyworkflow/config.py

@@ -171,6 +171,9 @@ class PyWorkflowConfig:
     event_hard_limit: int = 50_000  # Fail workflow at this count
     event_warning_interval: int = 100  # Log warning every N events after soft limit
 
+    # Tracing Configuration
+    TRACING_CONFIG: dict | None = None
+
 
 def _config_from_env_and_yaml() -> PyWorkflowConfig:
     """

{pyworkflow_engine-0.2.15 → pyworkflow_engine-0.3.1}/pyworkflow/context/base.py

@@ -133,6 +133,9 @@ class WorkflowContext(ABC):
         """
         self._run_id = run_id
         self._workflow_name = workflow_name
+        self._tracing: dict[str, Any] | None = None
+        self._tracing_provider: Any = None  # TracingProvider instance
+        self._trace_id: str | None = None  # Stable trace ID
 
     @property
     def run_id(self) -> str:
@@ -144,6 +147,20 @@ class WorkflowContext(ABC):
         """Get the current workflow name."""
         return self._workflow_name
 
+    @property
+    def tracing(self) -> dict[str, Any] | None:
+        """Get the tracing provider config (e.g. Langfuse credentials), or None."""
+        return self._tracing
+
+    @tracing.setter
+    def tracing(self, value: dict[str, Any] | None) -> None:
+        self._tracing = value
+
+    @property
+    def tracing_provider(self) -> Any:
+        """Get the active TracingProvider instance, or None."""
+        return self._tracing_provider
+
     # =========================================================================
     # Abstract methods - must be implemented by subclasses
     # =========================================================================

{pyworkflow_engine-0.2.15 → pyworkflow_engine-0.3.1}/pyworkflow/core/registry.py

@@ -23,6 +23,9 @@ class WorkflowMetadata:
     tags: list[str] | None = None
     description: str | None = None  # Docstring from the workflow function
     context_class: type | None = None  # StepContext subclass for step context access
+    tracing: dict[str, Any] | None = (
+        None  # Optional tracing provider config (e.g. Langfuse credentials)
+    )
 
     def __post_init__(self) -> None:
         if self.tags is None:
@@ -44,6 +47,9 @@ class StepMetadata:
     timeout: int | None = None
     metadata: dict[str, Any] | None = None
     force_local: bool = False
+    is_generator: bool = (
+        False  # If True, creates a Langfuse generation span instead of a regular span
+    )
 
     def __post_init__(self) -> None:
         if self.metadata is None:
@@ -73,6 +79,7 @@ class WorkflowRegistry:
         max_duration: str | None = None,
         tags: list[str] | None = None,
         context_class: type | None = None,
+        tracing: dict[str, Any] | None = None,
     ) -> None:
         """
         Register a workflow.
@@ -84,6 +91,7 @@ class WorkflowRegistry:
             max_duration: Optional maximum duration
             tags: Optional list of tags (max 3)
             context_class: Optional StepContext subclass for step context access
+            tracing: Optional tracing provider config dict (e.g. Langfuse credentials)
         """
         if name in self._workflows:
             existing = self._workflows[name]
@@ -101,6 +109,7 @@ class WorkflowRegistry:
             max_duration=max_duration,
             tags=tags or [],
             context_class=context_class,
+            tracing=tracing,
         )
 
         self._workflows[name] = workflow_meta
@@ -160,6 +169,7 @@ class WorkflowRegistry:
         timeout: int | None = None,
         metadata: dict[str, Any] | None = None,
         force_local: bool = False,
+        is_generator: bool = False,
     ) -> None:
         """
         Register a step.
@@ -173,6 +183,7 @@ class WorkflowRegistry:
             timeout: Optional timeout in seconds
             metadata: Optional metadata dict
             force_local: If True, always execute inline even in distributed runtimes
+            is_generator: If True, creates a Langfuse generation span instead of a regular span
         """
         if name in self._steps:
             existing = self._steps[name]
@@ -190,6 +201,7 @@ class WorkflowRegistry:
             timeout=timeout,
             metadata=metadata or {},
             force_local=force_local,
+            is_generator=is_generator,
         )
 
         self._steps[name] = step_meta
@@ -259,9 +271,12 @@ def register_workflow(
     max_duration: str | None = None,
     tags: list[str] | None = None,
     context_class: type | None = None,
+    tracing: dict[str, Any] | None = None,
 ) -> None:
     """Register a workflow in the global registry."""
-    _registry.register_workflow(name, func, original_func, max_duration, tags, context_class)
+    _registry.register_workflow(
+        name, func, original_func, max_duration, tags, context_class, tracing
+    )
 
 
 def get_workflow(name: str) -> WorkflowMetadata | None:
@@ -293,10 +308,19 @@ def register_step(
     timeout: int | None = None,
     metadata: dict[str, Any] | None = None,
     force_local: bool = False,
+    is_generator: bool = False,
 ) -> None:
     """Register a step in the global registry."""
     _registry.register_step(
-        name, func, original_func, max_retries, retry_delay, timeout, metadata, force_local
+        name,
+        func,
+        original_func,
+        max_retries,
+        retry_delay,
+        timeout,
+        metadata,
+        force_local,
+        is_generator,
     )
 
 

{pyworkflow_engine-0.2.15 → pyworkflow_engine-0.3.1}/pyworkflow/core/step.py

@@ -58,6 +58,7 @@ def step(
     timeout: int | None = None,
     metadata: dict[str, Any] | None = None,
     force_local: bool = False,
+    is_generator: bool = False,
 ) -> Callable:
     """
     Decorator to mark functions as workflow steps.
@@ -77,6 +78,8 @@ def step(
         force_local: If True, always execute inline even in distributed
             runtimes (e.g., Celery). Useful for lightweight steps where
             the overhead of message broker dispatch is undesirable.
+        is_generator: If True, creates a Langfuse generation span instead of
+            a regular span. Use for LLM/AI steps that produce generated content.
 
     Returns:
         Decorated step function
@@ -209,6 +212,7 @@ def step(
                     max_retries=max_retries,
                     retry_delay=retry_delay,
                     timeout=timeout,
+                    is_generator=is_generator,
                 )
 
             # Check if we're resuming from a retry
@@ -297,6 +301,9 @@ def step(
                     step_id=step_id,
                 )
 
+                # Record tracing span for locally executed step
+                _record_step_tracing(ctx, step_name, step_id, is_generator, result, args, kwargs)
+
                 return result
 
             except SuspensionSignal:
@@ -458,6 +465,7 @@ def step(
             timeout=timeout,
             metadata=metadata,
             force_local=force_local,
+            is_generator=is_generator,
         )
 
         # Store metadata on wrapper
@@ -468,6 +476,7 @@ def step(
         wrapper.__step_timeout__ = timeout  # type: ignore[attr-defined]
         wrapper.__step_metadata__ = metadata or {}  # type: ignore[attr-defined]
         wrapper.__step_force_local__ = force_local  # type: ignore[attr-defined]
+        wrapper.__step_is_generator__ = is_generator  # type: ignore[attr-defined]
 
         return wrapper
 
@@ -627,6 +636,7 @@ async def _dispatch_step_to_celery(
     max_retries: int,
     retry_delay: str | int | list[int],
     timeout: int | None,
+    is_generator: bool = False,
 ) -> Any:
     """
     Dispatch step execution to Celery step worker.
@@ -736,6 +746,8 @@ async def _dispatch_step_to_celery(
         context_data=context_data,
         context_class_name=context_class_name,
         workflow_name=ctx.workflow_name,
+        tracing=ctx.tracing,
+        is_generator=is_generator,
     )
 
     logger.info(
@@ -753,3 +765,38 @@ async def _dispatch_step_to_celery(
         step_name=step_name,
         task_id=task_result.id,
     )
+
+
+def _record_step_tracing(
+    ctx: Any,
+    step_name: str,
+    step_id: str,
+    is_generator: bool,
+    result: Any,
+    step_args: tuple = (),
+    step_kwargs: dict | None = None,
+) -> None:
+    """Record a tracing span for a completed step. No-op if tracing is not configured."""
+    tp = getattr(ctx, "_tracing_provider", None)
+    if not tp:
+        return
+    try:
+        trace_id = getattr(ctx, "_trace_id", None)
+        if not trace_id:
+            return
+
+        step_input = dict(step_kwargs) if step_kwargs else {}
+        if step_args:
+            step_input["_args"] = [str(a)[:500] for a in step_args]
+
+        tp.record_step_span(
+            trace_id=trace_id,
+            step_name=step_name,
+            step_id=step_id,
+            is_generator=is_generator,
+            result=result,
+            step_input=step_input or None,
+            trace_name=ctx.workflow_name,
+        )
+    except Exception:
+        pass  # Never fail the workflow because of tracing