Update services/pipeline_executor.py

services/pipeline_executor.py

@@ -479,6 +479,63 @@ def execute_pipeline_crewai_streaming(
 # UNIFIED EXECUTOR WITH FALLBACK
 # ========================
 
+# ========================
+# TOOL REGISTRY & DYNAMIC EXECUTION (UPDATED)
+# ========================
+
+# Import the master tools
+try:
+    from services.master_tools import get_master_tools
+    from langchain_core.tools import StructuredTool
+    
+    # Get all tools from master_tools
+    MASTER_TOOLS = get_master_tools()
+    
+    # Create tool registry mapping
+    TOOL_REGISTRY = {}
+    for tool in MASTER_TOOLS:
+        if hasattr(tool, 'name'):
+            TOOL_REGISTRY[tool.name] = tool
+        elif hasattr(tool, '__name__'):
+            TOOL_REGISTRY[tool.__name__] = tool
+    
+    print(f"✅ Loaded {len(TOOL_REGISTRY)} tools from master_tools.py")
+    
+except ImportError as e:
+    print(f"⚠️ Could not import master_tools: {e}")
+    TOOL_REGISTRY = {}
+
+
+def get_tool_executor(tool_name: str) -> Optional[Any]:
+    """Get tool from registry with intelligent name matching"""
+    # Direct match
+    if tool_name in TOOL_REGISTRY:
+        return TOOL_REGISTRY[tool_name]
+    
+    # Try variations
+    variations = [
+        tool_name,
+        f"{tool_name}_tool",
+        tool_name.replace("_", ""),
+        tool_name + "_tool"
+    ]
+    
+    for variation in variations:
+        if variation in TOOL_REGISTRY:
+            return TOOL_REGISTRY[variation]
+    
+    # Check partial matches
+    for registered_name, tool in TOOL_REGISTRY.items():
+        if tool_name in registered_name or registered_name in tool_name:
+            return tool
+    
+    return None
+
+
+# ========================
+# UNIFIED EXECUTOR WITH FALLBACK (UPDATED)
+# ========================
+
 def execute_pipeline_streaming(
     pipeline: Dict[str, Any],
     file_path: str,
@@ -486,74 +543,474 @@ def execute_pipeline_streaming(
     prefer_bedrock: bool = True
 ) -> Generator[Dict[str, Any], None, None]:
     """
-    Execute pipeline with fallback mechanism.
+    Execute pipeline with fallback mechanism using master_tools.
+    """
+    components_executed = []
+    final_output = None
+    executor_used = "unknown"
+    fallback_triggered = False
+    bedrock_error = None
     
-    Priority:
-    1. Try Bedrock + LangChain - if available
-    2. Fallback to CrewAI - if Bedrock fails
+    # Initialize pipeline info
+    pipeline_id = pipeline.get("pipeline_id")
+    pipeline_name = pipeline.get("pipeline_name", "Unnamed Pipeline")
+    steps = pipeline.get("pipeline_steps", [])
+    
+    yield {
+        "type": "info",
+        "message": f"Starting pipeline: {pipeline_name}",
+        "executor": "initializing"
+    }
+    
+    # Check if tools are available
+    if not TOOL_REGISTRY:
+        error_msg = "No tools available. master_tools.py not loaded correctly."
+        yield {
+            "type": "error",
+            "error": error_msg,
+            "data": {
+                "pipeline_id": pipeline_id,
+                "pipeline_name": pipeline_name,
+                "status": "failed",
+                "components_executed": [],
+                "error": error_msg
+            }
+        }
+        return
     
-    Yields:
-        Status updates and final results
-    """
     # Try Bedrock first (priority)
     if prefer_bedrock and BEDROCK_AVAILABLE:
         try:
-            print(f"🏆 Executing pipeline with Bedrock: {pipeline['pipeline_name']}")
+            print(f"🏆 Executing pipeline with Bedrock: {pipeline_name}")
             yield {
                 "type": "info",
                 "message": "Attempting execution with Bedrock LangChain...",
                 "executor": "bedrock"
             }
             
-            # Try to execute with Bedrock
-            error_occurred = False
-            for event in execute_pipeline_bedrock_streaming(pipeline, file_path, session_id):
-                yield event
+            # Execute step by step with Bedrock
+            for step_num, step_def in enumerate(steps, 1):
+                tool_name = step_def.get("tool_name", "unknown")
                 
-                # Check if error occurred
-                if event.get("type") == "error":
-                    error_occurred = True
-                    bedrock_error = event.get("error")
-                    print(f"❌ Bedrock execution failed: {bedrock_error}")
-                    print("🔄 Falling back to CrewAI...")
+                yield {
+                    "type": "step",
+                    "step": step_num,
+                    "tool": tool_name,
+                    "status": "executing",
+                    "executor": "bedrock"
+                }
+                
+                try:
+                    # Execute the step using master_tools
+                    result = _execute_step_with_master_tool(
+                        step_def=step_def,
+                        file_path=file_path,
+                        step_num=step_num,
+                        total_steps=len(steps),
+                        session_id=session_id,
+                        prefer_bedrock=True
+                    )
+                    
+                    executor_used = "bedrock"
+                    
+                    # Create component result object
+                    component_result = {
+                        **step_def,
+                        "result": result.get("output"),
+                        "status": "completed",
+                        "executor": executor_used,
+                        "execution_time": result.get("execution_time"),
+                        "step_number": step_num,
+                        "success": True,
+                        "tool_version": result.get("tool_version", "1.0")
+                    }
+                    
+                    components_executed.append(component_result)
                     
                     yield {
-                        "type": "info",
-                        "message": f"Bedrock failed: {bedrock_error}. Switching to CrewAI...",
-                        "executor": "fallback"
+                        "type": "step",
+                        "step": step_num,
+                        "tool": tool_name,
+                        "status": "completed",
+                        "observation": result.get("output"),
+                        "input": step_def,
+                        "executor": executor_used
+                    }
+                    
+                    # Update file_path for next step if needed
+                    file_path = _update_file_path(file_path, result)
+                    
+                except Exception as step_error:
+                    # Create failed component result
+                    component_result = {
+                        **step_def,
+                        "result": {"error": str(step_error)},
+                        "status": "failed",
+                        "error": str(step_error),
+                        "step_number": step_num,
+                        "success": False
+                    }
+                    
+                    components_executed.append(component_result)
+                    bedrock_error = str(step_error)
+                    
+                    yield {
+                        "type": "error",
+                        "step": step_num,
+                        "tool": tool_name,
+                        "error": str(step_error),
+                        "message": f"Step {step_num} failed with Bedrock"
                     }
+                    
+                    fallback_triggered = True
                     break
-                
-                # If final result, we're done
-                if event.get("type") == "final":
-                    print(f"✅ Bedrock execution completed: {pipeline['pipeline_name']}")
-                    return
             
-            # If we got here with error, fall back to CrewAI
-            if error_occurred:
-                # Fall through to CrewAI
-                pass
-            else:
-                # Successful completion (shouldn't reach here normally)
+            # If we completed all steps with Bedrock
+            if not fallback_triggered and len(components_executed) == len(steps):
+                final_output = _build_final_output(pipeline, components_executed, executor_used, "completed")
+                
+                yield {
+                    "type": "final",
+                    "data": final_output,
+                    "executor": executor_used
+                }
+                print(f"✅ Bedrock execution completed: {pipeline_name}")
                 return
                 
-        except Exception as bedrock_error:
-            print(f"❌ Bedrock execution exception: {str(bedrock_error)}")
-            print("🔄 Falling back to CrewAI...")
+        except Exception as bedrock_exception:
+            print(f"❌ Bedrock execution exception: {str(bedrock_exception)}")
+            bedrock_error = str(bedrock_exception)
+            fallback_triggered = True
+    
+    # If Bedrock failed or wasn't preferred, try CrewAI
+    if fallback_triggered or not prefer_bedrock:
+        print(f"🔄 Executing pipeline with CrewAI: {pipeline_name}")
+        
+        if fallback_triggered and bedrock_error:
             yield {
                 "type": "info",
-                "message": f"Bedrock exception: {str(bedrock_error)}. Switching to CrewAI...",
+                "message": f"Bedrock failed: {bedrock_error}. Switching to CrewAI...",
                 "executor": "fallback"
             }
+        else:
+            yield {
+                "type": "info",
+                "message": "Using CrewAI execution...",
+                "executor": "crewai"
+            }
+        
+        # Start from where Bedrock left off, or from beginning
+        start_step = len(components_executed) + 1 if components_executed else 1
+        
+        for step_num in range(start_step, len(steps) + 1):
+            step_def = steps[step_num - 1]
+            tool_name = step_def.get("tool_name", "unknown")
+            
+            yield {
+                "type": "step",
+                "step": step_num,
+                "tool": tool_name,
+                "status": "executing",
+                "executor": "crewai"
+            }
+            
+            try:
+                # Execute the step using master_tools
+                result = _execute_step_with_master_tool(
+                    step_def=step_def,
+                    file_path=file_path,
+                    step_num=step_num,
+                    total_steps=len(steps),
+                    session_id=session_id,
+                    prefer_bedrock=False
+                )
+                
+                executor_used = "crewai"
+                
+                # Create component result object
+                component_result = {
+                    **step_def,
+                    "result": result.get("output"),
+                    "status": "completed",
+                    "executor": executor_used,
+                    "execution_time": result.get("execution_time"),
+                    "step_number": step_num,
+                    "success": True,
+                    "tool_version": result.get("tool_version", "1.0")
+                }
+                
+                # Add or replace in components_executed
+                if len(components_executed) >= step_num:
+                    components_executed[step_num - 1] = component_result
+                else:
+                    components_executed.append(component_result)
+                
+                yield {
+                    "type": "step",
+                    "step": step_num,
+                    "tool": tool_name,
+                    "status": "completed",
+                    "observation": result.get("output"),
+                    "input": step_def,
+                    "executor": executor_used
+                }
+                
+                # Update file_path for next step if needed
+                file_path = _update_file_path(file_path, result)
+                    
+            except Exception as step_error:
+                # Create failed component result
+                component_result = {
+                    **step_def,
+                    "result": {"error": str(step_error)},
+                    "status": "failed",
+                    "error": str(step_error),
+                    "step_number": step_num,
+                    "success": False
+                }
+                
+                # Add or replace in components_executed
+                if len(components_executed) >= step_num:
+                    components_executed[step_num - 1] = component_result
+                else:
+                    components_executed.append(component_result)
+                
+                yield {
+                    "type": "error",
+                    "step": step_num,
+                    "tool": tool_name,
+                    "error": str(step_error),
+                    "message": f"Step {step_num} failed with CrewAI"
+                }
+                break
+        
+        # Check if we completed all steps
+        completed_steps = [c for c in components_executed if c.get("status") == "completed"]
+        
+        if len(completed_steps) == len(steps):
+            # All steps completed
+            final_output = _build_final_output(pipeline, components_executed, executor_used, "completed")
+            
+            yield {
+                "type": "final",
+                "data": final_output,
+                "executor": executor_used
+            }
+            print(f"✅ CrewAI execution completed: {pipeline_name}")
+        else:
+            # Partial completion or failure
+            final_output = _build_final_output(pipeline, components_executed, executor_used, "partial")
+            final_output["error"] = f"Pipeline execution incomplete. Completed {len(completed_steps)} of {len(steps)} steps."
+            
+            yield {
+                "type": "error",
+                "error": "Pipeline execution incomplete",
+                "data": final_output
+            }
+            print(f"⚠️ CrewAI execution incomplete for: {pipeline_name}")
+
+
+# ========================
+# DYNAMIC STEP EXECUTION WITH MASTER_TOOLS
+# ========================
+
+def _execute_step_with_master_tool(
+    step_def: Dict[str, Any],
+    file_path: str,
+    step_num: int,
+    total_steps: int,
+    session_id: Optional[str] = None,
+    prefer_bedrock: bool = True
+) -> Dict[str, Any]:
+    """
+    Execute a pipeline step using master_tools.
+    """
+    import time
+    import inspect
+    
+    tool_name = step_def.get("tool_name", "unknown")
+    start_time = time.time()
     
-    # Fallback to CrewAI
-    print(f"🔄 Executing pipeline with CrewAI: {pipeline['pipeline_name']}")
-    for event in execute_pipeline_crewai_streaming(pipeline, file_path, session_id):
-        yield event
+    # Get tool from registry
+    tool = get_tool_executor(tool_name)
+    
+    if not tool:
+        raise ValueError(f"Tool '{tool_name}' not found in registry. Available tools: {list(TOOL_REGISTRY.keys())}")
+    
+    # Prepare arguments
+    args = {}
+    
+    # For StructuredTool (LangChain tools)
+    if hasattr(tool, 'args_schema') and hasattr(tool, 'invoke'):
+        # Get the args schema
+        args_schema = tool.args_schema
         
-        if event.get("type") == "final":
-            print(f"✅ CrewAI execution completed: {pipeline['pipeline_name']}")
-            return
+        # Build arguments from step_def
+        for field_name, field in args_schema.__fields__.items():
+            # Check if parameter is in step_def
+            if field_name in step_def:
+                args[field_name] = step_def[field_name]
+            # Special handling for file_path
+            elif field_name == "file_path" and file_path:
+                args[field_name] = file_path
+            # Special handling for session_id
+            elif field_name == "session_id" and session_id:
+                args[field_name] = session_id
+            # Handle text parameter if not provided but we have previous output
+            elif field_name == "text" and field_name not in step_def and step_num > 1:
+                # Try to get text from previous step's output
+                if components_executed and len(components_executed) >= step_num - 1:
+                    prev_result = components_executed[step_num - 2].get("result")
+                    if isinstance(prev_result, dict) and "text" in prev_result:
+                        args["text"] = prev_result["text"]
+        
+        try:
+            # Execute the tool
+            output = tool.invoke(args)
+            execution_time = time.time() - start_time
+            
+            return {
+                "output": output,
+                "executor": "bedrock" if prefer_bedrock else "crewai",
+                "execution_time": execution_time,
+                "tool_version": "master_tools_1.0",
+                "args_used": list(args.keys())
+            }
+            
+        except Exception as e:
+            # Try with minimal arguments
+            print(f"⚠️ Tool {tool_name} failed with full args, trying minimal: {e}")
+            
+            # Try with just file_path if available
+            if file_path and "file_path" in args_schema.__fields__:
+                minimal_args = {"file_path": file_path}
+                try:
+                    output = tool.invoke(minimal_args)
+                    execution_time = time.time() - start_time
+                    
+                    return {
+                        "output": output,
+                        "executor": "bedrock" if prefer_bedrock else "crewai",
+                        "execution_time": execution_time,
+                        "tool_version": "master_tools_1.0",
+                        "args_used": list(minimal_args.keys()),
+                        "warning": "Used minimal arguments"
+                    }
+                except Exception as inner_error:
+                    raise RuntimeError(f"Tool '{tool_name}' failed with minimal args: {inner_error}")
+    
+    # For regular Python functions
+    elif callable(tool):
+        try:
+            # Get function signature
+            sig = inspect.signature(tool)
+            
+            # Build arguments based on signature
+            call_args = {}
+            for param_name, param in sig.parameters.items():
+                # Try step_def first
+                if param_name in step_def:
+                    call_args[param_name] = step_def[param_name]
+                # Special handling for file_path
+                elif param_name == "file_path" and file_path:
+                    call_args[param_name] = file_path
+                # Special handling for session_id
+                elif param_name == "session_id" and session_id:
+                    call_args[param_name] = session_id
+                # Handle text parameter
+                elif param_name == "text" and param_name not in step_def and step_num > 1:
+                    # Try to get text from previous step
+                    if components_executed and len(components_executed) >= step_num - 1:
+                        prev_result = components_executed[step_num - 2].get("result")
+                        if isinstance(prev_result, dict) and "text" in prev_result:
+                            call_args["text"] = prev_result["text"]
+            
+            # Execute the function
+            output = tool(**call_args)
+            execution_time = time.time() - start_time
+            
+            return {
+                "output": output,
+                "executor": "bedrock" if prefer_bedrock else "crewai",
+                "execution_time": execution_time,
+                "tool_version": "function_1.0",
+                "args_used": list(call_args.keys())
+            }
+            
+        except Exception as e:
+            raise RuntimeError(f"Failed to execute function {tool_name}: {e}")
+    
+    else:
+        raise ValueError(f"Tool '{tool_name}' is not callable or a valid StructuredTool")
+
+
+def _update_file_path(current_file_path: str, result: Dict[str, Any]) -> str:
+    """
+    Update file path based on tool result.
+    Some tools might generate new files.
+    """
+    output = result.get("output")
+    
+    if isinstance(output, dict):
+        # Check for file references in output
+        for key in ["file_path", "output_file", "new_file", "generated_file"]:
+            if key in output and isinstance(output[key], str):
+                return output[key]
+    
+    return current_file_path
+
+
+# ========================
+# HELPER FUNCTIONS
+# ========================
+
+def _build_final_output(
+    pipeline: Dict[str, Any],
+    components_executed: List[Dict[str, Any]],
+    executor_used: str,
+    status: str
+) -> Dict[str, Any]:
+    """
+    Build final output with components_executed array.
+    """
+    # Find the finalize step result if present
+    final_result = None
+    for component in components_executed:
+        if component.get("tool_name") == "finalize":
+            final_result = component.get("result")
+            break
+    
+    final_output = {
+        "pipeline_id": pipeline.get("pipeline_id"),
+        "pipeline_name": pipeline.get("pipeline_name"),
+        "status": status,
+        "components_executed": components_executed,
+        "executor": executor_used,
+        "summary": f"Pipeline execution {status} with {executor_used}",
+        "total_steps": len(pipeline.get("pipeline_steps", [])),
+        "completed_steps": len([c for c in components_executed if c.get("status") == "completed"]),
+        "final_result": final_result
+    }
+    
+    # Extract text for user-facing output
+    if final_result:
+        # Use finalize tool's output
+        final_output["text"] = final_result
+    elif components_executed:
+        # Find last completed component with text
+        for component in reversed(components_executed):
+            if component.get("status") == "completed" and component.get("result"):
+                result = component["result"]
+                if isinstance(result, str):
+                    final_output["text"] = result
+                    break
+                elif isinstance(result, dict):
+                    for field in ["text", "summary", "content", "translation"]:
+                        if field in result and isinstance(result[field], str):
+                            final_output["text"] = result[field]
+                            break
+    
+    return final_output
 
 
 # ========================
@@ -575,9 +1032,18 @@ def execute_pipeline(
         if event.get("type") == "final":
             final_result = event.get("data")
             break
+        elif event.get("type") == "error" and event.get("data"):
+            final_result = event.get("data")
+            break
     
     if final_result is None:
-        raise RuntimeError("Pipeline execution completed without final result")
+        final_result = {
+            "pipeline_id": pipeline.get("pipeline_id"),
+            "pipeline_name": pipeline.get("pipeline_name"),
+            "status": "failed",
+            "components_executed": [],
+            "error": "Pipeline execution completed without final result"
+        }
     
     return final_result