lambdagent/lambdagent_guard/crewai.py

"""
lambdagent_guard.crewai — CrewAI guard wrapper (I12).

Usage:
    from lambdagent_guard import guard_crewai

    crew = Crew(agents=[researcher, writer], tasks=[...])
    guarded = guard_crewai(crew, cost_budget=10.0)
    result = guarded.kickoff()

Catches: CrewAI #3847 (max_iter broken), #737 (tool result lost), #1355 (10x cost)
"""

from __future__ import annotations
from typing import Any, Callable, Optional
from .core import GuardConfig, RuntimeMonitor, run_compile_checks
import warnings


def guard_crewai(
    crew,
    cost_budget: float = float("inf"),
    parallel_safety: bool = True,
    terminate_check: bool = True,
    loop_detection: bool = True,
    cost_alert: Optional[Callable[[float], None]] = None,
):
    """
    Wrap a CrewAI Crew with lambdagent safety guards.

    Args:
        crew: CrewAI Crew object
        cost_budget: Max cost in USD
        parallel_safety: Check agent store independence
        terminate_check: Verify termination conditions exist
    """
    guard_cfg = GuardConfig(
        cost_budget=cost_budget,
        parallel_safety=parallel_safety,
        terminate_check=terminate_check,
        loop_detection=loop_detection,
        cost_alert=cost_alert,
    )

    # Phase 1: Compile-time checks
    try:
        from lambdagent.extractors import extract_config
        config = extract_config(crew, framework="crewai")
        issues = run_compile_checks(config, guard_cfg)
        for issue in issues:
            warnings.warn(f"[lambdagent-guard] {issue}")
    except Exception as e:
        warnings.warn(f"[lambdagent-guard] Compile check skipped: {e}")

    # Phase 2: Runtime hooks — wrap kickoff()
    monitor = RuntimeMonitor(guard_cfg)
    original_kickoff = crew.kickoff

    def guarded_kickoff(*args, **kwargs):
        # Hook into each agent's execution
        for agent in getattr(crew, 'agents', []):
            _hook_crewai_agent(agent, monitor)

        result = original_kickoff(*args, **kwargs)
        return result

    crew.kickoff = guarded_kickoff
    crew._lambdagent_monitor = monitor

    return crew


def _hook_crewai_agent(agent, monitor: RuntimeMonitor):
    """Hook into a CrewAI agent's execute_task if available."""
    if hasattr(agent, 'execute_task'):
        original = agent.execute_task

        def guarded_execute(*args, **kwargs):
            result = original(*args, **kwargs)
            monitor.on_step({
                "output": str(result)[:500],
                "tokens": 0,
            })
            return result

        agent.execute_task = guarded_execute