codekingpro/portable-devtools

1"""Human in the loop middleware."""2 3from typing import Any, Literal, Protocol4 5from langchain_core.messages import AIMessage, ToolCall, ToolMessage6from langgraph.runtime import Runtime7from langgraph.types import interrupt8from typing_extensions import NotRequired, TypedDict9 10from langchain.agents.middleware.types import (11    AgentMiddleware,12    AgentState,13    ContextT,14    ResponseT,15    StateT,16)17 18 19class Action(TypedDict):20    """Represents an action with a name and args."""21 22    name: str23    """The type or name of action being requested (e.g., `'add_numbers'`)."""24 25    args: dict[str, Any]26    """Key-value pairs of args needed for the action (e.g., `{"a": 1, "b": 2}`)."""27 28 29class ActionRequest(TypedDict):30    """Represents an action request with a name, args, and description."""31 32    name: str33    """The name of the action being requested."""34 35    args: dict[str, Any]36    """Key-value pairs of args needed for the action (e.g., `{"a": 1, "b": 2}`)."""37 38    description: NotRequired[str]39    """The description of the action to be reviewed."""40 41 42DecisionType = Literal["approve", "edit", "reject", "respond"]43 44 45class ReviewConfig(TypedDict):46    """Policy for reviewing a HITL request."""47 48    action_name: str49    """Name of the action associated with this review configuration."""50 51    allowed_decisions: list[DecisionType]52    """The decisions that are allowed for this request."""53 54    args_schema: NotRequired[dict[str, Any]]55    """JSON schema for the args associated with the action, if edits are allowed."""56 57 58class HITLRequest(TypedDict):59    """Request for human feedback on a sequence of actions requested by a model."""60 61    action_requests: list[ActionRequest]62    """A list of agent actions for human review."""63 64    review_configs: list[ReviewConfig]65    """Review configuration for all possible actions."""66 67 68class ApproveDecision(TypedDict):69    """Response when a human approves the action."""70 71    type: Literal["approve"]72    """The type of response when a human approves the action."""73 74 75class EditDecision(TypedDict):76    """Response when a human edits the action."""77 78    type: Literal["edit"]79    """The type of response when a human edits the action."""80 81    edited_action: Action82    """Edited action for the agent to perform.83 84    Ex: for a tool call, a human reviewer can edit the tool name and args.85    """86 87 88class RejectDecision(TypedDict):89    """Response when a human rejects the action."""90 91    type: Literal["reject"]92    """The type of response when a human rejects the action."""93 94    message: NotRequired[str]95    """The message sent to the model explaining why the action was rejected."""96 97 98class RespondDecision(TypedDict):99    """Response when a human answers on behalf of the tool, skipping execution.100 101    Used for "ask user" style tools whose real implementation is the human's102    response. The tool is not executed; instead, a synthetic `ToolMessage` with103    `status="success"` and the provided `message` is returned to the model.104    """105 106    type: Literal["respond"]107    """The type of response when a human responds on behalf of the tool."""108 109    message: str110    """Content of the synthetic `ToolMessage` returned to the model."""111 112 113Decision = ApproveDecision | EditDecision | RejectDecision | RespondDecision114 115 116class HITLResponse(TypedDict):117    """Response payload for a HITLRequest."""118 119    decisions: list[Decision]120    """The decisions made by the human."""121 122 123class _DescriptionFactory(Protocol):124    """Callable that generates a description for a tool call."""125 126    def __call__(127        self, tool_call: ToolCall, state: AgentState[Any], runtime: Runtime[ContextT]128    ) -> str:129        """Generate a description for a tool call."""130        ...131 132 133class InterruptOnConfig(TypedDict):134    """Configuration for an action requiring human in the loop.135 136    This is the configuration format used in the `HumanInTheLoopMiddleware.__init__`137    method.138    """139 140    allowed_decisions: list[DecisionType]141    """The decisions that are allowed for this action."""142 143    description: NotRequired[str | _DescriptionFactory]144    """The description attached to the request for human input.145 146    Can be either:147 148    - A static string describing the approval request149    - A callable that dynamically generates the description based on agent state,150        runtime, and tool call information151 152    Example:153        ```python154        # Static string description155        config = ToolConfig(156            allowed_decisions=["approve", "reject"],157            description="Please review this tool execution"158        )159 160        # Dynamic callable description161        def format_tool_description(162            tool_call: ToolCall,163            state: AgentState,164            runtime: Runtime[ContextT]165        ) -> str:166            import json167            return (168                f"Tool: {tool_call['name']}\\n"169                f"Arguments:\\n{json.dumps(tool_call['args'], indent=2)}"170            )171 172        config = InterruptOnConfig(173            allowed_decisions=["approve", "edit", "reject"],174            description=format_tool_description175        )176        ```177    """178    args_schema: NotRequired[dict[str, Any]]179    """JSON schema for the args associated with the action, if edits are allowed."""180 181 182class HumanInTheLoopMiddleware(AgentMiddleware[StateT, ContextT, ResponseT]):183    """Human in the loop middleware."""184 185    def __init__(186        self,187        interrupt_on: dict[str, bool | InterruptOnConfig],188        *,189        description_prefix: str = "Tool execution requires approval",190    ) -> None:191        """Initialize the human in the loop middleware.192 193        Args:194            interrupt_on: Mapping of tool name to allowed actions.195 196                If a tool doesn't have an entry, it's auto-approved by default.197 198                * `True` indicates all decisions are allowed: approve, edit, reject,199                    and respond.200                * `False` indicates that the tool is auto-approved.201                * `InterruptOnConfig` indicates the specific decisions allowed for this202                    tool.203 204                    The `InterruptOnConfig` can include a `description` field (`str` or205                    `Callable`) for custom formatting of the interrupt description.206            description_prefix: The prefix to use when constructing action requests.207 208                This is used to provide context about the tool call and the action being209                requested.210 211                Not used if a tool has a `description` in its `InterruptOnConfig`.212        """213        super().__init__()214        resolved_configs: dict[str, InterruptOnConfig] = {}215        for tool_name, tool_config in interrupt_on.items():216            if isinstance(tool_config, bool):217                if tool_config is True:218                    resolved_configs[tool_name] = InterruptOnConfig(219                        allowed_decisions=["approve", "edit", "reject", "respond"]220                    )221            elif tool_config.get("allowed_decisions"):222                resolved_configs[tool_name] = tool_config223        self.interrupt_on = resolved_configs224        self.description_prefix = description_prefix225 226    def _create_action_and_config(227        self,228        tool_call: ToolCall,229        config: InterruptOnConfig,230        state: AgentState[Any],231        runtime: Runtime[ContextT],232    ) -> tuple[ActionRequest, ReviewConfig]:233        """Create an ActionRequest and ReviewConfig for a tool call."""234        tool_name = tool_call["name"]235        tool_args = tool_call["args"]236 237        # Generate description using the description field (str or callable)238        description_value = config.get("description")239        if callable(description_value):240            description = description_value(tool_call, state, runtime)241        elif description_value is not None:242            description = description_value243        else:244            description = f"{self.description_prefix}\n\nTool: {tool_name}\nArgs: {tool_args}"245 246        # Create ActionRequest with description247        action_request = ActionRequest(248            name=tool_name,249            args=tool_args,250            description=description,251        )252 253        # Create ReviewConfig254        # eventually can get tool information and populate args_schema from there255        review_config = ReviewConfig(256            action_name=tool_name,257            allowed_decisions=config["allowed_decisions"],258        )259 260        return action_request, review_config261 262    @staticmethod263    def _process_decision(264        decision: Decision,265        tool_call: ToolCall,266        config: InterruptOnConfig,267    ) -> tuple[ToolCall | None, ToolMessage | None]:268        """Process a single decision and return the revised tool call and optional tool message."""269        allowed_decisions = config["allowed_decisions"]270 271        if decision["type"] == "approve" and "approve" in allowed_decisions:272            return tool_call, None273        if decision["type"] == "edit" and "edit" in allowed_decisions:274            edited_action = decision["edited_action"]275            return (276                ToolCall(277                    type="tool_call",278                    name=edited_action["name"],279                    args=edited_action["args"],280                    id=tool_call["id"],281                ),282                None,283            )284        if decision["type"] == "reject" and "reject" in allowed_decisions:285            # Create a tool message with the human's text response286            content = decision.get("message") or (287                f"User rejected the tool call for `{tool_call['name']}` with id {tool_call['id']}"288            )289            tool_message = ToolMessage(290                content=content,291                name=tool_call["name"],292                tool_call_id=tool_call["id"],293                status="error",294            )295            return tool_call, tool_message296        if decision["type"] == "respond" and "respond" in allowed_decisions:297            # Skip tool execution; the human answers on behalf of the tool.298            tool_message = ToolMessage(299                content=decision["message"],300                name=tool_call["name"],301                tool_call_id=tool_call["id"],302                status="success",303            )304            return tool_call, tool_message305        msg = (306            f"Unexpected human decision: {decision}. "307            f"Decision type '{decision.get('type')}' "308            f"is not allowed for tool '{tool_call['name']}'. "309            f"Expected one of {allowed_decisions} based on the tool's configuration."310        )311        raise ValueError(msg)312 313    def after_model(314        self, state: AgentState[Any], runtime: Runtime[ContextT]315    ) -> dict[str, Any] | None:316        """Trigger interrupt flows for relevant tool calls after an `AIMessage`.317 318        Args:319            state: The current agent state.320            runtime: The runtime context.321 322        Returns:323            Updated message with the revised tool calls.324 325        Raises:326            ValueError: If the number of human decisions does not match the number of327                interrupted tool calls.328        """329        messages = state["messages"]330        if not messages:331            return None332 333        last_ai_msg = next((msg for msg in reversed(messages) if isinstance(msg, AIMessage)), None)334        if not last_ai_msg or not last_ai_msg.tool_calls:335            return None336 337        # Create action requests and review configs for tools that need approval338        action_requests: list[ActionRequest] = []339        review_configs: list[ReviewConfig] = []340        interrupt_indices: list[int] = []341 342        for idx, tool_call in enumerate(last_ai_msg.tool_calls):343            if (config := self.interrupt_on.get(tool_call["name"])) is not None:344                action_request, review_config = self._create_action_and_config(345                    tool_call, config, state, runtime346                )347                action_requests.append(action_request)348                review_configs.append(review_config)349                interrupt_indices.append(idx)350 351        # If no interrupts needed, return early352        if not action_requests:353            return None354 355        # Create single HITLRequest with all actions and configs356        hitl_request = HITLRequest(357            action_requests=action_requests,358            review_configs=review_configs,359        )360 361        # Send interrupt and get response362        decisions = interrupt(hitl_request)["decisions"]363 364        # Validate that the number of decisions matches the number of interrupt tool calls365        if (decisions_len := len(decisions)) != (interrupt_count := len(interrupt_indices)):366            msg = (367                f"Number of human decisions ({decisions_len}) does not match "368                f"number of hanging tool calls ({interrupt_count})."369            )370            raise ValueError(msg)371 372        # Process decisions and rebuild tool calls in original order373        revised_tool_calls: list[ToolCall] = []374        artificial_tool_messages: list[ToolMessage] = []375        decision_idx = 0376 377        for idx, tool_call in enumerate(last_ai_msg.tool_calls):378            if idx in interrupt_indices:379                # This was an interrupt tool call - process the decision380                config = self.interrupt_on[tool_call["name"]]381                decision = decisions[decision_idx]382                decision_idx += 1383 384                revised_tool_call, tool_message = self._process_decision(385                    decision, tool_call, config386                )387                if revised_tool_call is not None:388                    revised_tool_calls.append(revised_tool_call)389                if tool_message:390                    artificial_tool_messages.append(tool_message)391            else:392                # This was auto-approved - keep original393                revised_tool_calls.append(tool_call)394 395        # Update the AI message to only include approved tool calls396        last_ai_msg.tool_calls = revised_tool_calls397 398        return {"messages": [last_ai_msg, *artificial_tool_messages]}399 400    async def aafter_model(401        self, state: AgentState[Any], runtime: Runtime[ContextT]402    ) -> dict[str, Any] | None:403        """Async trigger interrupt flows for relevant tool calls after an `AIMessage`.404 405        Args:406            state: The current agent state.407            runtime: The runtime context.408 409        Returns:410            Updated message with the revised tool calls.411        """412        return self.after_model(state, runtime)413