add tracing cookbook
This commit is contained in:
@@ -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
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user