agentmesh-lightning 2.3.0__py3-none-any.whl

agent_lightning_gov/__init__.py

@@ -0,0 +1,36 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""
+Agent-Lightning Governance Integration
+=======================================
+
+Provides kernel-level safety during Agent-Lightning RL training.
+
+Key components:
+- GovernedRunner: Agent-Lightning runner with policy enforcement
+- PolicyReward: Convert policy violations to RL penalties
+- FlightRecorderEmitter: Export audit logs to LightningStore
+- GovernedEnvironment: Training environment with governance constraints
+
+Example:
+    >>> from agent_lightning_gov import GovernedRunner, PolicyReward
+    >>> from agent_os import KernelSpace
+    >>> from agent_os.policies import SQLPolicy
+    >>>
+    >>> kernel = KernelSpace(policy=SQLPolicy())
+    >>> runner = GovernedRunner(kernel)
+    >>> reward_fn = PolicyReward(kernel, base_reward_fn=accuracy)
+"""
+
+from .emitter import FlightRecorderEmitter
+from .environment import GovernedEnvironment
+from .reward import PolicyReward, policy_penalty
+from .runner import GovernedRunner
+
+__all__ = [
+    "GovernedRunner",
+    "PolicyReward",
+    "policy_penalty",
+    "FlightRecorderEmitter",
+    "GovernedEnvironment",
+]

agent_lightning_gov/emitter.py

@@ -0,0 +1,314 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""
+FlightRecorderEmitter - Export Audit Logs to LightningStore
+============================================================
+
+Adapts Agent OS Flight Recorder logs to Agent-Lightning's
+span format for unified training and audit trail.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from collections.abc import Iterator
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class LightningSpan:
+    """
+    Span compatible with Agent-Lightning's LightningStore.
+
+    Maps Agent OS Flight Recorder entries to the span format
+    expected by Agent-Lightning's training and analysis tools.
+    """
+    span_id: str
+    trace_id: str
+    name: str
+    start_time: datetime
+    end_time: datetime | None = None
+    attributes: dict[str, Any] = field(default_factory=dict)
+    events: list[dict[str, Any]] = field(default_factory=list)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for serialization."""
+        return {
+            "span_id": self.span_id,
+            "trace_id": self.trace_id,
+            "name": self.name,
+            "start_time": self.start_time.isoformat(),
+            "end_time": self.end_time.isoformat() if self.end_time else None,
+            "attributes": self.attributes,
+            "events": self.events,
+        }
+
+    def to_json(self) -> str:
+        """Convert to JSON string."""
+        return json.dumps(self.to_dict())
+
+
+class FlightRecorderEmitter:
+    """
+    Emits Agent OS Flight Recorder entries to Agent-Lightning.
+
+    This adapter enables:
+    1. Complete audit trail from training to production
+    2. RL algorithms learning from policy violations
+    3. Compliance-friendly training logs
+
+    Example:
+        >>> from agent_os import FlightRecorder
+        >>> from agent_os.integrations.agent_lightning import FlightRecorderEmitter
+        >>>
+        >>> recorder = FlightRecorder()
+        >>> emitter = FlightRecorderEmitter(recorder)
+        >>>
+        >>> # Emit all logs to LightningStore
+        >>> emitter.emit_to_store(store)
+        >>>
+        >>> # Or stream continuously
+        >>> async for span in emitter.stream():
+        ...     store.emit_span(span)
+    """
+
+    def __init__(
+        self,
+        flight_recorder: Any,  # FlightRecorder
+        *,
+        include_policy_checks: bool = True,
+        include_signals: bool = True,
+        include_tool_calls: bool = True,
+        trace_id_prefix: str = "agentos",
+    ):
+        """
+        Initialize the emitter.
+
+        Args:
+            flight_recorder: Agent OS FlightRecorder instance
+            include_policy_checks: Include policy check spans
+            include_signals: Include signal dispatch spans
+            include_tool_calls: Include tool call spans
+            trace_id_prefix: Prefix for generated trace IDs
+        """
+        self.recorder = flight_recorder
+        self.include_policy_checks = include_policy_checks
+        self.include_signals = include_signals
+        self.include_tool_calls = include_tool_calls
+        self.trace_id_prefix = trace_id_prefix
+
+        self._emitted_count = 0
+        self._last_position = 0
+
+    def _convert_entry(self, entry: Any) -> LightningSpan | None:
+        """Convert a Flight Recorder entry to a Lightning span."""
+        entry_type = getattr(entry, 'type', getattr(entry, 'entry_type', 'unknown'))
+
+        # Filter by entry type
+        if entry_type == 'policy_check' and not self.include_policy_checks:
+            return None
+        if entry_type == 'signal' and not self.include_signals:
+            return None
+        if entry_type == 'tool_call' and not self.include_tool_calls:
+            return None
+
+        # Extract common fields
+        span_id = getattr(entry, 'id', getattr(entry, 'entry_id', str(self._emitted_count)))
+        timestamp = getattr(entry, 'timestamp', datetime.now(timezone.utc))
+        agent_id = getattr(entry, 'agent_id', 'unknown')
+
+        # Build attributes
+        attributes = {
+            "agent_os.entry_type": entry_type,
+            "agent_os.agent_id": agent_id,
+        }
+
+        # Add type-specific attributes
+        if entry_type == 'policy_check':
+            attributes.update({
+                "agent_os.policy_name": getattr(entry, 'policy_name', 'unknown'),
+                "agent_os.policy_result": getattr(entry, 'result', 'unknown'),
+                "agent_os.policy_violated": getattr(entry, 'violated', False),
+            })
+        elif entry_type == 'signal':
+            attributes.update({
+                "agent_os.signal_type": getattr(entry, 'signal', 'unknown'),
+                "agent_os.signal_target": getattr(entry, 'target', 'unknown'),
+            })
+        elif entry_type == 'tool_call':
+            attributes.update({
+                "agent_os.tool_name": getattr(entry, 'tool_name', 'unknown'),
+                "agent_os.tool_args": str(getattr(entry, 'args', {}))[:1000],  # Truncate
+                "agent_os.tool_result": str(getattr(entry, 'result', None))[:1000],
+            })
+
+        # Copy any additional attributes
+        if hasattr(entry, 'metadata'):
+            for key, value in entry.metadata.items():
+                attributes[f"agent_os.{key}"] = value
+
+        return LightningSpan(
+            span_id=str(span_id),
+            trace_id=f"{self.trace_id_prefix}-{agent_id}",
+            name=f"agent_os.{entry_type}",
+            start_time=timestamp,
+            end_time=timestamp,
+            attributes=attributes,
+        )
+
+    def get_spans(self) -> list[LightningSpan]:
+        """
+        Get all Flight Recorder entries as Lightning spans.
+
+        Returns:
+            List of LightningSpan objects
+        """
+        spans = []
+
+        # Get entries from recorder
+        if hasattr(self.recorder, 'get_entries'):
+            entries = self.recorder.get_entries()
+        elif hasattr(self.recorder, 'entries'):
+            entries = self.recorder.entries
+        elif hasattr(self.recorder, 'get_logs'):
+            entries = self.recorder.get_logs()
+        else:
+            logger.warning("Flight recorder has no recognized entry accessor")
+            return []
+
+        for entry in entries:
+            span = self._convert_entry(entry)
+            if span:
+                spans.append(span)
+                self._emitted_count += 1
+
+        return spans
+
+    def get_new_spans(self) -> list[LightningSpan]:
+        """
+        Get only new entries since last call.
+
+        Returns:
+            List of new LightningSpan objects
+        """
+        all_spans = self.get_spans()
+        new_spans = all_spans[self._last_position:]
+        self._last_position = len(all_spans)
+        return new_spans
+
+    async def stream(self) -> Iterator[LightningSpan]:
+        """
+        Stream spans as they are recorded.
+
+        Yields:
+            LightningSpan objects as they become available
+        """
+        import asyncio
+
+        while True:
+            new_spans = self.get_new_spans()
+            for span in new_spans:
+                yield span
+
+            # Wait before checking again
+            await asyncio.sleep(0.1)
+
+    def emit_to_store(self, store: Any) -> int:
+        """
+        Emit all spans to a LightningStore.
+
+        Args:
+            store: Agent-Lightning LightningStore instance
+
+        Returns:
+            Number of spans emitted
+        """
+        spans = self.get_spans()
+
+        for span in spans:
+            try:
+                if hasattr(store, 'emit_span'):
+                    store.emit_span(span.to_dict())
+                elif hasattr(store, 'add_span'):
+                    store.add_span(span.to_dict())
+                else:
+                    logger.warning("Store has no recognized span emitter")
+                    break
+            except Exception as e:
+                logger.error(f"Failed to emit span: {e}")
+
+        logger.info(f"Emitted {len(spans)} spans to LightningStore")
+        return len(spans)
+
+    def export_to_file(self, filepath: str) -> int:
+        """
+        Export spans to a JSON file.
+
+        Args:
+            filepath: Path to output file
+
+        Returns:
+            Number of spans exported
+        """
+        spans = self.get_spans()
+
+        with open(filepath, 'w') as f:
+            json.dump([s.to_dict() for s in spans], f, indent=2)
+
+        logger.info(f"Exported {len(spans)} spans to {filepath}")
+        return len(spans)
+
+    def get_violation_summary(self) -> dict[str, Any]:
+        """
+        Get summary of policy violations from recorded entries.
+
+        Returns:
+            Summary dictionary with violation statistics
+        """
+        spans = self.get_spans()
+
+        violations = [
+            s for s in spans
+            if s.attributes.get("agent_os.policy_violated", False)
+        ]
+
+        policies_violated = {}
+        for v in violations:
+            policy = v.attributes.get("agent_os.policy_name", "unknown")
+            policies_violated[policy] = policies_violated.get(policy, 0) + 1
+
+        return {
+            "total_entries": len(spans),
+            "total_violations": len(violations),
+            "violation_rate": len(violations) / max(len(spans), 1),
+            "policies_violated": policies_violated,
+        }
+
+    def get_stats(self) -> dict[str, Any]:
+        """Get emitter statistics."""
+        return {
+            "emitted_count": self._emitted_count,
+            "last_position": self._last_position,
+        }
+
+
+def create_emitter(
+    flight_recorder: Any,
+    **kwargs: Any,
+) -> FlightRecorderEmitter:
+    """
+    Factory function to create a FlightRecorderEmitter.
+
+    Args:
+        flight_recorder: Agent OS FlightRecorder
+        **kwargs: Additional configuration
+
+    Returns:
+        Configured FlightRecorderEmitter
+    """
+    return FlightRecorderEmitter(flight_recorder, **kwargs)

agent_lightning_gov/environment.py

@@ -0,0 +1,330 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""
+GovernedEnvironment - Training Environment with Governance
+==========================================================
+
+Wraps an Agent OS kernel as a training environment for
+Agent-Lightning's RL algorithms.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Any, Callable, Generic, TypeVar
+
+logger = logging.getLogger(__name__)
+
+T_state = TypeVar("T_state")
+T_action = TypeVar("T_action")
+
+
+@dataclass
+class EnvironmentConfig:
+    """Configuration for governed training environment."""
+
+    # Maximum steps per episode
+    max_steps: int = 100
+
+    # Penalty for policy violations
+    violation_penalty: float = -10.0
+
+    # Terminate episode on critical violation
+    terminate_on_critical: bool = True
+
+    # Reward shaping
+    step_penalty: float = -0.1  # Small penalty per step to encourage efficiency
+    success_bonus: float = 10.0
+
+    # Reset behavior
+    reset_kernel_state: bool = True
+
+
+@dataclass
+class EnvironmentState:
+    """State of the governed environment."""
+
+    step_count: int = 0
+    total_reward: float = 0.0
+    violations: list = field(default_factory=list)
+    terminated: bool = False
+    truncated: bool = False
+    info: dict = field(default_factory=dict)
+
+
+class GovernedEnvironment(Generic[T_state, T_action]):
+    """
+    RL training environment with Agent OS governance.
+
+    This environment wraps an Agent OS kernel and can be used
+    directly with Agent-Lightning or other RL frameworks.
+
+    The environment:
+    1. Enforces policies on each action
+    2. Converts violations to negative rewards
+    3. Optionally terminates on critical violations
+    4. Tracks compliance metrics during training
+
+    Example:
+        >>> from agent_os import KernelSpace
+        >>> from agent_os.policies import SQLPolicy
+        >>>
+        >>> kernel = KernelSpace(policy=SQLPolicy())
+        >>> env = GovernedEnvironment(kernel)
+        >>>
+        >>> state = env.reset()
+        >>> while not env.terminated:
+        ...     action = agent.get_action(state)
+        ...     state, reward, terminated, truncated, info = env.step(action)
+
+    Compatible with:
+    - Agent-Lightning trainers
+    - OpenAI Gym / Gymnasium
+    - Stable Baselines3
+    - Any environment with step/reset interface
+    """
+
+    def __init__(
+        self,
+        kernel: Any,  # KernelSpace
+        *,
+        task_generator: Callable[[], T_state] | None = None,
+        reward_fn: Callable[[T_state, T_action, Any], float] | None = None,
+        config: EnvironmentConfig | None = None,
+    ):
+        """
+        Initialize the governed environment.
+
+        Args:
+            kernel: Agent OS KernelSpace with loaded policies
+            task_generator: Optional function to generate initial states
+            reward_fn: Optional custom reward function
+            config: Environment configuration
+        """
+        self.kernel = kernel
+        self.task_generator = task_generator
+        self.reward_fn = reward_fn or self._default_reward
+        self.config = config or EnvironmentConfig()
+
+        # Current episode state
+        self._state = EnvironmentState()
+        self._current_task: T_state | None = None
+        self._current_violations: list = []
+
+        # Metrics
+        self._total_episodes = 0
+        self._total_steps = 0
+        self._total_violations = 0
+        self._successful_episodes = 0
+
+        # Set up kernel hooks
+        self._setup_hooks()
+
+        logger.info("GovernedEnvironment initialized")
+
+    def _setup_hooks(self) -> None:
+        """Set up hooks to capture violations."""
+        if hasattr(self.kernel, 'on_policy_violation'):
+            self.kernel.on_policy_violation(self._handle_violation)
+
+    def _handle_violation(
+        self,
+        policy_name: str,
+        description: str,
+        severity: str,
+        blocked: bool,
+    ) -> None:
+        """Handle policy violation during step."""
+        violation = {
+            "policy": policy_name,
+            "description": description,
+            "severity": severity,
+            "blocked": blocked,
+            "step": self._state.step_count,
+            "timestamp": datetime.now(timezone.utc).isoformat(),
+        }
+        self._current_violations.append(violation)
+        self._state.violations.append(violation)
+        self._total_violations += 1
+
+    def _default_reward(
+        self,
+        state: T_state,
+        action: T_action,
+        result: Any,
+    ) -> float:
+        """Default reward function."""
+        # Base reward for task completion
+        if result is not None:
+            return 1.0
+        return 0.0
+
+    def reset(
+        self,
+        *,
+        seed: int | None = None,
+        options: dict[str, Any] | None = None,
+    ) -> tuple[T_state, dict[str, Any]]:
+        """
+        Reset environment for new episode.
+
+        Args:
+            seed: Random seed (for compatibility)
+            options: Additional options
+
+        Returns:
+            Tuple of (initial_state, info_dict)
+        """
+        # Reset episode state
+        self._state = EnvironmentState()
+        self._current_violations = []
+        self._total_episodes += 1
+
+        # Reset kernel state if configured
+        if self.config.reset_kernel_state and hasattr(self.kernel, 'reset'):
+            self.kernel.reset()
+
+        # Generate initial task
+        if self.task_generator:
+            self._current_task = self.task_generator()
+        else:
+            self._current_task = None
+
+        info = {
+            "episode": self._total_episodes,
+            "kernel_policies": self._get_policy_names(),
+        }
+
+        return self._current_task, info
+
+    def step(
+        self,
+        action: T_action,
+    ) -> tuple[T_state, float, bool, bool, dict[str, Any]]:
+        """
+        Execute one step in the environment.
+
+        Args:
+            action: Agent's action
+
+        Returns:
+            Tuple of (next_state, reward, terminated, truncated, info)
+        """
+        self._current_violations = []
+        self._state.step_count += 1
+        self._total_steps += 1
+
+        # Execute action through kernel
+        try:
+            if hasattr(self.kernel, 'execute'):
+                result = self.kernel.execute(action)
+            else:
+                result = action  # No kernel execution, passthrough
+
+            success = True
+        except Exception as e:
+            logger.error(f"Step failed: {e}")
+            result = None
+            success = False
+
+        # Calculate reward
+        reward = self.reward_fn(self._current_task, action, result)
+
+        # Apply step penalty
+        reward += self.config.step_penalty
+
+        # Apply violation penalties
+        for violation in self._current_violations:
+            penalty = self.config.violation_penalty
+            if violation["severity"] == "critical":
+                penalty *= 10
+            elif violation["severity"] == "high":
+                penalty *= 5
+            reward += penalty
+
+        self._state.total_reward += reward
+
+        # Check termination conditions
+        terminated = False
+        truncated = False
+
+        # Terminate on critical violation if configured
+        if self.config.terminate_on_critical:
+            if any(v["severity"] == "critical" for v in self._current_violations):
+                terminated = True
+                logger.info("Episode terminated due to critical violation")
+
+        # Truncate on max steps
+        if self._state.step_count >= self.config.max_steps:
+            truncated = True
+
+        # Mark success
+        if success and not self._current_violations:
+            reward += self.config.success_bonus
+            self._successful_episodes += 1
+
+        self._state.terminated = terminated
+        self._state.truncated = truncated
+
+        info = {
+            "violations": self._current_violations,
+            "step": self._state.step_count,
+            "total_reward": self._state.total_reward,
+            "success": success,
+        }
+        self._state.info = info
+
+        return self._current_task, reward, terminated, truncated, info
+
+    def _get_policy_names(self) -> list[str]:
+        """Get names of loaded policies."""
+        if hasattr(self.kernel, 'get_policies'):
+            return [p.name for p in self.kernel.get_policies()]
+        if hasattr(self.kernel, 'policies'):
+            return [getattr(p, 'name', str(p)) for p in self.kernel.policies]
+        return []
+
+    @property
+    def terminated(self) -> bool:
+        """Whether current episode is terminated."""
+        return self._state.terminated or self._state.truncated
+
+    def get_metrics(self) -> dict[str, Any]:
+        """Get environment metrics."""
+        return {
+            "total_episodes": self._total_episodes,
+            "total_steps": self._total_steps,
+            "total_violations": self._total_violations,
+            "successful_episodes": self._successful_episodes,
+            "success_rate": self._successful_episodes / max(self._total_episodes, 1),
+            "violations_per_episode": self._total_violations / max(self._total_episodes, 1),
+            "steps_per_episode": self._total_steps / max(self._total_episodes, 1),
+        }
+
+    def close(self) -> None:
+        """Clean up environment resources."""
+        logger.info(f"Environment closed. Metrics: {self.get_metrics()}")
+
+
+def create_governed_env(
+    kernel: Any,
+    **kwargs: Any,
+) -> GovernedEnvironment:
+    """
+    Factory function to create a GovernedEnvironment.
+
+    Args:
+        kernel: Agent OS KernelSpace
+        **kwargs: Environment configuration
+
+    Returns:
+        Configured GovernedEnvironment
+    """
+    config = EnvironmentConfig()
+    for key, value in kwargs.items():
+        if hasattr(config, key):
+            setattr(config, key, value)
+
+    return GovernedEnvironment(kernel, config=config)