lambdagent/lambdagent_guard/autogen.py

"""
lambdagent_guard.autogen — AutoGen guard wrapper (I13).

Usage:
    from lambdagent_guard import guard_autogen

    manager = GroupChatManager(groupchat=chat)
    guarded = guard_autogen(manager, cost_budget=5.0)
    proxy.initiate_chat(guarded, message="Write quicksort")

Catches: AutoGen #108 (blank msg loop), #391 (TERMINATE bypass), #2702 (cost tracking)
"""

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


def guard_autogen(
    manager,
    cost_budget: float = float("inf"),
    empty_message_detection: bool = True,
    terminate_robustness: bool = True,
    loop_detection: bool = True,
    cost_alert: Optional[Callable[[float], None]] = None,
):
    """
    Wrap an AutoGen GroupChatManager with lambdagent safety guards.

    Args:
        manager: AutoGen GroupChatManager
        cost_budget: Max cost in USD
        empty_message_detection: Catch AutoGen #108 blank-message loops
        terminate_robustness: Don't rely on exact "TERMINATE" string matching
    """
    guard_cfg = GuardConfig(
        cost_budget=cost_budget,
        empty_message_detection=empty_message_detection,
        terminate_robustness=terminate_robustness,
        loop_detection=loop_detection,
        cost_alert=cost_alert,
    )

    # Phase 1: Compile-time checks
    try:
        from lambdagent.extractors import extract_config
        config = extract_config(manager, framework="autogen")

        # Check for fragile termination (AutoGen #108, #391)
        term_warning = config.get("_termination_warning")
        if term_warning:
            warnings.warn(f"[lambdagent-guard] {term_warning}")

        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
    monitor = RuntimeMonitor(guard_cfg)

    # Hook into GroupChatManager's run_chat or equivalent
    if hasattr(manager, 'run_chat'):
        original_run = manager.run_chat

        def guarded_run(messages, sender, config_override=None):
            for msg in messages if isinstance(messages, list) else [messages]:
                content = ""
                if isinstance(msg, dict):
                    content = msg.get("content", "")
                elif isinstance(msg, str):
                    content = msg

                monitor.on_step({
                    "output": str(content)[:500],
                    "tokens": 0,
                })

            return original_run(messages, sender, config_override)

        manager.run_chat = guarded_run

    # Hook terminate robustness: patch is_termination_msg if needed
    if terminate_robustness and hasattr(manager, 'groupchat'):
        chat = manager.groupchat
        original_term = getattr(chat, 'is_termination_msg', None)

        if original_term is not None:
            def robust_termination(msg):
                """Case-insensitive, substring-based termination check."""
                content = ""
                if isinstance(msg, dict):
                    content = msg.get("content", "")
                elif isinstance(msg, str):
                    content = msg

                if content and "terminate" in content.lower():
                    return True

                # Fall back to original check
                try:
                    return original_term(msg)
                except Exception:
                    return False

            chat.is_termination_msg = robust_termination

    manager._lambdagent_monitor = monitor
    return manager