agentic-blocks 0.1.22__tar.gz → 0.1.24__tar.gz

{agentic_blocks-0.1.22/src/agentic_blocks.egg-info → agentic_blocks-0.1.24}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agentic-blocks
-Version: 0.1.22
+Version: 0.1.24
 Summary: Simple building blocks for agentic AI systems with MCP client and conversation management
 Author-email: Magnus Bjelkenhed <bjelkenhed@gmail.com>
 License: MIT
@@ -27,6 +27,10 @@ Requires-Dist: langchain-core
 Requires-Dist: langfuse
 Requires-Dist: pocketflow
 Requires-Dist: pocketflow-tracing
+Requires-Dist: agno
+Requires-Dist: ipywidgets
+Requires-Dist: rich[jupyter]
+Requires-Dist: fastapi
 Provides-Extra: test
 Requires-Dist: pytest; extra == "test"
 Provides-Extra: dev

{agentic_blocks-0.1.22 → agentic_blocks-0.1.24}/pyproject.toml

@@ -14,7 +14,7 @@ agentic_blocks = []
 
 [project]
 name = "agentic-blocks"
-version = "0.1.22"
+version = "0.1.24"
 description = "Simple building blocks for agentic AI systems with MCP client and conversation management"
 readme = "README.md"
 requires-python = ">=3.11"
@@ -43,6 +43,10 @@ dependencies = [
     "langfuse",
     "pocketflow",
     "pocketflow-tracing",
+    "agno",
+    "ipywidgets", 
+    "rich[jupyter]",
+    "fastapi",
 ]
 
 [project.urls]

{agentic_blocks-0.1.22 → agentic_blocks-0.1.24}/src/agentic_blocks/agent.py

@@ -4,6 +4,18 @@ from agentic_blocks.utils.tools_utils import (
     execute_pending_tool_calls,
 )
 from agentic_blocks import call_llm, Messages
+from rich.panel import Panel
+from rich.box import HEAVY
+from rich.console import Console
+from rich.console import Group
+
+console = Console(
+    style="black on bright_white",
+    force_terminal=True,
+    width=None,
+    legacy_windows=False,
+    color_system="truecolor",
+)
 
 
 class Agent:
@@ -11,6 +23,7 @@ class Agent:
         self.system_prompt = system_prompt
         self.tools = tools
         self.tool_registry = create_tool_registry(tools)
+        self.panels = []
 
         # Create nodes
         self.llm_node = self._create_llm_node()
@@ -32,7 +45,8 @@ class Agent:
                 self.tools = tools
 
             def prep(self, shared):
-                return shared["messages"]
+                messages = shared["messages"]
+                return messages
 
             def exec(self, messages) -> Messages:
                 response = call_llm(messages=messages, tools=self.tools)
@@ -49,14 +63,31 @@ class Agent:
 
     def _create_tool_node(self):
         class ToolNode(Node):
-            def __init__(self, tool_registry):
+            def __init__(self, tool_registry, agent):
                 super().__init__()
                 self.tool_registry = tool_registry
+                self.agent = agent
 
             def prep(self, shared):
                 return shared["messages"]
 
             def exec(self, messages) -> Messages:
+                tool_calls = messages.get_pending_tool_calls()[0]
+                tool_name = tool_calls["tool_name"]
+                tool_arguments = tool_calls["arguments"]
+
+                # Format arguments nicely
+                if isinstance(tool_arguments, dict):
+                    args_str = ", ".join(
+                        [f"{k}={v}" for k, v in tool_arguments.items()]
+                    )
+                    formatted_call = f"{tool_name}({args_str})"
+                else:
+                    formatted_call = f"{tool_name}({tool_arguments})"
+
+                tool_panel = self.agent.create_panel(formatted_call, "Tool Calls")
+                self.agent.panels.append(tool_panel)
+
                 tool_responses = execute_pending_tool_calls(
                     messages, self.tool_registry
                 )
@@ -66,7 +97,7 @@ class Agent:
             def post(self, shared, prep_res, messages):
                 return "llm_node"
 
-        return ToolNode(self.tool_registry)
+        return ToolNode(self.tool_registry, self)
 
     def _create_answer_node(self):
         class AnswerNode(Node):
@@ -86,3 +117,29 @@ class Agent:
         self.flow.run(shared)
 
         return shared["answer"]
+
+    def print_response(self, user_prompt: str, stream: bool = False):
+        # Reset panels and start with message
+        self.panels = []
+        message_panel = self.create_panel(user_prompt, "Message")
+        self.panels.append(message_panel)
+
+        # Always collect all panels first
+        response = self.invoke(user_prompt)
+        response_panel = self.create_panel(response, "Response")
+        self.panels.append(response_panel)
+
+        # Print all panels as a group (no gaps)
+        panel_group = Group(*self.panel)
+        console.print(panel_group)
+
+    def create_panel(self, content, title, border_style="blue"):
+        return Panel(
+            content,
+            title=title,
+            title_align="left",
+            border_style=border_style,
+            box=HEAVY,
+            expand=True,  # Full terminal width
+            padding=(1, 1),  # Internal padding
+        )

{agentic_blocks-0.1.22 → agentic_blocks-0.1.24}/src/agentic_blocks/messages.py

@@ -175,6 +175,13 @@ class Messages:
         """Get the current messages list."""
         return self.messages
 
+    def get_user_message(self) -> str:
+        """Get the user message."""
+        for message in reversed(self.messages):
+            if message.get("role") == "user":
+                return message.get("content") or ""
+        return ""
+
     def has_pending_tool_calls(self) -> bool:
         """
         Check if the last message has tool calls that need execution.
@@ -213,7 +220,7 @@ class Messages:
             List of dictionaries with 'tool_name', 'arguments', and 'tool_call_id' keys
         """
         pending_calls = []
-        
+
         if not self.messages:
             return pending_calls
 
@@ -234,19 +241,22 @@ class Messages:
                     function_info = tool_call.get("function", {})
                     tool_name = function_info.get("name")
                     arguments_str = function_info.get("arguments", "{}")
-                    
+
                     # Parse arguments JSON string to dict
                     import json
+
                     try:
                         arguments = json.loads(arguments_str)
                     except json.JSONDecodeError:
                         arguments = {}
-                    
-                    pending_calls.append({
-                        "tool_name": tool_name,
-                        "arguments": arguments,
-                        "tool_call_id": tool_call_id
-                    })
+
+                    pending_calls.append(
+                        {
+                            "tool_name": tool_name,
+                            "arguments": arguments,
+                            "tool_call_id": tool_call_id,
+                        }
+                    )
 
         return pending_calls
 

{agentic_blocks-0.1.22 → agentic_blocks-0.1.24/src/agentic_blocks.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agentic-blocks
-Version: 0.1.22
+Version: 0.1.24
 Summary: Simple building blocks for agentic AI systems with MCP client and conversation management
 Author-email: Magnus Bjelkenhed <bjelkenhed@gmail.com>
 License: MIT
@@ -27,6 +27,10 @@ Requires-Dist: langchain-core
 Requires-Dist: langfuse
 Requires-Dist: pocketflow
 Requires-Dist: pocketflow-tracing
+Requires-Dist: agno
+Requires-Dist: ipywidgets
+Requires-Dist: rich[jupyter]
+Requires-Dist: fastapi
 Provides-Extra: test
 Requires-Dist: pytest; extra == "test"
 Provides-Extra: dev

{agentic_blocks-0.1.22 → agentic_blocks-0.1.24}/src/agentic_blocks.egg-info/SOURCES.txt

@@ -11,9 +11,5 @@ src/agentic_blocks.egg-info/SOURCES.txt
 src/agentic_blocks.egg-info/dependency_links.txt
 src/agentic_blocks.egg-info/requires.txt
 src/agentic_blocks.egg-info/top_level.txt
-src/agentic_blocks/tracing/__init__.py
-src/agentic_blocks/tracing/config.py
-src/agentic_blocks/tracing/core.py
-src/agentic_blocks/tracing/decorator.py
 src/agentic_blocks/utils/tools_utils.py
 src/agentic_blocks/visualization/visualize.py

{agentic_blocks-0.1.22 → agentic_blocks-0.1.24}/src/agentic_blocks.egg-info/requires.txt

@@ -6,6 +6,10 @@ langchain-core
 langfuse
 pocketflow
 pocketflow-tracing
+agno
+ipywidgets
+rich[jupyter]
+fastapi
 
 [dev]
 pytest

agentic_blocks-0.1.22/src/agentic_blocks/tracing/__init__.py

@@ -1,13 +0,0 @@
-"""
-PocketFlow Tracing Module
-
-This module provides observability and tracing capabilities for PocketFlow workflows
-using Langfuse as the backend. It includes decorators and utilities to automatically
-trace node execution, inputs, and outputs.
-"""
-
-from .config import TracingConfig
-from .core import LangfuseTracer
-from .decorator import trace_flow
-
-__all__ = ["trace_flow", "TracingConfig", "LangfuseTracer"]

agentic_blocks-0.1.22/src/agentic_blocks/tracing/config.py

@@ -1,111 +0,0 @@
-"""
-Configuration module for PocketFlow tracing with Langfuse.
-"""
-
-import os
-from dataclasses import dataclass
-from typing import Optional
-from dotenv import load_dotenv
-
-
-@dataclass
-class TracingConfig:
-    """Configuration class for PocketFlow tracing with Langfuse."""
-    
-    # Langfuse configuration
-    langfuse_secret_key: Optional[str] = None
-    langfuse_public_key: Optional[str] = None
-    langfuse_host: Optional[str] = None
-    
-    # PocketFlow tracing configuration
-    debug: bool = False
-    trace_inputs: bool = True
-    trace_outputs: bool = True
-    trace_prep: bool = True
-    trace_exec: bool = True
-    trace_post: bool = True
-    trace_errors: bool = True
-    
-    # Session configuration
-    session_id: Optional[str] = None
-    user_id: Optional[str] = None
-    
-    @classmethod
-    def from_env(cls, env_file: Optional[str] = None) -> "TracingConfig":
-        """
-        Create TracingConfig from environment variables.
-        
-        Args:
-            env_file: Optional path to .env file. If None, looks for .env in current directory.
-            
-        Returns:
-            TracingConfig instance with values from environment variables.
-        """
-        # Load environment variables from .env file if it exists
-        if env_file:
-            load_dotenv(env_file)
-        else:
-            # Try to find .env file in current directory or parent directories
-            load_dotenv()
-        
-        return cls(
-            langfuse_secret_key=os.getenv("LANGFUSE_SECRET_KEY"),
-            langfuse_public_key=os.getenv("LANGFUSE_PUBLIC_KEY"),
-            langfuse_host=os.getenv("LANGFUSE_HOST"),
-            debug=os.getenv("POCKETFLOW_TRACING_DEBUG", "false").lower() == "true",
-            trace_inputs=os.getenv("POCKETFLOW_TRACE_INPUTS", "true").lower() == "true",
-            trace_outputs=os.getenv("POCKETFLOW_TRACE_OUTPUTS", "true").lower() == "true",
-            trace_prep=os.getenv("POCKETFLOW_TRACE_PREP", "true").lower() == "true",
-            trace_exec=os.getenv("POCKETFLOW_TRACE_EXEC", "true").lower() == "true",
-            trace_post=os.getenv("POCKETFLOW_TRACE_POST", "true").lower() == "true",
-            trace_errors=os.getenv("POCKETFLOW_TRACE_ERRORS", "true").lower() == "true",
-            session_id=os.getenv("POCKETFLOW_SESSION_ID"),
-            user_id=os.getenv("POCKETFLOW_USER_ID"),
-        )
-    
-    def validate(self) -> bool:
-        """
-        Validate that required configuration is present.
-        
-        Returns:
-            True if configuration is valid, False otherwise.
-        """
-        if not self.langfuse_secret_key:
-            if self.debug:
-                print("Warning: LANGFUSE_SECRET_KEY not set")
-            return False
-            
-        if not self.langfuse_public_key:
-            if self.debug:
-                print("Warning: LANGFUSE_PUBLIC_KEY not set")
-            return False
-            
-        if not self.langfuse_host:
-            if self.debug:
-                print("Warning: LANGFUSE_HOST not set")
-            return False
-            
-        return True
-    
-    def to_langfuse_kwargs(self) -> dict:
-        """
-        Convert configuration to kwargs for Langfuse client initialization.
-        
-        Returns:
-            Dictionary of kwargs for Langfuse client.
-        """
-        kwargs = {}
-        
-        if self.langfuse_secret_key:
-            kwargs["secret_key"] = self.langfuse_secret_key
-            
-        if self.langfuse_public_key:
-            kwargs["public_key"] = self.langfuse_public_key
-            
-        if self.langfuse_host:
-            kwargs["host"] = self.langfuse_host
-            
-        if self.debug:
-            kwargs["debug"] = True
-            
-        return kwargs

agentic_blocks-0.1.22/src/agentic_blocks/tracing/core.py

@@ -1,287 +0,0 @@
-"""
-Core tracing functionality for PocketFlow with Langfuse integration.
-"""
-
-import json
-import time
-import uuid
-from typing import Any, Dict, Optional, Union
-from datetime import datetime
-
-try:
-    from langfuse import Langfuse
-
-    LANGFUSE_AVAILABLE = True
-except ImportError:
-    LANGFUSE_AVAILABLE = False
-    print("Warning: langfuse package not installed. Install with: pip install langfuse")
-
-from .config import TracingConfig
-
-
-class LangfuseTracer:
-    """
-    Core tracer class that handles Langfuse integration for PocketFlow.
-    """
-
-    def __init__(self, config: TracingConfig):
-        """
-        Initialize the LangfuseTracer.
-
-        Args:
-            config: TracingConfig instance with Langfuse settings.
-        """
-        self.config = config
-        self.client = None
-        self.current_trace = None
-        self.spans = {}  # Store spans by node ID
-
-        if LANGFUSE_AVAILABLE and config.validate():
-            try:
-                # Initialize Langfuse client with proper parameters
-                kwargs = {}
-                if config.langfuse_secret_key:
-                    kwargs["secret_key"] = config.langfuse_secret_key
-                if config.langfuse_public_key:
-                    kwargs["public_key"] = config.langfuse_public_key
-                if config.langfuse_host:
-                    kwargs["host"] = config.langfuse_host
-                if config.debug:
-                    kwargs["debug"] = True
-
-                self.client = Langfuse(**kwargs)
-                if config.debug:
-                    print(
-                        f"✓ Langfuse client initialized with host: {config.langfuse_host}"
-                    )
-            except Exception as e:
-                if config.debug:
-                    print(f"✗ Failed to initialize Langfuse client: {e}")
-                self.client = None
-        else:
-            if config.debug:
-                print("✗ Langfuse not available or configuration invalid")
-
-    def start_trace(self, flow_name: str, input_data: Dict[str, Any]) -> Optional[str]:
-        """
-        Start a new trace for a flow execution.
-
-        Args:
-            flow_name: Name of the flow being traced.
-            input_data: Input data for the flow.
-
-        Returns:
-            Trace ID if successful, None otherwise.
-        """
-        if not self.client:
-            return None
-
-        try:
-            # Serialize input data safely
-            serialized_input = self._serialize_data(input_data)
-
-            # Use Langfuse v2 API to create a trace
-            self.current_trace = self.client.trace(
-                name=flow_name,
-                input=serialized_input,
-                metadata={
-                    "framework": "PocketFlow",
-                    "trace_type": "flow_execution",
-                    "timestamp": datetime.now().isoformat(),
-                },
-                session_id=self.config.session_id,
-                user_id=self.config.user_id,
-            )
-
-            # Get the trace ID
-            trace_id = self.current_trace.id
-
-            if self.config.debug:
-                print(f"✓ Started trace: {trace_id} for flow: {flow_name}")
-
-            return trace_id
-
-        except Exception as e:
-            if self.config.debug:
-                print(f"✗ Failed to start trace: {e}")
-            return None
-
-    def end_trace(self, output_data: Dict[str, Any], status: str = "success") -> None:
-        """
-        End the current trace.
-
-        Args:
-            output_data: Output data from the flow.
-            status: Status of the trace execution.
-        """
-        if not self.current_trace:
-            return
-
-        try:
-            # Serialize output data safely
-            serialized_output = self._serialize_data(output_data)
-
-            # Update the trace with output data using v2 API
-            self.current_trace.update(
-                output=serialized_output,
-                metadata={
-                    "status": status,
-                    "end_timestamp": datetime.now().isoformat(),
-                },
-            )
-
-            if self.config.debug:
-                print(f"✓ Ended trace with status: {status}")
-
-        except Exception as e:
-            if self.config.debug:
-                print(f"✗ Failed to end trace: {e}")
-        finally:
-            self.current_trace = None
-            self.spans.clear()
-
-    def start_node_span(
-        self, node_name: str, node_id: str, phase: str
-    ) -> Optional[str]:
-        """
-        Start a span for a node execution phase.
-
-        Args:
-            node_name: Name/type of the node.
-            node_id: Unique identifier for the node instance.
-            phase: Execution phase (prep, exec, post).
-
-        Returns:
-            Span ID if successful, None otherwise.
-        """
-        if not self.current_trace:
-            return None
-
-        try:
-            span_id = f"{node_id}_{phase}"
-
-            # Create a child span using v2 API
-            span = self.current_trace.span(
-                name=f"{node_name}.{phase}",
-                metadata={
-                    "node_type": node_name,
-                    "node_id": node_id,
-                    "phase": phase,
-                    "start_timestamp": datetime.now().isoformat(),
-                },
-            )
-
-            self.spans[span_id] = span
-
-            if self.config.debug:
-                print(f"✓ Started span: {span_id}")
-
-            return span_id
-
-        except Exception as e:
-            if self.config.debug:
-                print(f"✗ Failed to start span: {e}")
-            return None
-
-    def end_node_span(
-        self,
-        span_id: str,
-        input_data: Any = None,
-        output_data: Any = None,
-        error: Exception = None,
-    ) -> None:
-        """
-        End a node execution span.
-
-        Args:
-            span_id: ID of the span to end.
-            input_data: Input data for the phase.
-            output_data: Output data from the phase.
-            error: Exception if the phase failed.
-        """
-        if span_id not in self.spans:
-            return
-
-        try:
-            span = self.spans[span_id]
-
-            # Prepare update data
-            update_data = {}
-
-            if input_data is not None and self.config.trace_inputs:
-                update_data["input"] = self._serialize_data(input_data)
-            if output_data is not None and self.config.trace_outputs:
-                update_data["output"] = self._serialize_data(output_data)
-
-            if error and self.config.trace_errors:
-                update_data.update(
-                    {
-                        "level": "ERROR",
-                        "status_message": str(error),
-                        "metadata": {
-                            "error_type": type(error).__name__,
-                            "error_message": str(error),
-                            "end_timestamp": datetime.now().isoformat(),
-                        },
-                    }
-                )
-            else:
-                update_data.update(
-                    {
-                        "level": "DEFAULT",
-                        "metadata": {"end_timestamp": datetime.now().isoformat()},
-                    }
-                )
-
-            # Update the span with all data at once
-            span.update(**update_data)
-
-            # End the span
-            span.end()
-
-            if self.config.debug:
-                status = "ERROR" if error else "SUCCESS"
-                print(f"✓ Ended span: {span_id} with status: {status}")
-
-        except Exception as e:
-            if self.config.debug:
-                print(f"✗ Failed to end span: {e}")
-        finally:
-            if span_id in self.spans:
-                del self.spans[span_id]
-
-    def _serialize_data(self, data: Any) -> Any:
-        """
-        Safely serialize data for Langfuse.
-
-        Args:
-            data: Data to serialize.
-
-        Returns:
-            Serialized data that can be sent to Langfuse.
-        """
-        try:
-            # Handle common PocketFlow data types
-            if hasattr(data, "__dict__"):
-                # Convert objects to dict representation
-                return {"_type": type(data).__name__, "_data": str(data)}
-            elif isinstance(data, (dict, list, str, int, float, bool, type(None))):
-                # JSON-serializable types
-                return data
-            else:
-                # Fallback to string representation
-                return {"_type": type(data).__name__, "_data": str(data)}
-        except Exception:
-            # Ultimate fallback
-            return {"_type": "unknown", "_data": "<serialization_failed>"}
-
-    def flush(self) -> None:
-        """Flush any pending traces to Langfuse."""
-        if self.client:
-            try:
-                self.client.flush()
-                if self.config.debug:
-                    print("✓ Flushed traces to Langfuse")
-            except Exception as e:
-                if self.config.debug:
-                    print(f"✗ Failed to flush traces: {e}")

agentic_blocks-0.1.22/src/agentic_blocks/tracing/decorator.py

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