"""Middleware for intercepting clarification requests and presenting them to the user.""" import json import logging from collections.abc import Callable from hashlib import sha256 from typing import override from deerflow.tools.builtins.clarification_utils import resolve_allow_multiple from langchain.agents import AgentState from langchain.agents.middleware import AgentMiddleware from langchain_core.messages import ToolMessage from langgraph.graph import END from langgraph.prebuilt.tool_node import ToolCallRequest from langgraph.types import Command logger = logging.getLogger(__name__) class ClarificationMiddlewareState(AgentState): """Compatible with the `ThreadState` schema.""" pass class ClarificationMiddleware(AgentMiddleware[ClarificationMiddlewareState]): """Intercepts clarification tool calls and interrupts execution to present questions to the user. When the model calls the `ask_clarification` tool, this middleware: 1. Intercepts the tool call before execution 2. Extracts the clarification question and metadata 3. Formats a user-friendly message 4. Returns a Command that interrupts execution and presents the question 5. Waits for user response before continuing This replaces the tool-based approach where clarification continued the conversation flow. """ state_schema = ClarificationMiddlewareState def _stable_message_id(self, tool_call_id: str, formatted_message: str) -> str: """Build a deterministic message ID so retried clarification calls replace, not append.""" if tool_call_id: return f"clarification:{tool_call_id}" digest = sha256(formatted_message.encode("utf-8")).hexdigest()[:16] return f"clarification:{digest}" def _is_chinese(self, text: str) -> bool: """Check if text contains Chinese characters. Args: text: Text to check Returns: True if text contains Chinese characters """ return any("\u4e00" <= char <= "\u9fff" for char in text) def _format_clarification_message(self, args: dict) -> str: """Format the clarification arguments into a user-friendly message. Args: args: The tool call arguments containing clarification details Returns: Formatted message string """ question = args.get("question", "") clarification_type = args.get("clarification_type", "missing_info") context = args.get("context") options = self._normalize_options(args.get("options", [])) # Type-specific icons type_icons = { "missing_info": "❓", "ambiguous_requirement": "🤔", "approach_choice": "🔀", "risk_confirmation": "⚠️", "suggestion": "💡", } icon = type_icons.get(clarification_type, "❓") # Build the message naturally message_parts = [] # Add icon and question together for a more natural flow if context: # If there's context, present it first as background message_parts.append(f"{icon} {context}") message_parts.append(f"\n{question}") else: # Just the question with icon message_parts.append(f"{icon} {question}") # Add options in a cleaner format if options and len(options) > 0: message_parts.append("") # blank line for spacing for i, option in enumerate(options, 1): label = option.get("label") or option.get("id") or "" message_parts.append(f" {i}. {label}") return "\n".join(message_parts) def _normalize_options(self, raw_options: object) -> list[dict[str, str]]: """Normalize model-produced options into a stable UI-friendly shape.""" options = raw_options # Some models serialize array parameters as JSON strings instead of native arrays. if isinstance(options, str): try: options = json.loads(options) except (json.JSONDecodeError, TypeError): options = [options] if options is None: values: list[object] = [] elif isinstance(options, list): values = options else: values = [options] normalized: list[dict[str, str]] = [] for index, value in enumerate(values, 1): if isinstance(value, dict): raw_label = value.get("label") or value.get("title") or value.get("name") or value.get("value") or value.get("id") label = str(raw_label).strip() if raw_label is not None else "" if not label: continue raw_id = value.get("id") or value.get("value") or f"option-{index}" option = { "id": str(raw_id).strip() or f"option-{index}", "label": label, } description = value.get("description") if description is not None and str(description).strip(): option["description"] = str(description).strip() normalized.append(option) else: label = str(value).strip() if label: normalized.append({"id": f"option-{index}", "label": label}) return normalized def _build_payload(self, args: dict) -> dict: """Build structured clarification data for rich clients.""" return { "question": str(args.get("question", "")).strip(), "clarification_type": args.get("clarification_type", "missing_info"), "context": args.get("context"), "options": self._normalize_options(args.get("options", [])), "allow_custom": True, "allow_multiple": resolve_allow_multiple(args), } def _handle_clarification(self, request: ToolCallRequest) -> Command: """Handle clarification request and return command to interrupt execution. Args: request: Tool call request Returns: Command that interrupts execution with the formatted clarification message """ # Extract clarification arguments args = request.tool_call.get("args", {}) question = args.get("question", "") logger.info("Intercepted clarification request") logger.debug("Clarification question: %s", question) # Format the clarification message formatted_message = self._format_clarification_message(args) # Get the tool call ID tool_call_id = request.tool_call.get("id", "") # Create a ToolMessage with the formatted question # This will be added to the message history tool_message = ToolMessage( id=self._stable_message_id(tool_call_id, formatted_message), content=formatted_message, tool_call_id=tool_call_id, name="ask_clarification", additional_kwargs={"clarification": self._build_payload(args)}, ) # Return a Command that: # 1. Adds the formatted tool message # 2. Interrupts execution by going to __end__ # Note: We don't add an extra AIMessage here - the frontend will detect # and display ask_clarification tool messages directly return Command( update={"messages": [tool_message]}, goto=END, ) @override def wrap_tool_call( self, request: ToolCallRequest, handler: Callable[[ToolCallRequest], ToolMessage | Command], ) -> ToolMessage | Command: """Intercept ask_clarification tool calls and interrupt execution (sync version). Args: request: Tool call request handler: Original tool execution handler Returns: Command that interrupts execution with the formatted clarification message """ # Check if this is an ask_clarification tool call if request.tool_call.get("name") != "ask_clarification": # Not a clarification call, execute normally return handler(request) return self._handle_clarification(request) @override async def awrap_tool_call( self, request: ToolCallRequest, handler: Callable[[ToolCallRequest], ToolMessage | Command], ) -> ToolMessage | Command: """Intercept ask_clarification tool calls and interrupt execution (async version). Args: request: Tool call request handler: Original tool execution handler (async) Returns: Command that interrupts execution with the formatted clarification message """ # Check if this is an ask_clarification tool call if request.tool_call.get("name") != "ask_clarification": # Not a clarification call, execute normally return await handler(request) return self._handle_clarification(request)