genxai-framework 0.1.0__py3-none-any.whl

genxai/core/graph/executor.py

@@ -0,0 +1,516 @@
+"""Workflow execution engine for GenXAI."""
+
+import asyncio
+import copy
+from typing import Any, Dict, List, Optional
+import logging
+from pathlib import Path
+
+from genxai.core.graph.engine import Graph
+from genxai.core.graph.nodes import InputNode, OutputNode, AgentNode, NodeType
+from genxai.core.graph.edges import Edge, ConditionalEdge
+from genxai.core.agent.base import Agent, AgentFactory
+from genxai.core.agent.registry import AgentRegistry
+from genxai.tools.registry import ToolRegistry
+from genxai.core.execution import WorkerQueueEngine, ExecutionStore
+from genxai.tools.builtin.computation.calculator import CalculatorTool
+from genxai.tools.builtin.file.file_reader import FileReaderTool
+from genxai.security.rbac import get_current_user, Permission
+from genxai.security.policy_engine import get_policy_engine
+from genxai.security.audit import get_audit_log, AuditEvent
+
+logger = logging.getLogger(__name__)
+
+
+class EnhancedGraph(Graph):
+    """Enhanced graph with agent execution support.
+    
+    This extends the base Graph class to provide real agent execution
+    with tool integration. It's the recommended way to execute workflows
+    in GenXAI.
+    """
+
+    async def _execute_node_logic(self, node: Any, state: Dict[str, Any]) -> Any:
+        """Execute node logic with actual agent execution.
+
+        Args:
+            node: Node to execute
+            state: Current state
+
+        Returns:
+            Result of node execution
+        """
+        if node.type == NodeType.INPUT:
+            # IMPORTANT: Avoid returning the exact same `dict` object stored under
+            # `state["input"]`. If we do, the engine will store that same object
+            # under the input node id (e.g. "start"), creating shared references
+            # which Python's `json.dumps` treats as circular.
+            return copy.deepcopy(state.get("input"))
+        
+        elif node.type == NodeType.OUTPUT:
+            # IMPORTANT: never return the live `state` dict. The engine stores the
+            # node result back into `state[node_id]`, so returning `state` would
+            # create a self-referential structure (circular reference) that can't
+            # be JSON-serialized for persistence.
+            # Also deep-copy to avoid shared references (json can't encode those).
+            return copy.deepcopy(state)
+        
+        elif node.type == NodeType.AGENT:
+            # Get agent from registry
+            agent_id = node.config.data.get("agent_id")
+            if not agent_id:
+                raise ValueError(f"Agent node '{node.id}' missing agent_id in config.data")
+            
+            agent = AgentRegistry.get(agent_id)
+            if agent is None:
+                raise ValueError(f"Agent '{agent_id}' not found in registry")
+            
+            # Prepare task from state
+            task = state.get("task", "Process the input data")
+            
+            # Execute agent with tools if available
+            result = await self._execute_agent_with_tools(agent, task, state)
+            
+            return result
+        
+        else:
+            # Default behavior
+            return {"node_id": node.id, "type": node.type.value}
+
+    async def _execute_agent_with_tools(
+        self, agent: Agent, task: str, state: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """Execute agent with tool support using AgentRuntime.
+
+        Args:
+            agent: Agent to execute
+            task: Task description
+            state: Current state
+
+        Returns:
+            Execution result
+        """
+        logger.debug(f"Executing agent '{agent.id}' ({agent.config.role})")
+        
+        # Use AgentRuntime for full integration
+        from genxai.core.agent.runtime import AgentRuntime
+        
+        runtime = AgentRuntime(agent=agent, enable_memory=True)
+        
+        # Load tools from registry
+        if agent.config.tools:
+            tools = {}
+            for tool_name in agent.config.tools:
+                tool = ToolRegistry.get(tool_name)
+                if tool:
+                    tools[tool_name] = tool
+            runtime.set_tools(tools)
+            logger.debug(f"Loaded {len(tools)} tools for agent")
+        
+        # Execute agent with full runtime support
+        result = await runtime.execute(task, context=state)
+        
+        return result
+
+    async def _execute_tool_for_task(
+        self, tool: Any, tool_name: str, task: str, state: Dict[str, Any]
+    ) -> Any:
+        """Execute a tool based on the task.
+
+        Args:
+            tool: Tool instance
+            tool_name: Tool name
+            task: Task description
+            state: Current state
+
+        Returns:
+            Tool execution result or None
+        """
+        try:
+            # Calculator tool
+            if tool_name == "calculator":
+                # Check if task involves calculation
+                if any(op in task.lower() for op in ["calculate", "compute", "add", "multiply", "+"]):
+                    # Extract expression from state or use default
+                    expression = state.get("expression", "10 * 5 + 3")
+                    logger.debug(f"Executing calculator: {expression}")
+                    result = await tool.execute(expression=expression)
+                    if result.success:
+                        logger.debug(f"Calculator result: {result.data['result']}")
+                        return result.data
+            
+            # File reader tool
+            elif tool_name == "file_reader":
+                # Check if task involves file reading
+                if any(word in task.lower() for word in ["read", "file", "load"]):
+                    # Get file path from state
+                    file_path = state.get("file_path")
+                    
+                    if file_path:
+                        logger.debug(f"Reading file: {file_path}")
+                        result = await tool.execute(path=file_path)
+                        if result.success:
+                            logger.debug(f"Read {result.data['lines']} lines")
+                            return result.data
+        
+        except Exception as e:
+            logger.error(f"Tool execution failed: {e}")
+        
+        return None
+
+
+class WorkflowExecutor:
+    """Executes workflows using GenXAI engine.
+    
+    This is the main class for executing workflows programmatically.
+    It handles agent creation, tool registration, graph building,
+    and execution.
+    
+    Example:
+        ```python
+        executor = WorkflowExecutor(openai_api_key="sk-...")
+        result = await executor.execute(nodes, edges, input_data)
+        ```
+    """
+
+    def __init__(
+        self, 
+        openai_api_key: Optional[str] = None, 
+        anthropic_api_key: Optional[str] = None,
+        register_builtin_tools: bool = True,
+        queue_engine: Optional[WorkerQueueEngine] = None,
+        execution_store: Optional[ExecutionStore] = None,
+    ):
+        """Initialize workflow executor.
+
+        Args:
+            openai_api_key: OpenAI API key
+            anthropic_api_key: Anthropic API key
+            register_builtin_tools: Whether to register built-in tools
+        """
+        self.openai_api_key = openai_api_key
+        self.anthropic_api_key = anthropic_api_key
+        
+        self.queue_engine = queue_engine
+        self.execution_store = execution_store or ExecutionStore()
+
+        if register_builtin_tools:
+            self._setup_tools()
+
+    def _setup_tools(self) -> None:
+        """Register built-in tools."""
+        # Register calculator tool
+        if not ToolRegistry.get("calculator"):
+            calculator = CalculatorTool()
+            ToolRegistry.register(calculator)
+            logger.info("Registered calculator tool")
+
+        # Register file reader tool
+        if not ToolRegistry.get("file_reader"):
+            file_reader = FileReaderTool()
+            ToolRegistry.register(file_reader)
+            logger.info("Registered file_reader tool")
+
+    def _create_agents_from_nodes(self, nodes: List[Dict[str, Any]]) -> None:
+        """Create and register agents from workflow nodes.
+
+        Args:
+            nodes: List of workflow nodes
+        """
+        for node in nodes:
+            if node.get("type") == "agent":
+                agent_id = node.get("id")
+                config = node.get("config", {})
+                
+                # Extract agent configuration
+                role = config.get("role", "Agent")
+                goal = config.get("goal", "Process tasks")
+                backstory = config.get("backstory", "")
+                tools = config.get("tools", [])
+                llm_model = config.get("llm_model", "gpt-4")
+                temperature = config.get("temperature", 0.7)
+
+                # Create agent
+                agent = AgentFactory.create_agent(
+                    id=agent_id,
+                    role=role,
+                    goal=goal,
+                    backstory=backstory,
+                    tools=tools,
+                    llm_model=llm_model,
+                    temperature=temperature,
+                )
+
+                # Register agent
+                AgentRegistry.register(agent)
+                logger.info(f"Created and registered agent: {agent_id}")
+
+    def _build_graph(self, nodes: List[Dict[str, Any]], edges: List[Dict[str, Any]]) -> EnhancedGraph:
+        """Build GenXAI graph from workflow definition.
+
+        Args:
+            nodes: List of workflow nodes
+            edges: List of workflow edges
+
+        Returns:
+            Constructed graph
+        """
+        graph = EnhancedGraph(name="workflow")
+
+        # Add nodes
+        for node in nodes:
+            node_id = node.get("id")
+            node_type = node.get("type")
+
+            # Support some common aliases used by the Studio UI
+            # - "start" behaves like an input node
+            # - "end" behaves like an output node
+            if node_type in {"input", "start"}:
+                graph.add_node(InputNode(id=node_id))
+            elif node_type in {"output", "end"}:
+                graph.add_node(OutputNode(id=node_id))
+            elif node_type == "agent":
+                graph.add_node(AgentNode(id=node_id, agent_id=node_id))
+            else:
+                logger.warning(f"Unknown node type: {node_type}")
+
+        # Add edges
+        for edge in edges:
+            source = edge.get("source")
+            target = edge.get("target")
+            condition = edge.get("condition")
+
+            if condition:
+                # Conditional edge
+                graph.add_edge(
+                    ConditionalEdge(
+                        source=source,
+                        target=target,
+                        condition=lambda state, cond=condition: self._evaluate_condition(state, cond),
+                    )
+                )
+            else:
+                # Regular edge
+                graph.add_edge(Edge(source=source, target=target))
+
+        return graph
+
+    def _evaluate_condition(self, state: Dict[str, Any], condition: str) -> bool:
+        """Evaluate a condition string.
+
+        Args:
+            state: Current workflow state
+            condition: Condition expression
+
+        Returns:
+            Boolean result
+        """
+        # Simple condition evaluation (can be enhanced)
+        try:
+            # For now, just check if condition key exists in state
+            return condition in state
+        except Exception as e:
+            logger.error(f"Error evaluating condition: {e}")
+            return False
+
+    async def execute(
+        self, 
+        nodes: List[Dict[str, Any]], 
+        edges: List[Dict[str, Any]], 
+        input_data: Dict[str, Any],
+        run_id: Optional[str] = None,
+        checkpoint_dir: Optional[str] = None,
+        resume_from: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """Execute a workflow.
+
+        Args:
+            nodes: Workflow nodes
+            edges: Workflow edges
+            input_data: Input data for execution
+
+        Returns:
+            Execution result with status, result, and metadata
+        """
+        run_id = run_id or self.execution_store.generate_run_id()
+        record = self.execution_store.create(run_id, workflow="workflow", status="running")
+
+        try:
+            logger.info("Starting workflow execution")
+
+            # Create agents from nodes
+            self._create_agents_from_nodes(nodes)
+
+            # Build graph
+            graph = self._build_graph(nodes, edges)
+
+            # Validate graph
+            graph.validate()
+            logger.info(f"Graph validated: {len(graph.nodes)} nodes, {len(graph.edges)} edges")
+
+            checkpoint = None
+            if resume_from and checkpoint_dir:
+                checkpoint = graph.load_checkpoint(resume_from, Path(checkpoint_dir))
+
+            # Execute graph
+            user = get_current_user()
+            if user is not None:
+                get_policy_engine().check(user, "workflow:workflow", Permission.WORKFLOW_EXECUTE)
+                get_audit_log().record(
+                    AuditEvent(
+                        action="workflow.execute",
+                        actor_id=user.user_id,
+                        resource_id="workflow:workflow",
+                        status="allowed",
+                    )
+                )
+            result = await graph.run(input_data=input_data, resume_from=checkpoint)
+
+            logger.info("Workflow execution completed successfully")
+
+            self.execution_store.update(
+                run_id,
+                status="success",
+                result=result,
+                completed=True,
+            )
+            return {
+                "status": "success",
+                "run_id": run_id,
+                "result": result,
+                "nodes_executed": len(graph.nodes),
+                "message": "Workflow executed successfully"
+            }
+
+        except Exception as e:
+            logger.error(f"Workflow execution failed: {e}", exc_info=True)
+            self.execution_store.update(
+                run_id,
+                status="error",
+                error=str(e),
+                completed=True,
+            )
+            return {
+                "status": "error",
+                "run_id": run_id,
+                "error": str(e),
+                "message": f"Workflow execution failed: {str(e)}"
+            }
+
+        finally:
+            # Cleanup: Clear registries for next execution
+            AgentRegistry.clear()
+            logger.info("Cleared agent registry")
+
+    async def execute_queued(
+        self,
+        nodes: List[Dict[str, Any]],
+        edges: List[Dict[str, Any]],
+        input_data: Dict[str, Any],
+        run_id: Optional[str] = None,
+        checkpoint_dir: Optional[str] = None,
+        resume_from: Optional[str] = None,
+    ) -> str:
+        """Enqueue workflow execution using a worker queue engine."""
+        if not self.queue_engine:
+            self.queue_engine = WorkerQueueEngine()
+
+        run_id = run_id or self.execution_store.generate_run_id()
+        existing = self.execution_store.get(run_id)
+        if existing and existing.status in {"running", "success"}:
+            return run_id
+
+        self.execution_store.create(run_id, workflow="workflow", status="queued")
+
+        async def _handler(payload: Dict[str, Any]) -> None:
+            await self.execute(
+                nodes=payload["nodes"],
+                edges=payload["edges"],
+                input_data=payload["input_data"],
+                run_id=payload["run_id"],
+                checkpoint_dir=payload.get("checkpoint_dir"),
+                resume_from=payload.get("resume_from"),
+            )
+
+        await self.queue_engine.start()
+        return await self.queue_engine.enqueue(
+            {
+                "nodes": nodes,
+                "edges": edges,
+                "input_data": input_data,
+                "run_id": run_id,
+                "checkpoint_dir": checkpoint_dir,
+                "resume_from": resume_from,
+            },
+            _handler,
+            metadata={"workflow": "queued"},
+            run_id=run_id,
+        )
+
+
+def execute_workflow_sync(
+    nodes: List[Dict[str, Any]],
+    edges: List[Dict[str, Any]],
+    input_data: Dict[str, Any],
+    openai_api_key: Optional[str] = None,
+    anthropic_api_key: Optional[str] = None,
+) -> Dict[str, Any]:
+    """Synchronous wrapper for workflow execution.
+    
+    This is a convenience function for executing workflows in
+    synchronous contexts.
+
+    Args:
+        nodes: Workflow nodes
+        edges: Workflow edges
+        input_data: Input data
+        openai_api_key: OpenAI API key
+        anthropic_api_key: Anthropic API key
+
+    Returns:
+        Execution result
+    """
+    executor = WorkflowExecutor(
+        openai_api_key=openai_api_key,
+        anthropic_api_key=anthropic_api_key
+    )
+    
+    # Run async execution in sync context
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    try:
+        result = loop.run_until_complete(
+            executor.execute(nodes, edges, input_data)
+        )
+        return result
+    finally:
+        loop.close()
+
+
+async def execute_workflow_async(
+    nodes: List[Dict[str, Any]],
+    edges: List[Dict[str, Any]],
+    input_data: Dict[str, Any],
+    openai_api_key: Optional[str] = None,
+    anthropic_api_key: Optional[str] = None,
+) -> Dict[str, Any]:
+    """Async convenience function for workflow execution.
+
+    This is the correct entry point when you're *already* inside an asyncio
+    event loop (e.g. FastAPI/Uvicorn request handlers).
+
+    Args:
+        nodes: Workflow nodes
+        edges: Workflow edges
+        input_data: Input data
+        openai_api_key: OpenAI API key
+        anthropic_api_key: Anthropic API key
+
+    Returns:
+        Execution result
+    """
+    executor = WorkflowExecutor(
+        openai_api_key=openai_api_key,
+        anthropic_api_key=anthropic_api_key,
+    )
+    return await executor.execute(nodes, edges, input_data)

genxai/core/graph/nodes.py

@@ -0,0 +1,161 @@
+"""Node types and implementations for graph-based orchestration."""
+
+from enum import Enum
+from typing import Any, Dict, Optional, Protocol
+from pydantic import BaseModel, Field, ConfigDict
+
+
+class NodeType(str, Enum):
+    """Types of nodes in the graph."""
+
+    AGENT = "agent"
+    TOOL = "tool"
+    CONDITION = "condition"
+    SUBGRAPH = "subgraph"
+    HUMAN = "human"
+    INPUT = "input"
+    OUTPUT = "output"
+
+
+class NodeConfig(BaseModel):
+    """Configuration for a node."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    type: NodeType
+    data: Dict[str, Any] = Field(default_factory=dict)
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+
+
+class NodeStatus(str, Enum):
+    """Execution status of a node."""
+
+    PENDING = "pending"
+    RUNNING = "running"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    SKIPPED = "skipped"
+
+
+class Node(BaseModel):
+    """Base node in the execution graph."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    id: str
+    type: NodeType
+    config: NodeConfig
+    status: NodeStatus = NodeStatus.PENDING
+    result: Optional[Any] = None
+    error: Optional[str] = None
+
+
+    def __repr__(self) -> str:
+        """String representation of the node."""
+        return f"Node(id={self.id}, type={self.type}, status={self.status})"
+
+    def __hash__(self) -> int:
+        """Hash function for node."""
+        return hash(self.id)
+
+
+class NodeExecutor(Protocol):
+    """Protocol for node execution."""
+
+    async def execute(self, node: Node, context: Dict[str, Any]) -> Any:
+        """Execute the node with given context."""
+        ...
+
+
+class AgentNode(Node):
+    """Node that executes an agent."""
+
+    def __init__(
+        self,
+        id: str,
+        agent_id: Optional[str] = None,
+        agent: Optional[Any] = None,
+        task: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize agent node.
+
+        This supports two construction styles:
+        1) Internal graph API: AgentNode(id="x", agent_id="agent_x")
+        2) Integration-test/user API: AgentNode(id="x", agent=<Agent>, task="...")
+
+        When `agent` is provided, we register it into AgentRegistry to make
+        it discoverable by the execution layer.
+        """
+        # Lazy import to avoid circular imports
+        from genxai.core.agent.registry import AgentRegistry
+
+        resolved_agent_id = agent_id
+        if agent is not None:
+            resolved_agent_id = agent.id
+            # Ensure agent is registered so EnhancedGraph can look it up.
+            try:
+                AgentRegistry.register(agent)
+            except Exception:
+                # If already registered, ignore.
+                pass
+
+        if not resolved_agent_id:
+            raise TypeError("AgentNode requires either agent_id or agent")
+
+        data: Dict[str, Any] = {"agent_id": resolved_agent_id}
+        if task is not None:
+            data["task"] = task
+
+        super().__init__(
+            id=id,
+            type=NodeType.AGENT,
+            config=NodeConfig(type=NodeType.AGENT, data=data),
+            **kwargs,
+        )
+
+
+class ToolNode(Node):
+    """Node that executes a tool."""
+
+    def __init__(self, id: str, tool_name: str, **kwargs: Any) -> None:
+        """Initialize tool node."""
+        super().__init__(
+            id=id,
+            type=NodeType.TOOL,
+            config=NodeConfig(type=NodeType.TOOL, data={"tool_name": tool_name}),
+            **kwargs,
+        )
+
+
+class ConditionNode(Node):
+    """Node that evaluates a condition."""
+
+    def __init__(self, id: str, condition: str, **kwargs: Any) -> None:
+        """Initialize condition node."""
+        super().__init__(
+            id=id,
+            type=NodeType.CONDITION,
+            config=NodeConfig(type=NodeType.CONDITION, data={"condition": condition}),
+            **kwargs,
+        )
+
+
+class InputNode(Node):
+    """Node that receives input."""
+
+    def __init__(self, id: str = "input", **kwargs: Any) -> None:
+        """Initialize input node."""
+        super().__init__(
+            id=id, type=NodeType.INPUT, config=NodeConfig(type=NodeType.INPUT), **kwargs
+        )
+
+
+class OutputNode(Node):
+    """Node that produces output."""
+
+    def __init__(self, id: str = "output", **kwargs: Any) -> None:
+        """Initialize output node."""
+        super().__init__(
+            id=id, type=NodeType.OUTPUT, config=NodeConfig(type=NodeType.OUTPUT), **kwargs
+        )