genxai-framework 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

genxai/core/graph/executor.py

@@ -2,12 +2,24 @@
 
 import asyncio
 import copy
-from typing import Any, Dict, List, Optional
+from typing import Any, Callable, 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.memory.shared import SharedMemoryBus
+from genxai.core.graph.nodes import (
+    InputNode,
+    OutputNode,
+    AgentNode,
+    ConditionNode,
+    ToolNode,
+    SubgraphNode,
+    LoopNode,
+    Node,
+    NodeConfig,
+    NodeType,
+)
 from genxai.core.graph.edges import Edge, ConditionalEdge
 from genxai.core.agent.base import Agent, AgentFactory
 from genxai.core.agent.registry import AgentRegistry
@@ -30,7 +42,9 @@ class EnhancedGraph(Graph):
     in GenXAI.
     """
 
-    async def _execute_node_logic(self, node: Any, state: Dict[str, Any]) -> Any:
+    async def _execute_node_logic(
+        self, node: Any, state: Dict[str, Any], max_iterations: int = 100
+    ) -> Any:
         """Execute node logic with actual agent execution.
 
         Args:
@@ -74,8 +88,7 @@ class EnhancedGraph(Graph):
             return result
         
         else:
-            # Default behavior
-            return {"node_id": node.id, "type": node.type.value}
+            return await super()._execute_node_logic(node, state, max_iterations)
 
     async def _execute_agent_with_tools(
         self, agent: Agent, task: str, state: Dict[str, Any]
@@ -95,7 +108,15 @@ class EnhancedGraph(Graph):
         # Use AgentRuntime for full integration
         from genxai.core.agent.runtime import AgentRuntime
         
-        runtime = AgentRuntime(agent=agent, enable_memory=True)
+        # Pass both API keys to runtime so it can select the correct one based on model
+        runtime = AgentRuntime(
+            agent=agent,
+            llm_provider=getattr(self, "llm_provider", None),
+            openai_api_key=getattr(self, "openai_api_key", None),
+            anthropic_api_key=getattr(self, "anthropic_api_key", None),
+            enable_memory=True,
+            shared_memory=getattr(self, "shared_memory", None),
+        )
         
         # Load tools from registry
         if agent.config.tools:
@@ -108,7 +129,10 @@ class EnhancedGraph(Graph):
             logger.debug(f"Loaded {len(tools)} tools for agent")
         
         # Execute agent with full runtime support
-        result = await runtime.execute(task, context=state)
+        context = dict(state)
+        if getattr(self, "shared_memory", None) is not None:
+            context["shared_memory"] = getattr(self, "shared_memory")
+        result = await runtime.execute(task, context=context)
         
         return result
 
@@ -256,11 +280,14 @@ class WorkflowExecutor:
             Constructed graph
         """
         graph = EnhancedGraph(name="workflow")
+        graph.openai_api_key = self.openai_api_key
+        graph.anthropic_api_key = self.anthropic_api_key
 
         # Add nodes
         for node in nodes:
             node_id = node.get("id")
             node_type = node.get("type")
+            config = node.get("config", {})
 
             # Support some common aliases used by the Studio UI
             # - "start" behaves like an input node
@@ -271,6 +298,29 @@ class WorkflowExecutor:
                 graph.add_node(OutputNode(id=node_id))
             elif node_type == "agent":
                 graph.add_node(AgentNode(id=node_id, agent_id=node_id))
+            elif node_type == "tool":
+                tool_name = config.get("tool_name") or config.get("name") or "tool"
+                graph.add_node(ToolNode(id=node_id, tool_name=tool_name))
+            elif node_type == "decision":
+                condition = config.get("condition", "")
+                graph.add_node(ConditionNode(id=node_id, condition=condition))
+            elif node_type == "subgraph":
+                workflow_id = config.get("workflow_id") or config.get("subgraph_id") or config.get("workflow")
+                if workflow_id:
+                    graph.add_node(SubgraphNode(id=node_id, workflow_id=workflow_id))
+                else:
+                    graph.add_node(
+                        Node(
+                            id=node_id,
+                            type=NodeType.SUBGRAPH,
+                            config=NodeConfig(type=NodeType.SUBGRAPH, data={"workflow_id": ""}),
+                        )
+                    )
+                    logger.warning(f"Subgraph node '{node_id}' missing workflow_id")
+            elif node_type == "loop":
+                condition = config.get("condition", "")
+                max_iterations = int(config.get("max_iterations", 5))
+                graph.add_node(LoopNode(id=node_id, condition=condition, max_iterations=max_iterations))
             else:
                 logger.warning(f"Unknown node type: {node_type}")
 
@@ -321,6 +371,10 @@ class WorkflowExecutor:
         run_id: Optional[str] = None,
         checkpoint_dir: Optional[str] = None,
         resume_from: Optional[str] = None,
+        model_override: Optional[str] = None,
+        event_callback: Optional[Callable[[Dict[str, Any]], Any]] = None,
+        shared_memory: bool = False,
+        llm_provider: Optional[Any] = None,
     ) -> Dict[str, Any]:
         """Execute a workflow.
 
@@ -338,11 +392,23 @@ class WorkflowExecutor:
         try:
             logger.info("Starting workflow execution")
 
+            # Apply model override if provided
+            if model_override:
+                for node in nodes:
+                    if node.get("type") == "agent":
+                        config = node.setdefault("config", {})
+                        config["llm_model"] = model_override
+
             # Create agents from nodes
             self._create_agents_from_nodes(nodes)
 
             # Build graph
             graph = self._build_graph(nodes, edges)
+            graph.llm_provider = llm_provider
+            if shared_memory:
+                graph.set_shared_memory(SharedMemoryBus())
+                graph.shared_memory = graph.shared_memory
+                graph.shared_memory_enabled = True
 
             # Validate graph
             graph.validate()
@@ -364,7 +430,11 @@ class WorkflowExecutor:
                         status="allowed",
                     )
                 )
-            result = await graph.run(input_data=input_data, resume_from=checkpoint)
+            result = await graph.run(
+                input_data=input_data,
+                resume_from=checkpoint,
+                event_callback=event_callback,
+            )
 
             logger.info("Workflow execution completed successfully")
 
@@ -378,6 +448,7 @@ class WorkflowExecutor:
                 "status": "success",
                 "run_id": run_id,
                 "result": result,
+                "node_events": result.get("node_events", []),
                 "nodes_executed": len(graph.nodes),
                 "message": "Workflow executed successfully"
             }
@@ -410,6 +481,7 @@ class WorkflowExecutor:
         run_id: Optional[str] = None,
         checkpoint_dir: Optional[str] = None,
         resume_from: Optional[str] = None,
+        model_override: Optional[str] = None,
     ) -> str:
         """Enqueue workflow execution using a worker queue engine."""
         if not self.queue_engine:
@@ -430,6 +502,7 @@ class WorkflowExecutor:
                 run_id=payload["run_id"],
                 checkpoint_dir=payload.get("checkpoint_dir"),
                 resume_from=payload.get("resume_from"),
+                model_override=payload.get("model_override"),
             )
 
         await self.queue_engine.start()
@@ -441,6 +514,7 @@ class WorkflowExecutor:
                 "run_id": run_id,
                 "checkpoint_dir": checkpoint_dir,
                 "resume_from": resume_from,
+                "model_override": model_override,
             },
             _handler,
             metadata={"workflow": "queued"},
@@ -454,6 +528,8 @@ def execute_workflow_sync(
     input_data: Dict[str, Any],
     openai_api_key: Optional[str] = None,
     anthropic_api_key: Optional[str] = None,
+    model_override: Optional[str] = None,
+    shared_memory: bool = False,
 ) -> Dict[str, Any]:
     """Synchronous wrapper for workflow execution.
     
@@ -480,7 +556,13 @@ def execute_workflow_sync(
     asyncio.set_event_loop(loop)
     try:
         result = loop.run_until_complete(
-            executor.execute(nodes, edges, input_data)
+            executor.execute(
+                nodes,
+                edges,
+                input_data,
+                model_override=model_override,
+                shared_memory=shared_memory,
+            )
         )
         return result
     finally:
@@ -493,6 +575,9 @@ async def execute_workflow_async(
     input_data: Dict[str, Any],
     openai_api_key: Optional[str] = None,
     anthropic_api_key: Optional[str] = None,
+    model_override: Optional[str] = None,
+    event_callback: Optional[Callable[[Dict[str, Any]], Any]] = None,
+    shared_memory: bool = False,
 ) -> Dict[str, Any]:
     """Async convenience function for workflow execution.
 
@@ -513,4 +598,12 @@ async def execute_workflow_async(
         openai_api_key=openai_api_key,
         anthropic_api_key=anthropic_api_key,
     )
-    return await executor.execute(nodes, edges, input_data)
+    return await executor.execute(
+        nodes,
+        edges,
+        input_data,
+        model_override=model_override,
+        event_callback=event_callback,
+        shared_memory=shared_memory,
+    )
+    

genxai/core/graph/nodes.py

@@ -15,6 +15,7 @@ class NodeType(str, Enum):
     HUMAN = "human"
     INPUT = "input"
     OUTPUT = "output"
+    LOOP = "loop"
 
 
 class NodeConfig(BaseModel):
@@ -159,3 +160,30 @@ class OutputNode(Node):
         super().__init__(
             id=id, type=NodeType.OUTPUT, config=NodeConfig(type=NodeType.OUTPUT), **kwargs
         )
+
+
+class SubgraphNode(Node):
+    """Node that executes a nested workflow."""
+
+    def __init__(self, id: str, workflow_id: str, **kwargs: Any) -> None:
+        super().__init__(
+            id=id,
+            type=NodeType.SUBGRAPH,
+            config=NodeConfig(type=NodeType.SUBGRAPH, data={"workflow_id": workflow_id}),
+            **kwargs,
+        )
+
+
+class LoopNode(Node):
+    """Node that represents a loop with a termination condition."""
+
+    def __init__(self, id: str, condition: str, max_iterations: int = 5, **kwargs: Any) -> None:
+        super().__init__(
+            id=id,
+            type=NodeType.LOOP,
+            config=NodeConfig(
+                type=NodeType.LOOP,
+                data={"condition": condition, "max_iterations": max_iterations},
+            ),
+            **kwargs,
+        )

genxai/core/graph/workflow_io.py

@@ -0,0 +1,199 @@
+"""Workflow YAML loading utilities."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Dict, List
+
+import yaml
+
+from genxai.core.agent.base import Agent
+from genxai.core.agent.registry import AgentRegistry
+from genxai.core.agent.config_io import import_agents_yaml
+
+
+def load_workflow_yaml(path: Path) -> Dict[str, Any]:
+    """Load a workflow YAML file.
+
+    Supports `agents_ref` to pull reusable agent definitions from another YAML file.
+    Inline agents remain supported and will be merged with referenced agents.
+    """
+    payload = yaml.safe_load(path.read_text())
+    if not isinstance(payload, dict) or "workflow" not in payload:
+        raise ValueError("Workflow YAML must contain a top-level 'workflow' mapping")
+
+    workflow = payload["workflow"]
+    if not isinstance(workflow, dict):
+        raise ValueError("workflow must be a mapping")
+
+    _merge_agents_ref(workflow, base_path=path.parent)
+    _validate_workflow_schema(workflow)
+    return workflow
+
+
+def register_workflow_agents(workflow: Dict[str, Any]) -> List[Agent]:
+    """Register agents defined in a workflow dict and return them.
+
+    Accepts agent dictionaries with the `agents_ref` already resolved.
+    """
+    agents_payload = workflow.get("agents", [])
+    if not agents_payload:
+        return []
+
+    workflow_memory = workflow.get("memory") if isinstance(workflow.get("memory"), dict) else {}
+
+    agents: List[Agent] = []
+    for agent_data in agents_payload:
+        if not isinstance(agent_data, dict):
+            raise ValueError("Invalid agent definition in workflow")
+        merged_agent = _apply_workflow_memory_defaults(agent_data, workflow_memory)
+        agent = _agent_from_workflow_dict(merged_agent)
+        AgentRegistry.register(agent)
+        agents.append(agent)
+    return agents
+
+
+def _apply_workflow_memory_defaults(
+    agent_data: Dict[str, Any], workflow_memory: Dict[str, Any]
+) -> Dict[str, Any]:
+    if not workflow_memory:
+        return agent_data
+
+    defaults: Dict[str, Any] = {}
+    if "enabled" in workflow_memory:
+        defaults["enabled"] = workflow_memory.get("enabled")
+    if "type" in workflow_memory:
+        defaults["type"] = workflow_memory.get("type")
+
+    if not defaults:
+        return agent_data
+
+    if isinstance(agent_data.get("memory"), dict):
+        memory_block = dict(agent_data.get("memory") or {})
+        memory_block.setdefault("enabled", defaults.get("enabled"))
+        memory_block.setdefault("type", defaults.get("type"))
+        agent_data = {**agent_data, "memory": memory_block}
+        return agent_data
+
+    if "enable_memory" in agent_data or "memory_type" in agent_data:
+        return agent_data
+
+    return {**agent_data, "memory": {k: v for k, v in defaults.items() if v is not None}}
+
+
+def _merge_agents_ref(workflow: Dict[str, Any], base_path: Path) -> None:
+    agents_ref = workflow.get("agents_ref")
+    if not agents_ref:
+        return
+
+    agents_ref_path = (base_path / agents_ref).resolve()
+    referenced = import_agents_yaml(agents_ref_path)
+    referenced_dicts = [
+        {"id": agent.id, **agent.config.model_dump(mode="json")}
+        if isinstance(agent, Agent)
+        else dict(agent)
+        for agent in referenced
+    ]
+
+    inline_agents = workflow.get("agents", [])
+    if inline_agents is None:
+        inline_agents = []
+    if not isinstance(inline_agents, list):
+        raise ValueError("workflow.agents must be a list")
+
+    # Merge with inline agents taking precedence by id.
+    merged = {agent["id"]: agent for agent in referenced_dicts if isinstance(agent, dict)}
+    for agent in inline_agents:
+        if not isinstance(agent, dict) or "id" not in agent:
+            raise ValueError("Invalid inline agent definition")
+        merged[agent["id"]] = agent
+
+    workflow["agents"] = list(merged.values())
+
+
+def _validate_workflow_schema(workflow: Dict[str, Any]) -> None:
+    if not workflow.get("name"):
+        raise ValueError("workflow.name is required")
+
+    if "memory" in workflow and not isinstance(workflow.get("memory"), dict):
+        raise ValueError("workflow.memory must be a mapping when provided")
+
+    graph = workflow.get("graph")
+    if not isinstance(graph, dict):
+        raise ValueError("workflow.graph must be a mapping")
+
+    nodes = graph.get("nodes")
+    if not isinstance(nodes, list) or not nodes:
+        raise ValueError("workflow.graph.nodes must be a non-empty list")
+
+    edges = graph.get("edges")
+    if not isinstance(edges, list):
+        raise ValueError("workflow.graph.edges must be a list")
+
+    node_ids = set()
+    for node in nodes:
+        if not isinstance(node, dict):
+            raise ValueError("workflow.graph.nodes entries must be mappings")
+        if "id" not in node or "type" not in node:
+            raise ValueError("Each node requires 'id' and 'type'")
+        node_ids.add(node["id"])
+        if node["type"] not in {"input", "start", "output", "end", "agent", "tool", "condition"}:
+            raise ValueError(f"Unsupported node type: {node['type']}")
+
+    for edge in edges:
+        if not isinstance(edge, dict):
+            raise ValueError("workflow.graph.edges entries must be mappings")
+        if "from" not in edge or "to" not in edge:
+            raise ValueError("Each edge requires 'from' and 'to'")
+        if edge["from"] not in node_ids or edge["to"] not in node_ids:
+            raise ValueError("Edges must reference existing node ids")
+
+    agent_ids = {agent.get("id") for agent in workflow.get("agents", []) if isinstance(agent, dict)}
+    for node in nodes:
+        if node.get("type") == "agent" and node.get("id") not in agent_ids:
+            raise ValueError(f"Agent node '{node['id']}' has no matching agent definition")
+
+
+def _agent_from_workflow_dict(data: Dict[str, Any]) -> Agent:
+    config_payload = data.get("config") if isinstance(data.get("config"), dict) else {}
+    merged = {
+        **config_payload,
+        **{k: v for k, v in data.items() if k not in {"config"}},
+    }
+    return Agent(
+        id=data["id"],
+        config=_agent_config_from_workflow_dict(merged),
+    )
+
+
+def _agent_config_from_workflow_dict(data: Dict[str, Any]):
+    from genxai.core.agent.base import AgentConfig
+
+    # Allow workflow agents to specify either `llm` (legacy) or `llm_model`.
+    llm_model = data.get("llm_model") or data.get("llm") or "gpt-4"
+
+    return AgentConfig(
+        role=data.get("role", "Agent"),
+        goal=data.get("goal", "Process tasks"),
+        backstory=data.get("backstory", ""),
+        llm_provider=data.get("llm_provider", "openai"),
+        llm_model=llm_model,
+        llm_temperature=data.get("llm_temperature", 0.7),
+        tools=data.get("tools", []),
+        enable_memory=data.get("memory", {}).get("enabled", True)
+        if isinstance(data.get("memory"), dict)
+        else data.get("enable_memory", True),
+        memory_type=data.get("memory", {}).get("type", "short_term")
+        if isinstance(data.get("memory"), dict)
+        else data.get("memory_type", "short_term"),
+        agent_type=data.get("behavior", {}).get("agent_type", "reactive")
+        if isinstance(data.get("behavior"), dict)
+        else data.get("agent_type", "reactive"),
+        max_iterations=data.get("behavior", {}).get("max_iterations", 10)
+        if isinstance(data.get("behavior"), dict)
+        else data.get("max_iterations", 10),
+        verbose=data.get("behavior", {}).get("verbose", False)
+        if isinstance(data.get("behavior"), dict)
+        else data.get("verbose", False),
+        metadata=data.get("metadata", {}),
+    )

genxai/flows/__init__.py

@@ -0,0 +1,33 @@
+"""Flow orchestrators for common agent coordination patterns."""
+
+from genxai.flows.base import FlowOrchestrator
+from genxai.flows.round_robin import RoundRobinFlow
+from genxai.flows.selector import SelectorFlow
+from genxai.flows.p2p import P2PFlow
+from genxai.flows.parallel import ParallelFlow
+from genxai.flows.conditional import ConditionalFlow
+from genxai.flows.loop import LoopFlow
+from genxai.flows.router import RouterFlow
+from genxai.flows.ensemble_voting import EnsembleVotingFlow
+from genxai.flows.critic_review import CriticReviewFlow
+from genxai.flows.coordinator_worker import CoordinatorWorkerFlow
+from genxai.flows.map_reduce import MapReduceFlow
+from genxai.flows.subworkflow import SubworkflowFlow
+from genxai.flows.auction import AuctionFlow
+
+__all__ = [
+    "FlowOrchestrator",
+    "RoundRobinFlow",
+    "SelectorFlow",
+    "P2PFlow",
+    "ParallelFlow",
+    "ConditionalFlow",
+    "LoopFlow",
+    "RouterFlow",
+    "EnsembleVotingFlow",
+    "CriticReviewFlow",
+    "CoordinatorWorkerFlow",
+    "MapReduceFlow",
+    "SubworkflowFlow",
+    "AuctionFlow",
+]

genxai/flows/auction.py

@@ -0,0 +1,66 @@
+"""Auction flow orchestrator."""
+
+from typing import Any, Dict, List, Optional
+
+from genxai.core.agent.runtime import AgentRuntime
+from genxai.flows.base import FlowOrchestrator
+
+
+class AuctionFlow(FlowOrchestrator):
+    """Agents bid to handle a task; highest bid executes."""
+
+    def __init__(
+        self,
+        agents: List[Any],
+        name: str = "auction_flow",
+        llm_provider: Any = None,
+    ) -> None:
+        super().__init__(agents=agents, name=name, llm_provider=llm_provider)
+
+    async def run(
+        self,
+        input_data: Any,
+        state: Optional[Dict[str, Any]] = None,
+        max_iterations: int = 100,
+    ) -> Dict[str, Any]:
+        if state is None:
+            state = {}
+        state["input"] = input_data
+        state.setdefault("bids", {})
+
+        runtimes = {
+            agent.id: AgentRuntime(agent=agent, llm_provider=self.llm_provider)
+            for agent in self.agents
+        }
+
+        bid_task = state.get("bid_task", "Provide a numeric bid between 0 and 1")
+        tasks = [
+            self._execute_with_retry(
+                runtimes[agent.id],
+                task=bid_task,
+                context=state,
+            )
+            for agent in self.agents
+        ]
+        results = await self._gather_tasks(tasks)
+        for agent, bid_result in zip(self.agents, results):
+            bid_value = 0.0
+            try:
+                bid_value = float(bid_result.get("output", 0))
+            except (TypeError, ValueError, AttributeError):
+                bid_value = 0.0
+            state["bids"][agent.id] = bid_value
+
+        if not state["bids"]:
+            raise ValueError("AuctionFlow requires at least one bid")
+
+        winner_id = max(state["bids"], key=state["bids"].get)
+        winner_runtime = runtimes[winner_id]
+        execution = await self._execute_with_retry(
+            winner_runtime,
+            task=state.get("task", "Execute the task"),
+            context={**state, "winner_id": winner_id},
+        )
+        state["winner_id"] = winner_id
+        state["winner_result"] = execution
+        return state