agent-killswitch 0.1.0__py3-none-any.whl

agent_killswitch/integrations/crewai.py

@@ -0,0 +1,151 @@
+"""CrewAI integration for agent-killswitch.
+
+Provides a pre-action hook that checks the kill switch before
+each CrewAI agent action.
+
+Requires: pip install agent-killswitch[crewai]
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+class KillSwitchHook:
+    """CrewAI pre-action hook that checks the kill switch.
+
+    This hook is designed to be called before each agent action in
+    a CrewAI workflow. If the kill switch is active, it prevents
+    the action from executing.
+
+    Example::
+
+        from agent_killswitch import KillSwitch
+        from agent_killswitch.integrations.crewai import KillSwitchHook
+
+        ks = KillSwitch()
+        hook = KillSwitchHook(killswitch=ks, crew_id="my-crew")
+
+        # In your CrewAI setup:
+        # Use hook.before_action as a step callback
+        if hook.before_action(agent_id="researcher"):
+            # Action is allowed
+            pass
+
+    Args:
+        killswitch: The KillSwitch instance.
+        crew_id: The crew ID for scope checks.
+        on_blocked: Callback when an action is blocked by kill switch.
+            Signature: (agent_id: str, level: str, reason: str) -> None
+    """
+
+    def __init__(
+        self,
+        killswitch: Any,
+        crew_id: str | None = None,
+        on_blocked: Callable[[str, str, str], None] | None = None,
+    ) -> None:
+        self._ks = killswitch
+        self._crew_id = crew_id
+        self._on_blocked = on_blocked
+
+    def before_action(self, agent_id: str | None = None) -> bool:
+        """Check if the agent action should proceed.
+
+        Args:
+            agent_id: The agent attempting the action.
+
+        Returns:
+            True if the action should proceed, False if blocked.
+        """
+        status = self._ks.check(
+            scope_target=agent_id,
+            crew_id=self._crew_id,
+        )
+
+        if status.is_killed:
+            level_str = status.level.value if status.level else "unknown"
+            reason = status.reason or ""
+
+            logger.warning(
+                "crewai_action_blocked",
+                extra={
+                    "agent_id": agent_id,
+                    "crew_id": self._crew_id,
+                    "level": level_str,
+                    "reason": reason,
+                },
+            )
+
+            if self._on_blocked is not None:
+                try:
+                    self._on_blocked(agent_id or "", level_str, reason)
+                except Exception as exc:
+                    logger.error(
+                        "on_blocked callback error",
+                        extra={"error": str(exc)},
+                    )
+
+            return False
+
+        return True
+
+    def create_step_callback(
+        self,
+        agent_id: str | None = None,
+    ) -> Callable[[Any], Any]:
+        """Create a CrewAI step callback function.
+
+        Returns a callable suitable for use as a CrewAI step_callback.
+        The callback checks the kill switch and raises RuntimeError
+        if the switch is active.
+
+        Args:
+            agent_id: The agent ID.
+
+        Returns:
+            Callback function.
+        """
+
+        def callback(step_output: Any) -> Any:
+            if not self.before_action(agent_id=agent_id):
+                raise RuntimeError(f"Kill switch active for agent={agent_id}, crew={self._crew_id}")
+            return step_output
+
+        return callback
+
+    def wrap_crew(
+        self,
+        crew: Any,
+        agent_id_field: str = "role",
+    ) -> Any:
+        """Wrap a CrewAI crew with kill switch protection.
+
+        This method patches the crew to check the kill switch before
+        each agent action. Returns the modified crew.
+
+        Args:
+            crew: The CrewAI Crew instance.
+            agent_id_field: Agent attribute to use as agent_id.
+
+        Returns:
+            The crew with kill switch protection applied.
+        """
+        if hasattr(crew, "step_callback"):
+            original = crew.step_callback
+
+            def wrapped_callback(step_output: Any) -> Any:
+                # Check kill switch
+                if not self.before_action(agent_id=self._crew_id):
+                    raise RuntimeError(f"Kill switch active for crew={self._crew_id}")
+                if original is not None:
+                    return original(step_output)
+                return step_output
+
+            crew.step_callback = wrapped_callback
+
+        return crew

agent_killswitch/integrations/langchain.py

@@ -0,0 +1,135 @@
+"""LangChain integration for agent-killswitch.
+
+Provides a callback handler that checks the kill switch before and
+after each LLM call, tool invocation, and chain execution.
+
+Requires: pip install agent-killswitch[langchain]
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+from uuid import UUID
+
+logger = logging.getLogger(__name__)
+
+try:
+    from langchain_core.callbacks import BaseCallbackHandler
+except ImportError:
+    # Create a stub so the class can be defined without the dependency
+    class BaseCallbackHandler:  # type: ignore[no-redef]
+        """Stub for when langchain-core is not installed."""
+
+        pass
+
+
+class KillSwitchCallbackHandler(BaseCallbackHandler):
+    """LangChain callback handler that enforces kill switch checks.
+
+    Checks the kill switch before each LLM call and tool invocation.
+    If the kill switch is active, raises an exception to halt the chain.
+
+    Example::
+
+        from langchain_openai import ChatOpenAI
+        from agent_killswitch import KillSwitch
+        from agent_killswitch.integrations.langchain import KillSwitchCallbackHandler
+
+        ks = KillSwitch()
+        handler = KillSwitchCallbackHandler(killswitch=ks, agent_id="my-agent")
+
+        llm = ChatOpenAI(callbacks=[handler])
+        llm.invoke("Hello")  # Will check kill switch before calling
+
+    Args:
+        killswitch: The KillSwitch instance to check.
+        agent_id: The agent ID for scope-specific checks.
+        crew_id: Optional crew ID for crew scope checks.
+        raise_on_kill: If True, raise RuntimeError on kill. Default True.
+    """
+
+    def __init__(
+        self,
+        killswitch: Any,
+        agent_id: str | None = None,
+        crew_id: str | None = None,
+        raise_on_kill: bool = True,
+    ) -> None:
+        self._ks = killswitch
+        self._agent_id = agent_id
+        self._crew_id = crew_id
+        self._raise_on_kill = raise_on_kill
+
+    def _check(self, action: str) -> None:
+        """Check the kill switch and raise if active.
+
+        Args:
+            action: Description of the action being checked.
+        """
+        status = self._ks.check(
+            scope_target=self._agent_id,
+            crew_id=self._crew_id,
+        )
+        if status.is_killed:
+            msg = (
+                f"Kill switch active during {action}: level={status.level}, reason={status.reason}"
+            )
+            logger.warning("langchain_kill_switch_triggered", extra={"action": action})
+            if self._raise_on_kill:
+                raise RuntimeError(msg)
+
+    def on_llm_start(
+        self,
+        serialized: dict[str, Any],
+        prompts: list[str],
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Check kill switch before LLM call."""
+        self._check("llm_start")
+
+    def on_llm_end(
+        self,
+        response: Any,
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        tags: list[str] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Check kill switch after LLM call."""
+        self._check("llm_end")
+
+    def on_tool_start(
+        self,
+        serialized: dict[str, Any],
+        input_str: str,
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        inputs: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Check kill switch before tool invocation."""
+        self._check("tool_start")
+
+    def on_chain_start(
+        self,
+        serialized: dict[str, Any],
+        inputs: dict[str, Any],
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Check kill switch before chain execution."""
+        self._check("chain_start")

agent_killswitch/integrations/langgraph.py

@@ -0,0 +1,143 @@
+"""LangGraph integration for agent-killswitch.
+
+Provides a killswitch check node and conditional edge for use
+in LangGraph state graphs.
+
+Requires: pip install agent-killswitch[langgraph]
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+class KillSwitchNode:
+    """A LangGraph node that checks the kill switch.
+
+    Use this as a node in a LangGraph StateGraph. If the kill switch
+    is active, it sets a flag in the state that can be used with
+    conditional edges to route to a termination node.
+
+    Example::
+
+        from langgraph.graph import StateGraph
+        from agent_killswitch import KillSwitch
+        from agent_killswitch.integrations.langgraph import KillSwitchNode
+
+        ks = KillSwitch()
+        ks_node = KillSwitchNode(killswitch=ks, agent_id="graph-agent")
+
+        graph = StateGraph(dict)
+        graph.add_node("check_kill", ks_node)
+        graph.add_node("agent_work", agent_node)
+        graph.add_node("terminated", termination_node)
+        graph.add_conditional_edges(
+            "check_kill",
+            ks_node.should_continue,
+            {"continue": "agent_work", "killed": "terminated"},
+        )
+
+    Args:
+        killswitch: The KillSwitch instance.
+        agent_id: Agent ID for scope checks.
+        crew_id: Optional crew ID for crew scope checks.
+        kill_state_key: Key in state dict to set when killed.
+            Default "kill_switch_active".
+        kill_level_key: Key in state dict to set with kill level.
+            Default "kill_level".
+        kill_reason_key: Key in state dict to set with kill reason.
+            Default "kill_reason".
+    """
+
+    def __init__(
+        self,
+        killswitch: Any,
+        agent_id: str | None = None,
+        crew_id: str | None = None,
+        kill_state_key: str = "kill_switch_active",
+        kill_level_key: str = "kill_level",
+        kill_reason_key: str = "kill_reason",
+    ) -> None:
+        self._ks = killswitch
+        self._agent_id = agent_id
+        self._crew_id = crew_id
+        self._kill_state_key = kill_state_key
+        self._kill_level_key = kill_level_key
+        self._kill_reason_key = kill_reason_key
+
+    def __call__(self, state: dict[str, Any]) -> dict[str, Any]:
+        """Check the kill switch and update state.
+
+        Args:
+            state: The current LangGraph state.
+
+        Returns:
+            Updated state with kill switch status.
+        """
+        status = self._ks.check(
+            scope_target=self._agent_id,
+            crew_id=self._crew_id,
+        )
+
+        result = dict(state)
+        result[self._kill_state_key] = status.is_killed
+        if status.is_killed:
+            result[self._kill_level_key] = status.level.value if status.level else None
+            result[self._kill_reason_key] = status.reason or ""
+            logger.warning(
+                "langgraph_kill_switch_active",
+                extra={
+                    "agent_id": self._agent_id,
+                    "level": result[self._kill_level_key],
+                },
+            )
+        else:
+            result[self._kill_level_key] = None
+            result[self._kill_reason_key] = ""
+
+        return result
+
+    def should_continue(self, state: dict[str, Any]) -> str:
+        """Conditional edge function: 'continue' or 'killed'.
+
+        Use with graph.add_conditional_edges() to route based on
+        kill switch status.
+
+        Args:
+            state: The current LangGraph state.
+
+        Returns:
+            'killed' if kill switch is active, 'continue' otherwise.
+        """
+        if state.get(self._kill_state_key, False):
+            return "killed"
+        return "continue"
+
+
+def create_killswitch_node(
+    killswitch: Any,
+    agent_id: str | None = None,
+    crew_id: str | None = None,
+) -> tuple[Any, Any]:
+    """Convenience function to create a kill switch node and condition.
+
+    Returns a tuple of (node_function, condition_function) ready to
+    add to a LangGraph StateGraph.
+
+    Args:
+        killswitch: The KillSwitch instance.
+        agent_id: Agent ID for scope checks.
+        crew_id: Optional crew ID.
+
+    Returns:
+        Tuple of (node_callable, condition_callable).
+    """
+    node = KillSwitchNode(
+        killswitch=killswitch,
+        agent_id=agent_id,
+        crew_id=crew_id,
+    )
+    return node, node.should_continue

agent_killswitch/integrations/openai_agents.py

@@ -0,0 +1,151 @@
+"""OpenAI Agents SDK integration for agent-killswitch.
+
+Provides a guardrail function and input guard compatible with the
+OpenAI Agents SDK.
+
+Requires: pip install agent-killswitch[openai]
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+class KillSwitchGuardrail:
+    """OpenAI Agents SDK guardrail that enforces kill switch checks.
+
+    Can be used as an input guardrail or output guardrail in the
+    OpenAI Agents SDK.
+
+    Example::
+
+        from agent_killswitch import KillSwitch
+        from agent_killswitch.integrations.openai_agents import KillSwitchGuardrail
+
+        ks = KillSwitch()
+        guardrail = KillSwitchGuardrail(killswitch=ks, agent_id="openai-agent")
+
+        # Use as input guardrail
+        result = guardrail.check_input("user message")
+        if not result["allowed"]:
+            print(f"Blocked: {result['reason']}")
+
+    Args:
+        killswitch: The KillSwitch instance.
+        agent_id: Agent ID for scope checks.
+        crew_id: Optional crew ID for crew scope checks.
+    """
+
+    def __init__(
+        self,
+        killswitch: Any,
+        agent_id: str | None = None,
+        crew_id: str | None = None,
+    ) -> None:
+        self._ks = killswitch
+        self._agent_id = agent_id
+        self._crew_id = crew_id
+
+    def check_input(self, user_input: Any) -> dict[str, Any]:
+        """Check kill switch before processing input.
+
+        Args:
+            user_input: The user input being processed.
+
+        Returns:
+            Dictionary with keys:
+            - 'allowed': bool
+            - 'reason': str (if blocked)
+            - 'level': str (if blocked)
+        """
+        status = self._ks.check(
+            scope_target=self._agent_id,
+            crew_id=self._crew_id,
+        )
+
+        if status.is_killed:
+            level_str = status.level.value if status.level else "unknown"
+            return {
+                "allowed": False,
+                "reason": status.reason or "Kill switch is active",
+                "level": level_str,
+            }
+
+        return {"allowed": True, "reason": "", "level": ""}
+
+    def check_output(self, output: Any) -> dict[str, Any]:
+        """Check kill switch before returning output.
+
+        Args:
+            output: The agent output being validated.
+
+        Returns:
+            Dictionary with keys:
+            - 'allowed': bool
+            - 'reason': str (if blocked)
+            - 'level': str (if blocked)
+        """
+        return self.check_input(output)
+
+    def as_guardrail_function(self) -> Any:
+        """Return a guardrail function compatible with the OpenAI SDK.
+
+        Returns:
+            An async function suitable for use as a guardrail.
+        """
+
+        async def guardrail_fn(
+            context: Any,
+            agent: Any,
+            input_data: Any,
+        ) -> Any | None:
+            """Guardrail function that checks the kill switch.
+
+            Args:
+                context: The run context.
+                agent: The agent being guarded.
+                input_data: The input being processed.
+
+            Returns:
+                None if allowed, or a GuardrailFunctionOutput if blocked.
+            """
+            status = self._ks.check(
+                scope_target=self._agent_id,
+                crew_id=self._crew_id,
+            )
+
+            if status.is_killed:
+                level_str = status.level.value if status.level else "unknown"
+                reason = status.reason or "Kill switch is active"
+                logger.warning(
+                    "openai_agents_guardrail_blocked",
+                    extra={
+                        "agent_id": self._agent_id,
+                        "level": level_str,
+                        "reason": reason,
+                    },
+                )
+
+                # Try to use the SDK's GuardrailFunctionOutput if available
+                try:
+                    from openai.agents import GuardrailFunctionOutput, GuardrailTripwireTriggered
+
+                    output = GuardrailFunctionOutput(
+                        output_info={"blocked": True, "reason": reason, "level": level_str},
+                        tripwire_triggered=True,
+                    )
+                    raise GuardrailTripwireTriggered(output)
+                except ImportError:
+                    # The openai.agents extra is not installed. That is an expected,
+                    # supported configuration - chaining the ImportError here would
+                    # imply something went wrong, so suppress the context.
+                    raise RuntimeError(
+                        f"Kill switch active: level={level_str}, reason={reason}"
+                    ) from None
+
+            return None
+
+        return guardrail_fn