agentic-blocks 0.1.18__py3-none-any.whl → 0.1.20__py3-none-any.whl

agentic_blocks/tracing/decorator.py

@@ -0,0 +1,316 @@
+"""
+Decorator for tracing PocketFlow workflows with Langfuse.
+"""
+
+import functools
+import inspect
+import uuid
+from typing import Any, Callable, Dict, Optional, Union
+
+from .config import TracingConfig
+from .core import LangfuseTracer
+
+
+def trace_flow(
+    config: Optional[TracingConfig] = None,
+    flow_name: Optional[str] = None,
+    session_id: Optional[str] = None,
+    user_id: Optional[str] = None,
+):
+    """
+    Decorator to add Langfuse tracing to PocketFlow flows.
+
+    This decorator automatically traces:
+    - Flow execution start/end
+    - Each node's prep, exec, and post phases
+    - Input and output data for each phase
+    - Errors and exceptions
+
+    Args:
+        config: TracingConfig instance. If None, loads from environment.
+        flow_name: Custom name for the flow. If None, uses the flow class name.
+        session_id: Session ID for grouping related traces.
+        user_id: User ID for the trace.
+
+    Returns:
+        Decorated flow class or function.
+
+    Example:
+        ```python
+        from tracing import trace_flow
+
+        @trace_flow()
+        class MyFlow(Flow):
+            def __init__(self):
+                super().__init__(start=MyNode())
+
+        # Or with custom configuration
+        config = TracingConfig.from_env()
+
+        @trace_flow(config=config, flow_name="CustomFlow")
+        class MyFlow(Flow):
+            pass
+        ```
+    """
+
+    def decorator(flow_class_or_func):
+        # Handle both class and function decoration
+        if inspect.isclass(flow_class_or_func):
+            return _trace_flow_class(
+                flow_class_or_func, config, flow_name, session_id, user_id
+            )
+        else:
+            return _trace_flow_function(
+                flow_class_or_func, config, flow_name, session_id, user_id
+            )
+
+    return decorator
+
+
+def _trace_flow_class(flow_class, config, flow_name, session_id, user_id):
+    """Trace a Flow class by wrapping its methods."""
+
+    # Get or create config
+    if config is None:
+        config = TracingConfig.from_env()
+
+    # Override session/user if provided
+    if session_id:
+        config.session_id = session_id
+    if user_id:
+        config.user_id = user_id
+
+    # Get flow name
+    if flow_name is None:
+        flow_name = flow_class.__name__
+
+    # Store original methods
+    original_init = flow_class.__init__
+    original_run = getattr(flow_class, "run", None)
+    original_run_async = getattr(flow_class, "run_async", None)
+
+    def traced_init(self, *args, **kwargs):
+        """Initialize the flow with tracing capabilities."""
+        # Call original init
+        original_init(self, *args, **kwargs)
+
+        # Add tracing attributes
+        self._tracer = LangfuseTracer(config)
+        self._flow_name = flow_name
+        self._trace_id = None
+
+        # Patch all nodes in the flow
+        self._patch_nodes()
+
+    def traced_run(self, shared):
+        """Traced version of the run method."""
+        if not hasattr(self, "_tracer"):
+            # Fallback if not properly initialized
+            return original_run(self, shared) if original_run else None
+
+        # Start trace
+        self._trace_id = self._tracer.start_trace(self._flow_name, shared)
+
+        try:
+            # Run the original flow
+            result = original_run(self, shared) if original_run else None
+
+            # End trace successfully
+            self._tracer.end_trace(shared, "success")
+
+            return result
+
+        except Exception as e:
+            # End trace with error
+            self._tracer.end_trace(shared, "error")
+            raise
+        finally:
+            # Ensure cleanup
+            self._tracer.flush()
+
+    async def traced_run_async(self, shared):
+        """Traced version of the async run method."""
+        if not hasattr(self, "_tracer"):
+            # Fallback if not properly initialized
+            return (
+                await original_run_async(self, shared) if original_run_async else None
+            )
+
+        # Start trace
+        self._trace_id = self._tracer.start_trace(self._flow_name, shared)
+
+        try:
+            # Run the original flow
+            result = (
+                await original_run_async(self, shared) if original_run_async else None
+            )
+
+            # End trace successfully
+            self._tracer.end_trace(shared, "success")
+
+            return result
+
+        except Exception as e:
+            # End trace with error
+            self._tracer.end_trace(shared, "error")
+            raise
+        finally:
+            # Ensure cleanup
+            self._tracer.flush()
+
+    def patch_nodes(self):
+        """Patch all nodes in the flow to add tracing."""
+        if not hasattr(self, "start_node") or not self.start_node:
+            return
+
+        visited = set()
+        nodes_to_patch = [self.start_node]
+
+        while nodes_to_patch:
+            node = nodes_to_patch.pop(0)
+            if id(node) in visited:
+                continue
+
+            visited.add(id(node))
+
+            # Patch this node
+            self._patch_node(node)
+
+            # Add successors to patch list
+            if hasattr(node, "successors"):
+                for successor in node.successors.values():
+                    if successor and id(successor) not in visited:
+                        nodes_to_patch.append(successor)
+
+    def patch_node(self, node):
+        """Patch a single node to add tracing."""
+        if hasattr(node, "_pocketflow_traced"):
+            return  # Already patched
+
+        node_id = str(uuid.uuid4())
+        node_name = type(node).__name__
+
+        # Store original methods
+        original_prep = getattr(node, "prep", None)
+        original_exec = getattr(node, "exec", None)
+        original_post = getattr(node, "post", None)
+        original_prep_async = getattr(node, "prep_async", None)
+        original_exec_async = getattr(node, "exec_async", None)
+        original_post_async = getattr(node, "post_async", None)
+
+        # Create traced versions
+        if original_prep:
+            node.prep = self._create_traced_method(
+                original_prep, node_id, node_name, "prep"
+            )
+        if original_exec:
+            node.exec = self._create_traced_method(
+                original_exec, node_id, node_name, "exec"
+            )
+        if original_post:
+            node.post = self._create_traced_method(
+                original_post, node_id, node_name, "post"
+            )
+        if original_prep_async:
+            node.prep_async = self._create_traced_async_method(
+                original_prep_async, node_id, node_name, "prep"
+            )
+        if original_exec_async:
+            node.exec_async = self._create_traced_async_method(
+                original_exec_async, node_id, node_name, "exec"
+            )
+        if original_post_async:
+            node.post_async = self._create_traced_async_method(
+                original_post_async, node_id, node_name, "post"
+            )
+
+        # Mark as traced
+        node._pocketflow_traced = True
+
+    def create_traced_method(self, original_method, node_id, node_name, phase):
+        """Create a traced version of a synchronous method."""
+
+        @functools.wraps(original_method)
+        def traced_method(*args, **kwargs):
+            span_id = self._tracer.start_node_span(node_name, node_id, phase)
+
+            try:
+                result = original_method(*args, **kwargs)
+                self._tracer.end_node_span(span_id, input_data=args, output_data=result)
+                return result
+            except Exception as e:
+                self._tracer.end_node_span(span_id, input_data=args, error=e)
+                raise
+
+        return traced_method
+
+    def create_traced_async_method(self, original_method, node_id, node_name, phase):
+        """Create a traced version of an asynchronous method."""
+
+        @functools.wraps(original_method)
+        async def traced_async_method(*args, **kwargs):
+            span_id = self._tracer.start_node_span(node_name, node_id, phase)
+
+            try:
+                result = await original_method(*args, **kwargs)
+                self._tracer.end_node_span(span_id, input_data=args, output_data=result)
+                return result
+            except Exception as e:
+                self._tracer.end_node_span(span_id, input_data=args, error=e)
+                raise
+
+        return traced_async_method
+
+    # Replace methods on the class
+    flow_class.__init__ = traced_init
+    flow_class._patch_nodes = patch_nodes
+    flow_class._patch_node = patch_node
+    flow_class._create_traced_method = create_traced_method
+    flow_class._create_traced_async_method = create_traced_async_method
+
+    if original_run:
+        flow_class.run = traced_run
+    if original_run_async:
+        flow_class.run_async = traced_run_async
+
+    return flow_class
+
+
+def _trace_flow_function(flow_func, config, flow_name, session_id, user_id):
+    """Trace a flow function (for functional-style flows)."""
+
+    # Get or create config
+    if config is None:
+        config = TracingConfig.from_env()
+
+    # Override session/user if provided
+    if session_id:
+        config.session_id = session_id
+    if user_id:
+        config.user_id = user_id
+
+    # Get flow name
+    if flow_name is None:
+        flow_name = flow_func.__name__
+
+    tracer = LangfuseTracer(config)
+
+    @functools.wraps(flow_func)
+    def traced_flow_func(*args, **kwargs):
+        # Assume first argument is shared data
+        shared = args[0] if args else {}
+
+        # Start trace
+        trace_id = tracer.start_trace(flow_name, shared)
+
+        try:
+            result = flow_func(*args, **kwargs)
+            tracer.end_trace(shared, "success")
+            return result
+        except Exception as e:
+            tracer.end_trace(shared, "error")
+            raise
+        finally:
+            tracer.flush()
+
+    return traced_flow_func