AutoGen

import os
import autogen
from cheqpoint import CheqpointClient, RejectedError

client = CheqpointClient(api_key=os.environ["CHEQPOINT_CONNECTION_KEY"])

def with_approval(action: str, risk_score: float = 0.8):
    """Decorator that gates an AutoGen tool behind Cheqpoint human approval."""
    def decorator(fn):
        def wrapper(*args, **kwargs):
            # Build a human-readable summary
            summary = f"AutoGen Assistant wants to execute: {action}"
            if args:
                summary += f" with args: {args[:2]}"  # show first 2 args

            try:
                result = client.checkpoint(
                    action=action,
                    summary=summary,
                    details={"args": list(args), "kwargs": kwargs},
                    risk_score=risk_score,
                )
            except RejectedError as e:
                raise RuntimeError(f"Action rejected: {e}") from e

            # Use modified_details if reviewer changed the payload
            effective = result.modified_details or {}
            if effective:
                kwargs.update(effective)

            return fn(*args, **kwargs)
        wrapper.__name__ = fn.__name__
        return wrapper
    return decorator

@with_approval("process_refund", risk_score=0.9)
def process_refund(order_id: str, amount: float) -> dict:
    """Process a customer refund via Stripe."""
    return stripe_client.refunds.create(
        charge=lookup_charge(order_id),
        amount=int(amount * 100),
    )

@with_approval("send_bulk_email", risk_score=0.6)
def send_bulk_email(segment: str, subject: str, body: str) -> dict:
    """Send a marketing email to a customer segment."""
    return email_client.campaigns.send(
        segment=segment, subject=subject, body=body
    )

assistant = autogen.AssistantAgent(
    name="FinanceAssistant",
    llm_config={"model": "gpt-4o"},
    system_message=(
        "You handle finance tasks. "
        "Always call process_refund for refunds and send_bulk_email for campaigns."
    ),
)

user_proxy = autogen.UserProxyAgent(
    name="UserProxy",
    human_input_mode="NEVER",
    function_map={
        "process_refund": process_refund,
        "send_bulk_email": send_bulk_email,
    },
    code_execution_config=False,
)

# This will pause automatically when the Assistant calls a gated function
user_proxy.initiate_chat(
    assistant,
    message="Please process a $300 refund for order #54321",
)

# For AutoGen tools where you don't want to block the agent,
# use request_async() — returns immediately with status.

def process_refund_async(order_id: str, amount: float) -> dict:
    result = client.request_async(
        action="process_refund",
        summary=f"Refund ${amount} for order {order_id}",
        details={"orderId": order_id, "amount": amount},
    )
    if result["status"] == "approved":
        return {"success": True, "orderId": order_id}
    elif result["status"] == "rejected":
        raise RuntimeError(f"Refund rejected: {result.get('decisionNote', '')}")
    else:
        # "pending" — decision awaiting; use approvalId to follow up
        return {"pending": True, "approvalId": result["approvalId"]}