""" R12 Service: Workflow Pruner. Reduces workflow JSON size by tracing relevant upstream nodes from the error point. Implements deterministic traversal for consistent testing and A/B validation. """ import logging import json from dataclasses import dataclass, field from typing import Dict, Any, Optional, List, Set, Literal logger = logging.getLogger(__name__) @dataclass class PruneConfig: """Configuration for workflow pruning.""" max_depth: int = 3 max_nodes: int = 40 min_nodes_floor: int = 10 allow_cycle_break: bool = True fallback_recent_nodes_count: int = 20 @dataclass class PruneResult: """Result of workflow pruning.""" pruned_workflow_json: Dict[str, Any] kept_node_ids: List[str] dropped_nodes_count: int depth_used: int nodes_used: int mode: Literal["upstream_trace", "fallback_recent_nodes", "metadata_only", "full_pass"] reason: str class WorkflowPruner: """ Intelligent workflow pruning service. """ def __init__(self): pass def prune(self, workflow: Dict[str, Any], target_node_id: Optional[str], config: Optional[PruneConfig] = None) -> PruneResult: """ Prune the workflow to relevant subgraph ending at target_node_id. Args: workflow: The full ComfyUI workflow JSON (API format or UI format). Expected to be API format (dict of node_id -> node_data) for backend. target_node_id: The ID of the node where error occurred. config: Pruning configuration. Returns: PruneResult containing the reduced workflow. """ config = config or PruneConfig() # 1. Validation & Format Normalization if not workflow or not isinstance(workflow, dict): return self._return_empty_result(workflow, "invalid_format") # Detect prompt format (API format) vs UI format (with "nodes", "links") # Start with API format assumption (node_id -> keys) is_api_format = True nodes_dict = workflow # If it looks like UI format (has "nodes" list), we might need conversion or different handling. # R12 scope focuses on backend payload which is typically the prompt/API format. # Standard ComfyUI API prompt is key-value pairs of node_id -> node_data. # Example: {"3": {"inputs": {...}, "class_type": "KSampler"}} # Check if basic validation holds for API format if "nodes" in workflow and isinstance(workflow["nodes"], list): # It's UI format (saved workflow) # We treat this as "metadata_only" fallback for now or extraction # But usually doctor receives the executing prompt which is API format. # If we receive UI format, we might fallback to recent nodes or just return as is if small. # Let's assume API format for the primary logic as per "node_context" usage. is_api_format = False if not is_api_format: # Fallback for UI format (not structured for upstream tracing easily without conversion) # Just truncate if too large? # For MVP, we'll try to keep it simple. If not API format, return as is or minimal fallback. return self._fallback_pruning(workflow, target_node_id, config, "unsupported_format_fallback") # 2. Target Node Validation if not target_node_id or str(target_node_id) not in nodes_dict: return self._fallback_pruning(workflow, target_node_id, config, "target_node_missing") # 3. Upstream Tracing (BFS/DFS) try: kept_nodes, depth = self._trace_upstream(nodes_dict, str(target_node_id), config) # Construct pruned workflow pruned_workflow = {nid: nodes_dict[nid] for nid in kept_nodes} dropped_count = len(nodes_dict) - len(kept_nodes) return PruneResult( pruned_workflow_json=pruned_workflow, kept_node_ids=sorted(list(kept_nodes)), # Deterministic sort for metadata dropped_nodes_count=dropped_count, depth_used=depth, nodes_used=len(kept_nodes), mode="upstream_trace", reason="success" ) except Exception as e: logger.error(f"Pruning failed: {e}") return self._fallback_pruning(workflow, target_node_id, config, f"error: {str(e)}") def _trace_upstream(self, nodes: Dict[str, Any], start_node: str, config: PruneConfig) -> tuple[Set[str], int]: """ Trace upstream nodes starting from start_node. Returns (set of kept node_ids, max_depth_reached). """ visited = set() queue = [(start_node, 0)] # (node_id, depth) visited.add(start_node) kept_nodes = {start_node} max_depth_reached = 0 # Deterministic traversal: # We use a queue, but when we expand inputs, we must do so in a deterministic order. while queue: # Check node limit if len(kept_nodes) >= config.max_nodes: break current_id, current_depth = queue.pop(0) max_depth_reached = max(max_depth_reached, current_depth) if current_depth >= config.max_depth: continue node_data = nodes.get(current_id, {}) inputs = node_data.get("inputs", {}) # Find upstream dependencies # In API format, inputs usually look like: "key": ["other_node_id", 0] dependencies = [] # Sort input keys for deterministic traversal ordering sorted_keys = sorted(inputs.keys()) for key in sorted_keys: val = inputs[key] if isinstance(val, list) and len(val) == 2: # Potential link: [node_id, output_index] # Check if val[0] is a node_id in our graph upstream_id = str(val[0]) if upstream_id in nodes: dependencies.append(upstream_id) # Add unique dependencies to queue for dep_id in dependencies: if dep_id not in visited: if len(kept_nodes) < config.max_nodes: visited.add(dep_id) kept_nodes.add(dep_id) queue.append((dep_id, current_depth + 1)) return kept_nodes, max_depth_reached def _fallback_pruning(self, workflow: Dict[str, Any], target_id: Optional[str], config: PruneConfig, reason: str) -> PruneResult: """ Fallback strategy when tracing fails or target missing. Keeps the last N nodes based on some heuristic (ID sorting or insertion order). """ if not isinstance(workflow, dict): return self._return_empty_result(workflow, "invalid_type") # Heuristic: Sort by Node ID (assuming numeric IDs usually imply creation order) # fallback_recent_nodes_count try: # Try to sort keys numerically if possible, else alphanumeric sorted_ids = sorted(workflow.keys(), key=lambda x: int(x) if x.isdigit() else x) # Keep last N keep_ids = sorted_ids[-config.fallback_recent_nodes_count:] pruned_workflow = {nid: workflow[nid] for nid in keep_ids if nid in workflow} return PruneResult( pruned_workflow_json=pruned_workflow, kept_node_ids=keep_ids, dropped_nodes_count=len(workflow) - len(pruned_workflow), depth_used=0, nodes_used=len(pruned_workflow), mode="fallback_recent_nodes", reason=reason ) except Exception: # Absolute fallback: return as is or empty return PruneResult( pruned_workflow_json=workflow, # Return full as safe fallback? Or empty? # A safer budget-conscious fallback might be to return minimal info, # but to be safe for now, let's just return what we have and let budget truncator handle text. kept_node_ids=list(workflow.keys()), dropped_nodes_count=0, depth_used=0, nodes_used=len(workflow), mode="full_pass", reason=f"{reason}_fatal" ) def _return_empty_result(self, original_workflow: Any, reason: str) -> PruneResult: return PruneResult( pruned_workflow_json={}, kept_node_ids=[], dropped_nodes_count=len(original_workflow) if isinstance(original_workflow, dict) else 0, depth_used=0, nodes_used=0, mode="metadata_only", reason=reason ) def estimate_tokens(self, data: Any) -> int: """Estimate token count for JSON data (legacy helper).""" text = json.dumps(data) # Rough estimation: 1 token ~= 4 chars return len(text) // 4