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

genxai/core/agent/runtime.py

@@ -5,6 +5,7 @@ import asyncio
 import time
 import logging
 import json
+import copy
 
 from genxai.core.agent.base import Agent
 from genxai.llm.base import LLMProvider
@@ -16,6 +17,7 @@ from genxai.observability.tracing import span, add_event, record_exception
 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
+from genxai.core.memory.shared import SharedMemoryBus
 
 logger = logging.getLogger(__name__)
 
@@ -34,19 +36,25 @@ class AgentRuntime:
         agent: Agent,
         llm_provider: Optional[LLMProvider] = None,
         api_key: Optional[str] = None,
+        openai_api_key: Optional[str] = None,
+        anthropic_api_key: Optional[str] = None,
         enable_memory: bool = True,
+        shared_memory: Optional[SharedMemoryBus] = None,
     ) -> None:
         """Initialize agent runtime.
 
         Args:
             agent: Agent to execute
             llm_provider: LLM provider instance (optional, will be created if not provided)
-            api_key: API key for LLM provider (optional, will use env var if not provided)
+            api_key: API key for LLM provider (optional, deprecated - use openai_api_key or anthropic_api_key)
+            openai_api_key: OpenAI API key (for GPT models)
+            anthropic_api_key: Anthropic API key (for Claude models)
             enable_memory: Whether to initialize memory system
         """
         self.agent = agent
         self._tools: Dict[str, Any] = {}
         self._memory: Optional[Any] = None
+        self._shared_memory = shared_memory
 
         # Initialize LLM provider
         if llm_provider:
@@ -54,9 +62,25 @@ class AgentRuntime:
         else:
             # Create provider from agent config
             try:
+                # Determine which API key to use based on model
+                model = agent.config.llm_model.lower()
+                selected_api_key = api_key  # Fallback to deprecated api_key parameter
+                
+                if model.startswith('claude'):
+                    # Claude models use Anthropic API key
+                    selected_api_key = anthropic_api_key or api_key
+                    logger.info(f"Using Anthropic API key for Claude model: {agent.config.llm_model}")
+                elif model.startswith('gpt'):
+                    # GPT models use OpenAI API key
+                    selected_api_key = openai_api_key or api_key
+                    logger.info(f"Using OpenAI API key for GPT model: {agent.config.llm_model}")
+                else:
+                    # For other models, try to infer or use openai_api_key as default
+                    selected_api_key = openai_api_key or anthropic_api_key or api_key
+                
                 self._llm_provider = LLMProviderFactory.create_provider(
                     model=agent.config.llm_model,
-                    api_key=api_key,
+                    api_key=selected_api_key,
                     temperature=agent.config.llm_temperature,
                     max_tokens=agent.config.llm_max_tokens,
                 )
@@ -177,7 +201,13 @@ class AgentRuntime:
             memory_context = await self.get_memory_context(limit=5)
         
         # Build prompt (without memory context, as it's handled in _get_llm_response)
-        prompt = self._build_prompt(task, context, "")
+        prompt_context = dict(context)
+        if self._shared_memory is not None:
+            prompt_context["shared_memory"] = {
+                key: self._shared_memory.get(key)
+                for key in self._shared_memory.list_keys()
+            }
+        prompt = self._build_prompt(task, prompt_context, "")
         
         # Get LLM response with retry logic and memory context
         if self.agent.config.tools and self._tools and self._provider_supports_tools():
@@ -193,13 +223,20 @@ class AgentRuntime:
             await self._update_memory(task, response)
         
         # Build result
+        safe_context: Dict[str, Any]
+        try:
+            safe_context = copy.deepcopy(context)
+        except Exception:
+            safe_context = dict(context)
+        safe_context.pop("llm_provider", None)
+        safe_context.pop("shared_memory", None)
         result = {
             "agent_id": self.agent.id,
             "task": task,
             "status": "completed",
             "output": response,
-            "context": context,
             "tokens_used": self.agent._total_tokens,
+            "context": safe_context,
         }
         
         # Store episode in episodic memory

genxai/core/graph/__init__.py

@@ -14,6 +14,7 @@ from genxai.core.graph.checkpoints import (
 )
 from genxai.core.graph.trigger_runner import TriggerWorkflowRunner
 from genxai.core.execution import WorkerQueueEngine
+from genxai.core.graph.workflow_io import load_workflow_yaml, register_workflow_agents
 
 __all__ = [
     "Node",
@@ -27,4 +28,6 @@ __all__ = [
     "WorkflowCheckpointManager",
     "TriggerWorkflowRunner",
     "WorkerQueueEngine",
+    "load_workflow_yaml",
+    "register_workflow_agents",
 ]

genxai/core/graph/engine.py

@@ -1,18 +1,19 @@
 """Graph execution engine for orchestrating agent workflows."""
 
 import asyncio
-from typing import Any, Dict, List, Optional, Set
+from typing import Any, Callable, Dict, List, Optional, Set
 from collections import defaultdict, deque
 import logging
 import time
 import copy
 from pathlib import Path
 
-from genxai.core.graph.nodes import Node, NodeStatus, NodeType
+from genxai.core.graph.nodes import Node, NodeConfig, NodeStatus, NodeType
 from genxai.core.agent.registry import AgentRegistry
 from genxai.core.agent.runtime import AgentRuntime
 from genxai.tools.registry import ToolRegistry
 from genxai.core.graph.edges import Edge
+from genxai.core.memory.shared import SharedMemoryBus
 from genxai.core.graph.checkpoints import (
     WorkflowCheckpoint,
     WorkflowCheckpointManager,
@@ -47,6 +48,7 @@ class Graph:
         self.edges: List[Edge] = []
         self._adjacency_list: Dict[str, List[Edge]] = defaultdict(list)
         self._reverse_adjacency: Dict[str, List[str]] = defaultdict(list)
+        self.shared_memory: Optional[SharedMemoryBus] = None
 
     def add_node(self, node: Node) -> None:
         """Add a node to the graph.
@@ -82,6 +84,10 @@ class Graph:
         self._reverse_adjacency[edge.target].append(edge.source)
         logger.debug(f"Added edge: {edge.source} -> {edge.target}")
 
+    def set_shared_memory(self, shared_memory: Optional[SharedMemoryBus]) -> None:
+        """Attach a shared memory bus to the graph for agent execution."""
+        self.shared_memory = shared_memory
+
     def get_node(self, node_id: str) -> Optional[Node]:
         """Get a node by ID.
 
@@ -212,6 +218,7 @@ class Graph:
         state: Optional[Dict[str, Any]] = None,
         resume_from: Optional[WorkflowCheckpoint] = None,
         llm_provider: Any = None,
+        event_callback: Optional[Callable[[Dict[str, Any]], Any]] = None,
     ) -> Dict[str, Any]:
         """Execute the graph workflow.
 
@@ -244,6 +251,7 @@ class Graph:
                 state = {}
             state["input"] = input_data
             state["iterations"] = 0
+        state.setdefault("node_events", [])
 
         if resume_from:
             for node_id, status in resume_from.node_statuses.items():
@@ -274,7 +282,7 @@ class Graph:
         try:
             with span("genxai.workflow.execute", {"workflow_id": self.name}):
                 for entry_point in entry_points:
-                    await self._execute_node(entry_point, state, max_iterations)
+                    await self._execute_node(entry_point, state, max_iterations, event_callback)
         except Exception as exc:
             status = "error"
             record_exception(exc)
@@ -288,6 +296,7 @@ class Graph:
             )
 
         logger.info(f"Graph execution completed: {self.name}")
+        state["node_events"] = state.get("node_events", [])
         return state
 
     def create_checkpoint(self, name: str, state: Dict[str, Any]) -> WorkflowCheckpoint:
@@ -307,7 +316,11 @@ class Graph:
         return manager.load(name)
 
     async def _execute_node(
-        self, node_id: str, state: Dict[str, Any], max_iterations: int
+        self,
+        node_id: str,
+        state: Dict[str, Any],
+        max_iterations: int,
+        event_callback: Optional[Callable[[Dict[str, Any]], Any]] = None,
     ) -> None:
         """Execute a single node and its descendants.
 
@@ -334,6 +347,16 @@ class Graph:
         node.status = NodeStatus.RUNNING
         logger.debug(f"Executing node: {node_id}")
         node_start = time.time()
+        running_event = {
+            "node_id": node_id,
+            "status": NodeStatus.RUNNING.value,
+            "timestamp": time.time(),
+        }
+        state.setdefault("node_events", []).append(running_event)
+        if event_callback:
+            callback_result = event_callback(running_event)
+            if asyncio.iscoroutine(callback_result):
+                await callback_result
 
         try:
             # Execute node (placeholder - will be implemented with actual executors)
@@ -341,16 +364,35 @@ class Graph:
                 "genxai.workflow.node",
                 {"workflow_id": self.name, "node_id": node_id, "node_type": node.type.value},
             ):
-                result = await self._execute_node_logic(node, state)
+                result = await self._execute_node_logic(node, state, max_iterations)
             node.result = result
             node.status = NodeStatus.COMPLETED
             logger.debug(f"Node completed: {node_id}")
 
+            node_duration_ms = int((time.time() - node_start) * 1000)
+
             record_workflow_node_execution(
                 workflow_id=self.name,
                 node_id=node_id,
                 status="success",
             )
+            completed_event = {
+                "node_id": node_id,
+                "status": NodeStatus.COMPLETED.value,
+                "timestamp": time.time(),
+                "duration_ms": node_duration_ms,
+            }
+            state.setdefault("node_events", []).append(completed_event)
+            if event_callback:
+                callback_result = event_callback(completed_event)
+                if asyncio.iscoroutine(callback_result):
+                    await callback_result
+
+            state.setdefault("node_results", {})[node_id] = {
+                "output": result,
+                "status": NodeStatus.COMPLETED.value,
+                "duration_ms": node_duration_ms,
+            }
 
             # Update state with result
             state[node_id] = result
@@ -367,27 +409,48 @@ class Graph:
                 tasks = []
                 for edge in parallel_edges:
                     if edge.evaluate_condition(state):
-                        tasks.append(self._execute_node(edge.target, state, max_iterations))
+                        tasks.append(self._execute_node(edge.target, state, max_iterations, event_callback))
                 if tasks:
-                    await asyncio.gather(*tasks)
+                    await self._gather_with_config(tasks, state)
 
             # Execute sequential edges in order
             for edge in sorted(sequential_edges, key=lambda e: e.priority):
                 if edge.evaluate_condition(state):
-                    await self._execute_node(edge.target, state, max_iterations)
+                    await self._execute_node(edge.target, state, max_iterations, event_callback)
 
         except Exception as e:
             node.status = NodeStatus.FAILED
             node.error = str(e)
             logger.error(f"Node execution failed: {node_id} - {e}")
+            node_duration_ms = int((time.time() - node_start) * 1000)
             record_workflow_node_execution(
                 workflow_id=self.name,
                 node_id=node_id,
                 status="error",
             )
+            failed_event = {
+                "node_id": node_id,
+                "status": NodeStatus.FAILED.value,
+                "timestamp": time.time(),
+                "error": str(e),
+                "duration_ms": node_duration_ms,
+            }
+            state.setdefault("node_events", []).append(failed_event)
+            if event_callback:
+                callback_result = event_callback(failed_event)
+                if asyncio.iscoroutine(callback_result):
+                    await callback_result
+            state.setdefault("node_results", {})[node_id] = {
+                "output": None,
+                "status": NodeStatus.FAILED.value,
+                "duration_ms": node_duration_ms,
+                "error": str(e),
+            }
             raise GraphExecutionError(f"Node {node_id} failed: {e}") from e
 
-    async def _execute_node_logic(self, node: Node, state: Dict[str, Any]) -> Any:
+    async def _execute_node_logic(
+        self, node: Node, state: Dict[str, Any], max_iterations: int
+    ) -> Any:
         """Execute the actual logic of a node.
 
         Args:
@@ -409,6 +472,12 @@ class Graph:
         if node.type == NodeType.TOOL:
             return await self._execute_tool_node(node, state)
 
+        if node.type == NodeType.SUBGRAPH:
+            return await self._execute_subgraph_node(node, state, max_iterations)
+
+        if node.type == NodeType.LOOP:
+            return await self._execute_loop_node(node, state, max_iterations)
+
         # Default fallback for unsupported nodes
         return {"node_id": node.id, "type": node.type.value}
 
@@ -435,7 +504,12 @@ class Graph:
         task = node.config.data.get("task") or state.get("task") or "Process input"
 
         llm_provider = state.get("llm_provider")
-        runtime = AgentRuntime(agent=agent, llm_provider=llm_provider, enable_memory=True)
+        runtime = AgentRuntime(
+            agent=agent,
+            llm_provider=llm_provider,
+            enable_memory=True,
+            shared_memory=self.shared_memory,
+        )
         if agent.config.tools:
             tools: Dict[str, Any] = {}
             for tool_name in agent.config.tools:
@@ -444,7 +518,73 @@ class Graph:
                     tools[tool_name] = tool
             runtime.set_tools(tools)
 
-        return await runtime.execute(task, context=state)
+        context = dict(state)
+        if self.shared_memory is not None:
+            context["shared_memory"] = self.shared_memory
+        return await self._execute_with_config(runtime, task=task, context=context, state=state)
+
+    def _get_execution_config(self, state: Dict[str, Any]) -> Dict[str, Any]:
+        config = state.get("execution_config") or {}
+        return {
+            "timeout_seconds": config.get("timeout_seconds", 120.0),
+            "retry_count": config.get("retry_count", 3),
+            "backoff_base": config.get("backoff_base", 1.0),
+            "backoff_multiplier": config.get("backoff_multiplier", 2.0),
+            "cancel_on_failure": config.get("cancel_on_failure", True),
+        }
+
+    async def _execute_with_config(
+        self,
+        runtime: AgentRuntime,
+        task: str,
+        context: Dict[str, Any],
+        state: Dict[str, Any],
+    ) -> Any:
+        config = self._get_execution_config(state)
+        delay = config["backoff_base"]
+        for attempt in range(config["retry_count"] + 1):
+            try:
+                coro = runtime.execute(task=task, context=context)
+                timeout = config["timeout_seconds"]
+                if timeout:
+                    return await asyncio.wait_for(coro, timeout=timeout)
+                return await coro
+            except asyncio.CancelledError:
+                raise
+            except Exception:
+                if attempt >= config["retry_count"]:
+                    raise
+                await asyncio.sleep(delay)
+                delay *= config["backoff_multiplier"]
+
+    async def _gather_with_config(self, coros: List[Any], state: Dict[str, Any]) -> List[Any]:
+        config = self._get_execution_config(state)
+        tasks = [asyncio.create_task(coro) for coro in coros]
+        if not tasks:
+            return []
+        if not config["cancel_on_failure"]:
+            return await asyncio.gather(*tasks, return_exceptions=True)
+
+        results: List[Any] = [None] * len(tasks)
+        index_map = {task: idx for idx, task in enumerate(tasks)}
+        done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_EXCEPTION)
+
+        for task in done:
+            idx = index_map[task]
+            exc = task.exception()
+            if exc:
+                for pending_task in pending:
+                    pending_task.cancel()
+                await asyncio.gather(*pending, return_exceptions=True)
+                raise exc
+            results[idx] = task.result()
+
+        if pending:
+            pending_results = await asyncio.gather(*pending, return_exceptions=True)
+            for task, result in zip(pending, pending_results):
+                results[index_map[task]] = result
+
+        return results
 
     async def _execute_tool_node(self, node: Node, state: Dict[str, Any]) -> Any:
         """Execute a ToolNode using ToolRegistry.
@@ -477,6 +617,73 @@ class Graph:
         result = await tool.execute(**tool_params)
         return result.model_dump() if hasattr(result, "model_dump") else result
 
+    async def _execute_subgraph_node(
+        self, node: Node, state: Dict[str, Any], max_iterations: int
+    ) -> Any:
+        """Execute a nested workflow defined in the state metadata."""
+        workflow_id = node.config.data.get("workflow_id")
+        if not workflow_id:
+            raise GraphExecutionError(
+                f"Subgraph node '{node.id}' missing workflow_id in config.data"
+            )
+
+        subgraphs = state.get("subgraphs", {})
+        workflow_def = subgraphs.get(workflow_id)
+        if not workflow_def and "subgraphs" in state:
+            workflow_def = state["subgraphs"].get(workflow_id)
+        if not workflow_def and "metadata" in state:
+            workflow_def = state.get("metadata", {}).get("subgraphs", {}).get(workflow_id)
+        if not workflow_def:
+            raise GraphExecutionError(
+                f"Subgraph workflow '{workflow_id}' not found in state.subgraphs"
+            )
+
+        subgraph = Graph(name=f"subgraph:{workflow_id}")
+        for node_def in workflow_def.get("nodes", []):
+            node_type = node_def.get("type")
+            node_id = node_def.get("id")
+            if node_type == "input":
+                subgraph.add_node(Node(id=node_id, type=NodeType.INPUT, config=NodeConfig(type=NodeType.INPUT)))
+            elif node_type == "output":
+                subgraph.add_node(Node(id=node_id, type=NodeType.OUTPUT, config=NodeConfig(type=NodeType.OUTPUT)))
+            elif node_type == "agent":
+                subgraph.add_node(Node(id=node_id, type=NodeType.AGENT, config=NodeConfig(type=NodeType.AGENT, data=node_def.get("config", {}))))
+            elif node_type == "tool":
+                subgraph.add_node(Node(id=node_id, type=NodeType.TOOL, config=NodeConfig(type=NodeType.TOOL, data=node_def.get("config", {}))))
+            else:
+                subgraph.add_node(Node(id=node_id, type=NodeType.CONDITION, config=NodeConfig(type=NodeType.CONDITION, data=node_def.get("config", {}))))
+
+        for edge_def in workflow_def.get("edges", []):
+            subgraph.add_edge(Edge(source=edge_def["source"], target=edge_def["target"], condition=edge_def.get("condition")))
+
+        result_state = await subgraph.run(
+            input_data=state.get("input"),
+            max_iterations=max_iterations,
+            state={"parent_state": state},
+        )
+        return {"workflow_id": workflow_id, "state": result_state}
+
+    async def _execute_loop_node(
+        self, node: Node, state: Dict[str, Any], max_iterations: int
+    ) -> Any:
+        """Execute a loop node by iterating until condition is met."""
+        condition_key = node.config.data.get("condition")
+        loop_limit = int(node.config.data.get("max_iterations", 5))
+        loop_iterations = 0
+        results = []
+
+        while loop_iterations < loop_limit:
+            loop_iterations += 1
+            state_key = f"loop_{node.id}_iteration"
+            state[state_key] = loop_iterations
+            results.append({"iteration": loop_iterations})
+            if condition_key and state.get(condition_key):
+                break
+            if state.get("iterations", 0) >= max_iterations:
+                break
+
+        return {"iterations": loop_iterations, "results": results}
+
     def to_dict(self) -> Dict[str, Any]:
         """Convert graph to dictionary representation.