ag2 0.9.9__py3-none-any.whl → 0.9.10__py3-none-any.whl

autogen/agentchat/group/safeguards/api.py

@@ -0,0 +1,224 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from ....llm_config import LLMConfig
+
+# Import types that are safe to import at runtime
+from ...agent import Agent
+from ...conversable_agent import ConversableAgent
+from ...groupchat import GroupChatManager
+
+if TYPE_CHECKING:
+    from .enforcer import SafeguardEnforcer
+
+
+def reset_safeguard_policy(
+    *,
+    agents: list[ConversableAgent] | None = None,
+    groupchat_manager: GroupChatManager | None = None,
+) -> None:
+    """Reset/remove all safeguards from agents and groupchat managers.
+
+    This function removes all safeguard hooks and inter-agent guardrails that were
+    previously applied by apply_safeguard_policy. Events are sent to track the reset process.
+
+    Args:
+        agents: List of agents to remove safeguards from (optional if groupchat_manager provided)
+        groupchat_manager: GroupChatManager to remove safeguards from (optional if agents provided)
+
+    Example:
+        ```python
+        from autogen.agentchat.group.safeguards import reset_safeguard_policy
+
+        # Remove safeguards from agents
+        reset_safeguard_policy(agents=[agent1, agent2, agent3])
+
+        # Or remove from GroupChatManager
+        reset_safeguard_policy(groupchat_manager=manager)
+        ```
+    """
+    from ....events.print_event import PrintEvent
+    from ....io.base import IOStream
+
+    iostream = IOStream.get_default()
+
+    # Send initial reset event
+    iostream.send(PrintEvent("Resetting safeguards..."))
+    # Determine which agents to remove safeguards from
+    target_agents: list[ConversableAgent | Agent] = []
+
+    if groupchat_manager:
+        if not isinstance(groupchat_manager, GroupChatManager):
+            raise ValueError("groupchat_manager must be an instance of GroupChatManager")
+
+        target_agents.extend([agent for agent in groupchat_manager.groupchat.agents if hasattr(agent, "hook_lists")])
+
+        agent_names = [agent.name for agent in target_agents]
+        iostream.send(PrintEvent(f"📋 Found {len(target_agents)} agents in GroupChat: {', '.join(agent_names)}"))
+
+        # Clear inter-agent guardrails from the groupchat
+        if hasattr(groupchat_manager.groupchat, "_inter_agent_guardrails"):
+            guardrail_count = len(groupchat_manager.groupchat._inter_agent_guardrails)
+            if guardrail_count > 0:
+                iostream.send(PrintEvent(f"🔗 Clearing {guardrail_count} inter-agent guardrails from GroupChat"))
+            groupchat_manager.groupchat._inter_agent_guardrails.clear()
+
+    elif agents:
+        target_agents.extend(agents)
+        agent_names = [agent.name for agent in target_agents]
+        iostream.send(PrintEvent(f"📋 Resetting safeguards for {len(target_agents)} agents: {', '.join(agent_names)}"))
+    else:
+        raise ValueError("Either agents or groupchat_manager must be provided")
+
+    # Remove safeguard hooks from each agent
+    safeguard_hook_names = [
+        "safeguard_tool_inputs",
+        "safeguard_tool_outputs",
+        "safeguard_llm_inputs",
+        "safeguard_llm_outputs",
+        "safeguard_human_inputs",
+        "process_message_before_send",  # Inter-agent communication hooks for direct agent communication
+    ]
+
+    for agent in target_agents:
+        if hasattr(agent, "hook_lists"):
+            # Use the agent's reset_safeguards method for agent-to-environment safeguards
+            if hasattr(agent, "reset_safeguards"):
+                agent.reset_safeguards()
+            else:
+                # Fallback to manual clearing for older agent versions
+                for hook_name in safeguard_hook_names:
+                    if hook_name in agent.hook_lists:
+                        # Clear all hooks in safeguard-specific hook lists
+                        agent.hook_lists[hook_name].clear()
+
+            # Manually clear inter-agent safeguards (process_message_before_send)
+            if "process_message_before_send" in agent.hook_lists:
+                agent.hook_lists["process_message_before_send"].clear()
+        else:
+            raise ValueError(
+                f"Agent {agent.name} does not support hooks. Please ensure it inherits from ConversableAgent."
+            )
+
+    # Send completion event
+    iostream.send(PrintEvent(f"✅ Safeguard reset completed for {len(target_agents)} agents"))
+
+
+def apply_safeguard_policy(
+    *,
+    agents: list[ConversableAgent] | None = None,
+    groupchat_manager: GroupChatManager | None = None,
+    policy: dict[str, Any] | str,
+    safeguard_llm_config: LLMConfig | None = None,
+    mask_llm_config: LLMConfig | None = None,
+) -> SafeguardEnforcer:
+    """Apply safeguards to agents using a policy file.
+
+    This is the main function for applying safeguards. It supports the policy format
+    with 'inter_agent_safeguards' and 'agent_environment_safeguards' sections.
+
+    Args:
+        agents: List of agents to apply safeguards to (optional if groupchat_manager provided)
+        groupchat_manager: GroupChatManager to apply safeguards to (optional if agents provided)
+        policy: Safeguard policy dict or path to JSON file
+        safeguard_llm_config: LLM configuration for safeguard checks
+        mask_llm_config: LLM configuration for masking
+
+    Returns:
+        SafeguardEnforcer instance for further configuration
+
+    Example:
+        ```python
+        from autogen.agentchat.group.safeguards import apply_safeguard_policy
+
+        # Apply safeguards to agents
+        safeguard_enforcer = apply_safeguard_policy(
+            agents=[agent1, agent2, agent3],
+            policy="path/to/policy.json",
+            safeguard_llm_config=safeguard_llm_config,
+        )
+
+        # Or apply to GroupChatManager
+        safeguard_enforcer = apply_safeguard_policy(
+            groupchat_manager=manager,
+            policy="path/to/policy.json",
+            safeguard_llm_config=safeguard_llm_config,
+            mask_llm_config=mask_llm_config,
+        )
+        ```
+    """
+    from .enforcer import SafeguardEnforcer
+
+    enforcer = SafeguardEnforcer(
+        policy=policy,
+        safeguard_llm_config=safeguard_llm_config,
+        mask_llm_config=mask_llm_config,
+    )
+
+    # Determine which agents to apply safeguards to
+    target_agents: list[ConversableAgent | Agent] = []
+    all_agent_names = []
+
+    if groupchat_manager:
+        if not isinstance(groupchat_manager, GroupChatManager):
+            raise ValueError("groupchat_manager must be an instance of GroupChatManager")
+
+        target_agents.extend([agent for agent in groupchat_manager.groupchat.agents if hasattr(agent, "hook_lists")])
+        all_agent_names = [agent.name for agent in groupchat_manager.groupchat.agents]
+        all_agent_names.append(groupchat_manager.name)
+
+        # Register inter-agent guardrails with the groupchat
+        # Ensure the list exists and append our enforcer
+        if not hasattr(groupchat_manager.groupchat, "_inter_agent_guardrails"):
+            groupchat_manager.groupchat._inter_agent_guardrails = []
+        groupchat_manager.groupchat._inter_agent_guardrails.clear()  # Clear any existing
+        groupchat_manager.groupchat._inter_agent_guardrails.append(enforcer)
+    elif agents:
+        target_agents.extend(agents)
+        all_agent_names = [agent.name for agent in agents]
+    else:
+        raise ValueError("Either agents or groupchat_manager must be provided")
+
+    # Build agent-to-tool mapping for validation
+    agent_tool_mapping = {}
+    for agent in target_agents:
+        agent_tools = []
+
+        # Get tools from the tools property (Tool objects)
+        if hasattr(agent, "tools"):
+            for tool in agent.tools:
+                agent_tools.append(tool.name)
+
+        # Get tools from function_map (functions registered with @register_for_execution)
+        if hasattr(agent, "function_map"):
+            agent_tools.extend(agent.function_map.keys())
+
+        agent_tool_mapping[agent.name] = agent_tools
+
+    # Validate policy including agent names and tool names
+    try:
+        from .validator import SafeguardValidator
+
+        validator = SafeguardValidator(enforcer.policy)  # Use enforcer's loaded policy dict
+        validator.validate_policy_complete(agent_names=all_agent_names, agent_tool_mapping=agent_tool_mapping)
+    except ValueError as e:
+        raise ValueError(f"Policy validation failed: {e}")
+
+    # Apply hooks to each agent
+    for agent in target_agents:
+        if hasattr(agent, "hook_lists"):
+            hooks = enforcer.create_agent_hooks(agent.name)
+            for hook_name, hook_func in hooks.items():
+                if hook_name in agent.hook_lists:
+                    agent.hook_lists[hook_name].append(hook_func)
+        else:
+            raise ValueError(
+                f"Agent {agent.name} does not support hooks. Please ensure it inherits from ConversableAgent."
+            )
+
+    return enforcer