Spaces:
Sleeping
Sleeping
File size: 4,218 Bytes
c0cb280 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | import datetime
import functools
import logging
from advanced_workflow_orchestrator import AdvancedWorkflowOrchestrator
from analytics.collector import AsyncAnalyticsCollector
logger = logging.getLogger(__name__)
def activate_analytics(orchestrator_instance: AdvancedWorkflowOrchestrator):
"""
Monkey-patches the AdvancedWorkflowOrchestrator's execution method
to automatically record execution metrics.
"""
logger.info("🧬 Activating Workflow DNA (Analytics Instrumentation) for AdvancedWorkflowOrchestrator...")
# Capture the original method
# Note: simple assignment like `original = inst.method` captures a bound method.
# We need to be careful not to create infinite recursion.
original_method = orchestrator_instance._execute_workflow_step
# Create the wrapper
async def instrumented_execute_step(self, workflow, step_id, context):
start_time = datetime.datetime.now()
# Resolve step details
step = next((s for s in workflow.steps if s.step_id == step_id), None)
step_type = step.step_type.value if step and hasattr(step.step_type, 'value') else "unknown"
if step_type == "unknown" and step:
step_type = str(step.step_type)
status = "COMPLETED"
error = None
try:
# Call original method
# Since original_method is already bound to the instance (if we captured it from instance),
# we might not need to pass 'self' again if we just call it?
# actually, if we replace the method on the instance, 'self' will be passed
# to our wrapper.
# But 'original_method' is the OLD bound method.
# So we should call `original_method(workflow, step_id, context)` directly depending on how it was captured.
# If we captured `inst._execute_workflow_step`, it IS a bound method.
# So we don't pass self.
await original_method(workflow, step_id, context)
# Check context for status (it might have failed inside)
result = context.results.get(step_id)
if result and result.get("status") == "failed":
status = "FAILED"
error = result.get("error")
except Exception as e:
status = "FAILED"
error = str(e)
raise e
finally:
end_time = datetime.datetime.now()
# Log to sidecar
# MAPPING FIX:
# 1. WorkflowContext.workflow_id IS the execution_id (e.g. exec_123)
# 2. Real Workflow Definition ID is usually in input_data OR we treat the execution ID as the workflow ID if missing.
execution_id = getattr(context, 'workflow_id', 'unknown')
# Try to find the Definition ID
workflow_def_id = context.input_data.get("_ui_workflow_id") if context.input_data else None
# Fallback: if we can't find a definition ID, use the execution ID or 'ad-hoc'
if not workflow_def_id:
workflow_def_id = "ad-hoc"
if execution_id != "unknown":
await AsyncAnalyticsCollector.get_instance().log_step(
execution_id=execution_id,
workflow_id=workflow_def_id,
step_id=step_id,
step_type=step_type,
start_time=start_time,
end_time=end_time,
status=status,
error=error,
results=context.results.get(step_id) if hasattr(context, 'results') else None
)
# Apply the patch
# We are replacing a BOUND method on the instance.
# The wrapper function `instrumented_execute_step` expects `self` as first arg.
# We need to bind it to the instance manually or use partial.
orchestrator_instance._execute_workflow_step = functools.partial(instrumented_execute_step, orchestrator_instance)
logger.info("✅ Workflow DNA Active: Instrumentation applied to AdvancedWorkflowOrchestrator.")
|