agent_os_kernel 3.1.0__py3-none-any.whl

agent_kernel/simulator.py

@@ -0,0 +1,449 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""
+Path simulation system to test alternative solutions.
+"""
+
+import logging
+import uuid
+from typing import List, Dict, Any, Optional
+
+from .models import FailureAnalysis, SimulationResult, DiagnosisJSON, ShadowAgentResult, AgentFailure, CognitiveGlitch
+
+logger = logging.getLogger(__name__)
+
+
+class ShadowAgent:
+    """
+    Shadow Agent for counterfactual simulation.
+    
+    Replays the user prompt with an injected hint in a sandbox environment.
+    This is "The Scientist" that verifies if the hint actually fixes the problem.
+    """
+    
+    def __init__(self):
+        self.execution_history: List[ShadowAgentResult] = []
+    
+    def replay_with_hint(
+        self,
+        original_prompt: str,
+        hint: str,
+        diagnosis: DiagnosisJSON,
+        failure: AgentFailure
+    ) -> ShadowAgentResult:
+        """
+        Replay the original prompt with an injected hint.
+        
+        This simulates the agent executing with additional context/guidance.
+        
+        Args:
+            original_prompt: Original user prompt that led to failure
+            hint: Hint to inject based on diagnosis
+            diagnosis: Cognitive glitch diagnosis
+            failure: Original failure for context
+            
+        Returns:
+            ShadowAgentResult with execution outcome
+        """
+        shadow_id = f"shadow-{uuid.uuid4().hex[:8]}"
+        
+        logger.info(f"Shadow agent {shadow_id} replaying with hint")
+        
+        # Construct modified prompt with hint
+        modified_prompt = f"{original_prompt}\n\n{hint}"
+        
+        # Simulate execution (in real system, this would run actual agent in sandbox)
+        execution_success, output, reasoning, action = self._simulate_execution(
+            modified_prompt, diagnosis, failure
+        )
+        
+        # Verify the fix
+        verified = self._verify_fix(execution_success, action, failure)
+        
+        result = ShadowAgentResult(
+            shadow_id=shadow_id,
+            original_prompt=original_prompt,
+            injected_hint=hint,
+            modified_prompt=modified_prompt,
+            execution_success=execution_success,
+            output=output,
+            reasoning_chain=reasoning,
+            action_taken=action,
+            verified=verified
+        )
+        
+        self.execution_history.append(result)
+        logger.info(f"Shadow execution complete. Success: {execution_success}, Verified: {verified}")
+        
+        return result
+    
+    def _simulate_execution(
+        self,
+        prompt: str,
+        diagnosis: DiagnosisJSON,
+        failure: AgentFailure
+    ) -> tuple:
+        """
+        Simulate agent execution with the modified prompt.
+        
+        In a real system, this would:
+        1. Spin up a sandboxed agent instance
+        2. Inject the hint into the system prompt or context
+        3. Execute the agent with the original user prompt
+        4. Capture the reasoning chain and action
+        5. Check if the action would pass control plane checks
+        """
+        # Simulate reasoning chain with hint consideration
+        reasoning = [
+            "Parse user request",
+            f"Consider hint: {diagnosis.hint[:50]}...",
+            "Validate assumptions against provided context",
+            "Construct safe action"
+        ]
+        
+        # Simulate a corrected action
+        # In reality, this would be the actual agent's output
+        if diagnosis.cognitive_glitch == CognitiveGlitch.HALLUCINATION:
+            # Agent would verify schema
+            action = {
+                "action": "execute_with_validation",
+                "validation": "schema_check_passed",
+                "safe_mode": True
+            }
+            success = True
+            output = "Action validated and executed successfully"
+        elif diagnosis.cognitive_glitch == CognitiveGlitch.PERMISSION_ERROR:
+            # Agent would check permissions first
+            action = {
+                "action": "check_permissions_then_execute",
+                "permission_validation": True
+            }
+            success = True
+            output = "Permissions validated, action executed"
+        else:
+            # Generic safe action
+            action = {
+                "action": "safe_execute",
+                "hint_applied": True
+            }
+            # Success based on diagnosis confidence (deterministic)
+            success = diagnosis.confidence > 0.7
+            output = "Action executed with safety checks" if success else "Action still failed"
+        
+        return success, output, reasoning, action
+    
+    def _verify_fix(self, success: bool, action: Optional[Dict], failure: AgentFailure) -> bool:
+        """
+        Verify that the fix actually works.
+        
+        This is the key validation step - confirming the hint flips the outcome
+        from Fail to Pass.
+        """
+        if not success:
+            return False
+        
+        # Check if action has safety mechanisms that original lacked
+        if action:
+            has_validation = any(
+                key in action for key in ["validation", "permission_validation", "schema_check", "safe_mode"]
+            )
+            return has_validation
+        
+        return success
+
+
+class PathSimulator:
+    """Simulates alternative paths to avoid failures."""
+    
+    def __init__(self):
+        self.simulation_history: List[SimulationResult] = []
+        self.shadow_agent = ShadowAgent()
+        self.mcts_iterations = 5  # MCTS search iterations
+    
+    def simulate(self, analysis: FailureAnalysis) -> SimulationResult:
+        """
+        Simulate an alternative path based on failure analysis.
+        
+        Args:
+            analysis: The failure analysis containing suggested fixes
+            
+        Returns:
+            SimulationResult with the alternative path and predicted outcome
+        """
+        logger.info(f"Simulating alternative path for agent {analysis.failure.agent_id}")
+        
+        # Generate simulation ID
+        simulation_id = str(uuid.uuid4())
+        
+        # Build alternative path from suggested fixes
+        alternative_path = self._build_alternative_path(analysis)
+        
+        # Predict outcome
+        expected_outcome = self._predict_outcome(analysis, alternative_path)
+        
+        # Calculate risk score
+        risk_score = self._calculate_risk(analysis, alternative_path)
+        
+        # Estimate success rate
+        estimated_success_rate = self._estimate_success_rate(analysis, risk_score)
+        
+        # Determine if simulation is successful
+        success = risk_score < 0.5 and estimated_success_rate > 0.7
+        
+        result = SimulationResult(
+            simulation_id=simulation_id,
+            success=success,
+            alternative_path=alternative_path,
+            expected_outcome=expected_outcome,
+            risk_score=risk_score,
+            estimated_success_rate=estimated_success_rate
+        )
+        
+        self.simulation_history.append(result)
+        
+        if success:
+            logger.info(f"Simulation successful. Success rate: {estimated_success_rate:.2f}, Risk: {risk_score:.2f}")
+        else:
+            logger.warning(f"Simulation failed. Success rate: {estimated_success_rate:.2f}, Risk: {risk_score:.2f}")
+        
+        return result
+    
+    def simulate_counterfactual(
+        self,
+        diagnosis: DiagnosisJSON,
+        failure: AgentFailure
+    ) -> ShadowAgentResult:
+        """
+        Counterfactual simulation using Shadow Agent.
+        
+        This is "The Simulator" - it replays the user prompt but injects a hint
+        based on the DiagnosisJSON. Uses MCTS-like approach to find minimal
+        change required to flip outcome from Fail to Pass.
+        
+        Args:
+            diagnosis: Cognitive glitch diagnosis with hint
+            failure: Original failure with trace
+            
+        Returns:
+            ShadowAgentResult showing if hint fixes the problem
+        """
+        logger.info("Starting counterfactual simulation with Shadow Agent")
+        
+        if not failure.failure_trace:
+            logger.warning("No failure trace available for counterfactual simulation")
+            # Create a dummy result
+            return ShadowAgentResult(
+                shadow_id="shadow-no-trace",
+                original_prompt="No prompt available",
+                injected_hint=diagnosis.hint,
+                modified_prompt="No prompt available",
+                execution_success=False,
+                output="No trace available for simulation",
+                reasoning_chain=[],
+                action_taken=None,
+                verified=False
+            )
+        
+        # Use MCTS-inspired approach: try multiple hint variations
+        best_result = self._mcts_search_minimal_hint(
+            failure.failure_trace.user_prompt,
+            diagnosis,
+            failure
+        )
+        
+        return best_result
+    
+    def _mcts_search_minimal_hint(
+        self,
+        prompt: str,
+        diagnosis: DiagnosisJSON,
+        failure: AgentFailure
+    ) -> ShadowAgentResult:
+        """
+        MCTS-inspired search for minimal hint that fixes the problem.
+        
+        This searches for the minimal change required to flip the outcome.
+        In full MCTS, we'd build a tree of hint variations and explore/exploit,
+        but here we do a simplified version with multiple trials.
+        """
+        logger.info(f"MCTS search across {self.mcts_iterations} iterations")
+        
+        hint_variations = self._generate_hint_variations(diagnosis.hint)
+        results = []
+        
+        for i, hint in enumerate(hint_variations[:self.mcts_iterations]):
+            logger.debug(f"MCTS iteration {i+1}: Testing hint variation")
+            result = self.shadow_agent.replay_with_hint(
+                prompt, hint, diagnosis, failure
+            )
+            results.append(result)
+            
+            # Early exit if we find a verified solution
+            if result.verified and result.execution_success:
+                logger.info(f"Found verified solution at iteration {i+1}")
+                break
+        
+        # Select best result (verified + successful, or highest success)
+        best = max(
+            results,
+            key=lambda r: (r.verified, r.execution_success, len(r.reasoning_chain))
+        )
+        
+        logger.info(f"MCTS search complete. Best result verified: {best.verified}")
+        return best
+    
+    def _generate_hint_variations(self, base_hint: str) -> List[str]:
+        """
+        Generate variations of the hint for MCTS exploration.
+        
+        This finds different ways to provide the same guidance,
+        searching for the minimal effective intervention.
+        """
+        variations = [base_hint]  # Original hint
+        
+        # Variation 1: More concise
+        if len(base_hint) > 50:
+            concise = base_hint.split(".")[0] + "."
+            variations.append(concise)
+        
+        # Variation 2: More explicit
+        explicit = base_hint + " Double-check all assumptions."
+        variations.append(explicit)
+        
+        # Variation 3: Focus on specific action
+        if "validate" in base_hint.lower():
+            variations.append("VALIDATION REQUIRED: " + base_hint)
+        
+        # Variation 4: Minimal version
+        if "HINT:" in base_hint:
+            minimal = base_hint.replace("HINT: ", "")
+            variations.append(minimal)
+        
+        return variations
+    
+    def _build_alternative_path(self, analysis: FailureAnalysis) -> List[Dict[str, Any]]:
+        """Build an alternative execution path from suggested fixes."""
+        path = []
+        failure = analysis.failure
+        
+        # Add validation step for control plane blocks
+        if failure.failure_type.value == "blocked_by_control_plane":
+            path.append({
+                "step": 1,
+                "action": "validate_permissions",
+                "description": "Check permissions before attempting action",
+                "params": {
+                    "resource": failure.context.get("resource", "unknown"),
+                    "action": failure.context.get("action", "unknown")
+                }
+            })
+            
+            path.append({
+                "step": 2,
+                "action": "request_authorization",
+                "description": "Request proper authorization if needed",
+                "params": {
+                    "required_permission": "resource_access"
+                }
+            })
+            
+            path.append({
+                "step": 3,
+                "action": "safe_execute",
+                "description": "Execute action with safety checks",
+                "params": {
+                    "original_action": failure.context.get("action", "unknown"),
+                    "safety_mode": "enabled"
+                }
+            })
+        
+        # Add timeout handling for timeout failures
+        elif failure.failure_type.value == "timeout":
+            path.append({
+                "step": 1,
+                "action": "set_timeout",
+                "description": "Configure appropriate timeout",
+                "params": {"timeout_seconds": 30}
+            })
+            
+            path.append({
+                "step": 2,
+                "action": "add_progress_monitoring",
+                "description": "Add progress monitoring",
+                "params": {"check_interval_seconds": 5}
+            })
+            
+            path.append({
+                "step": 3,
+                "action": "execute_with_timeout",
+                "description": "Execute with timeout handling",
+                "params": {"allow_partial_results": True}
+            })
+        
+        # Generic alternative path for other failures
+        else:
+            for i, fix in enumerate(analysis.suggested_fixes[:3], 1):
+                path.append({
+                    "step": i,
+                    "action": "apply_fix",
+                    "description": fix,
+                    "params": {"fix": fix}
+                })
+        
+        return path
+    
+    def _predict_outcome(self, analysis: FailureAnalysis, alternative_path: List[Dict[str, Any]]) -> str:
+        """Predict the outcome of executing the alternative path."""
+        failure = analysis.failure
+        
+        if failure.failure_type.value == "blocked_by_control_plane":
+            return "Action will be executed with proper authorization and safety checks"
+        elif failure.failure_type.value == "timeout":
+            return "Operation will complete within timeout with progress monitoring"
+        else:
+            return f"Failure {failure.failure_type.value} will be prevented by applying suggested fixes"
+    
+    def _calculate_risk(self, analysis: FailureAnalysis, alternative_path: List[Dict[str, Any]]) -> float:
+        """Calculate risk score for the alternative path."""
+        risk = 0.3  # Base risk
+        
+        # Lower risk if confidence is high
+        risk -= (analysis.confidence_score * 0.2)
+        
+        # Lower risk if we have multiple steps (more thorough)
+        if len(alternative_path) >= 3:
+            risk -= 0.1
+        
+        # Higher risk for unknown failure types
+        if analysis.failure.failure_type.value == "unknown":
+            risk += 0.2
+        
+        return max(0.0, min(1.0, risk))
+    
+    def _estimate_success_rate(self, analysis: FailureAnalysis, risk_score: float) -> float:
+        """Estimate success rate of the alternative path."""
+        # Base success rate from confidence
+        success_rate = analysis.confidence_score
+        
+        # Adjust based on risk
+        success_rate = success_rate * (1.0 - risk_score * 0.5)
+        
+        # Bonus for having similar failures (we've seen this before)
+        if len(analysis.similar_failures) > 0:
+            success_rate += 0.1
+        
+        return max(0.0, min(1.0, success_rate))
+    
+    def get_best_simulation(self, simulations: List[SimulationResult]) -> SimulationResult:
+        """Get the best simulation from a list based on success rate and risk."""
+        if not simulations:
+            raise ValueError("No simulations provided")
+        
+        # Sort by estimated success rate (desc) and risk score (asc)
+        sorted_sims = sorted(
+            simulations,
+            key=lambda s: (s.estimated_success_rate, -s.risk_score),
+            reverse=True
+        )
+        
+        return sorted_sims[0]

agent_kernel/teacher.py

@@ -0,0 +1,85 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""
+Teacher - Simplified reference implementation for the Shadow Teacher.
+
+This is a reference implementation showing the core concept of using a
+"Teacher Model" (stronger reasoning model) to diagnose why an agent failed.
+
+The production implementation is integrated throughout the analyzer.py and
+completeness_auditor.py modules with full trace capture and cognitive diagnosis.
+"""
+
+
+def _sanitize_input(text: str, max_length: int = 1000) -> str:
+    """
+    Sanitize input to prevent prompt injection.
+    
+    Args:
+        text: Input text to sanitize
+        max_length: Maximum allowed length
+        
+    Returns:
+        Sanitized text
+    """
+    if not text:
+        return ""
+    
+    # Truncate to max length
+    text = str(text)[:max_length]
+    
+    # Remove potential prompt injection patterns
+    # In production, use more sophisticated sanitization
+    dangerous_patterns = ["ignore previous", "ignore all", "disregard", "new instructions"]
+    text_lower = text.lower()
+    for pattern in dangerous_patterns:
+        if pattern in text_lower:
+            text = text.replace(pattern, "[FILTERED]")
+    
+    return text
+
+
+async def diagnose_failure(prompt, failed_response, tool_trace):
+    """
+    Uses a 'Reasoning Model' (e.g., o1 or Claude 3.5 Sonnet) 
+    to find the Root Cause.
+    
+    Args:
+        prompt: The original task/prompt that failed
+        failed_response: The agent's failed response
+        tool_trace: Trace of tools/actions the agent attempted
+        
+    Returns:
+        dict: Diagnosis with cause and lesson_patch
+    """
+    # Sanitize inputs to prevent prompt injection
+    safe_prompt = _sanitize_input(prompt)
+    safe_response = _sanitize_input(failed_response)
+    safe_trace = _sanitize_input(tool_trace)
+    
+    teacher_prompt = f"""
+    The Agent failed to complete this task: '{safe_prompt}'.
+    
+    Agent Output: {safe_response}
+    Tool Trace: {safe_trace}
+    
+    Task:
+    1. Did the agent try hard enough? (Laziness)
+    2. Did the agent hallucinate a tool parameter? (Skill Issue)
+    3. Write a 1-sentence 'Lesson' that fixes this specific error.
+    
+    Output Format: JSON {{ "cause": "...", "lesson_patch": "..." }}
+    """
+    
+    # In production, this would call the "Expensive" Model only on failure
+    # For this reference implementation, we simulate the response
+    # diagnosis = await llm_client.generate(model="o1-preview", prompt=teacher_prompt)
+    
+    # Simulated diagnosis for reference
+    diagnosis = {
+        "cause": "Agent gave up without exhaustive search",
+        "lesson_patch": "Before reporting 'not found', check all data sources including archived partitions"
+    }
+    
+    return diagnosis

agent_kernel/triage.py

@@ -0,0 +1,152 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""
+Failure Triage Engine - Decides sync (JIT) vs async (batch) correction strategy.
+
+This is the missing component that sits between Failure Detection and Correction.
+It analyzes the context to determine if a failure should be fixed immediately
+(blocking the user) or asynchronously (returning error quickly, fixing later).
+"""
+
+from enum import Enum
+from typing import Dict, Any, Optional
+
+
+class FixStrategy(Enum):
+    """Strategy for fixing agent failures."""
+    SYNC_JIT = "jit_retry"       # High Latency, High Reliability - Fix NOW and wait
+    ASYNC_BATCH = "async_patch"  # Low Latency, Eventual Consistency - Fix LATER
+
+
+class FailureTriage:
+    """
+    Decision engine for routing failures to sync (JIT) or async (batch) correction.
+    
+    The triage engine applies three rules to determine criticality:
+    1. Safety/Write Operations → SYNC_JIT (must fix immediately)
+    2. High Effort Prompts → SYNC_JIT (user expects deep thinking)
+    3. Read/Query Operations → ASYNC_BATCH (save user time)
+    
+    This enables "thinking fast" (async) for trivial failures and
+    "thinking slow" (sync) for critical failures.
+    """
+    
+    def __init__(self, config: Optional[Dict[str, Any]] = None):
+        """
+        Initialize the triage engine.
+        
+        Args:
+            config: Optional configuration for custom critical tools and keywords
+        """
+        self.config = config or {}
+        
+        # Critical tools that require synchronous fixing (can be customized via config)
+        self.critical_tools = self.config.get("critical_tools", [
+            "delete_resource",
+            "update_db",
+            "execute_payment",
+            "drop_table",
+            "refund_user",
+            "delete_file",
+            "execute_sql",
+            "write_file",
+            "modify_permissions",
+            "delete_user"  # User deletion is critical
+        ])
+        
+        # Keywords indicating high-effort prompts requiring deep thinking
+        self.high_effort_keywords = self.config.get("high_effort_keywords", [
+            "carefully",
+            "critical",
+            "important",
+            "urgent",
+            "must",
+            "required",
+            "ensure"
+        ])
+    
+    def decide_strategy(
+        self,
+        prompt: str,
+        tool_name: Optional[str] = None,
+        user_metadata: Optional[Dict[str, Any]] = None,
+        context: Optional[Dict[str, Any]] = None
+    ) -> FixStrategy:
+        """
+        Decide whether to fix this failure sync (JIT) or async (batch).
+        
+        Decision Rules (in priority order):
+        1. Cognitive failures with full trace → SYNC_JIT (deep diagnosis needed)
+        2. Safety/Write Operations → SYNC_JIT
+        3. High Effort Prompts → SYNC_JIT  
+        4. VIP Users → SYNC_JIT
+        5. Default (Read/Query) → ASYNC_BATCH
+        
+        Args:
+            prompt: The user prompt that led to the failure
+            tool_name: Name of the tool that failed (if available)
+            user_metadata: Metadata about the user (e.g., VIP status)
+            context: Additional context about the failure
+            
+        Returns:
+            FixStrategy indicating sync (JIT) or async (batch) correction
+        """
+        # Rule 0: Cognitive failures with full trace (chain_of_thought + failed_action)
+        # These warrant immediate deep analysis with Shadow Teacher
+        if context:
+            has_chain = context.get("chain_of_thought") is not None
+            has_failed_action = context.get("failed_action") is not None
+            if has_chain and has_failed_action:
+                return FixStrategy.SYNC_JIT
+        
+        # Rule 1: Safety/Write Operations are always Critical
+        if tool_name and tool_name in self.critical_tools:
+            return FixStrategy.SYNC_JIT
+        
+        # Check context for critical actions (fallback if tool_name not provided)
+        if context:
+            action = context.get("action", "")
+            if action in self.critical_tools:
+                return FixStrategy.SYNC_JIT
+            
+            # Also check failed_action if present
+            failed_action = context.get("failed_action")
+            if failed_action and isinstance(failed_action, dict):
+                failed_action_name = failed_action.get("action", "")
+                if failed_action_name in self.critical_tools:
+                    return FixStrategy.SYNC_JIT
+        
+        # Rule 2: "High Effort" prompts request deep thinking
+        prompt_lower = prompt.lower()
+        if any(keyword in prompt_lower for keyword in self.high_effort_keywords):
+            return FixStrategy.SYNC_JIT
+        
+        # Rule 3: VIP users get priority treatment (optional)
+        if user_metadata and user_metadata.get("is_vip", False):
+            return FixStrategy.SYNC_JIT
+        
+        # Rule 4: Default to Async for "Read/Query" failures to save user time
+        return FixStrategy.ASYNC_BATCH
+    
+    def is_critical(
+        self,
+        prompt: str,
+        tool_name: Optional[str] = None,
+        user_metadata: Optional[Dict[str, Any]] = None,
+        context: Optional[Dict[str, Any]] = None
+    ) -> bool:
+        """
+        Convenience method to check if a failure is critical (needs SYNC_JIT).
+        
+        Args:
+            prompt: The user prompt
+            tool_name: Name of the tool that failed
+            user_metadata: User metadata
+            context: Additional context
+            
+        Returns:
+            True if critical (SYNC_JIT), False if non-critical (ASYNC_BATCH)
+        """
+        strategy = self.decide_strategy(prompt, tool_name, user_metadata, context)
+        return strategy == FixStrategy.SYNC_JIT