add tracing cookbook

This commit is contained in:
remy
2025-06-27 12:09:19 +10:00
parent 42523a90f1
commit ba574d530e
17 changed files with 1722 additions and 1 deletions
@@ -0,0 +1,13 @@
"""
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"]
@@ -0,0 +1,111 @@
"""
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
+287
View File
@@ -0,0 +1,287 @@
"""
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}")
@@ -0,0 +1,293 @@
"""
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