| 1 | """Ask-user tool: send structured question cards to the WebUI.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | import uuid |
| 6 | from contextvars import ContextVar |
| 7 | from typing import Any, Awaitable, Callable |
| 8 | |
| 9 | from nanobot.agent.tools.base import Tool, tool_parameters |
| 10 | from nanobot.agent.tools.schema import ( |
| 11 | ArraySchema, |
| 12 | BooleanSchema, |
| 13 | ObjectSchema, |
| 14 | StringSchema, |
| 15 | tool_parameters_schema, |
| 16 | ) |
| 17 | from nanobot.bus.events import OutboundMessage |
| 18 | |
| 19 | REFERENCE_IMAGE_EDIT_OPTION = "我想修改/增删参考图" |
| 20 | STORY_CONFIRM_OPTION = "可以，按这个来" |
| 21 | STORY_REVISE_OPTION = "需要修改" |
| 22 | _REFERENCE_IMAGE_EDIT_ALIASES = frozenset( |
| 23 | { |
| 24 | REFERENCE_IMAGE_EDIT_OPTION, |
| 25 | "我要修改/增删参考图", |
| 26 | } |
| 27 | ) |
| 28 | _STORY_DIRECTION_CARD_IDS = frozenset({"confirm_story", "story_direction"}) |
| 29 | _STORY_DIRECTION_OPTION_MARKERS = frozenset({STORY_CONFIRM_OPTION, STORY_REVISE_OPTION}) |
| 30 | _STORY_DIRECTION_QUESTION_MARKERS = ( |
| 31 | "故事方向", |
| 32 | "什么样的故事", |
| 33 | "什么故事", |
| 34 | "想拍什么", |
| 35 | "故事题材", |
| 36 | ) |
| 37 | _NOT_STORY_DIRECTION_QUESTION_MARKERS = ( |
| 38 | "上传首帧", |
| 39 | "上传参考图", |
| 40 | "还需上传", |
| 41 | "确认使用此参考图", |
| 42 | "完成参考图修改", |
| 43 | "几个镜头", |
| 44 | "多少镜头", |
| 45 | ) |
| 46 | _UPLOAD_GATE_OPTIONS = frozenset({"需要上传,已上传完毕", "不上传"}) |
| 47 | |
| 48 | _QUESTION_SCHEMA = ObjectSchema( |
| 49 | { |
| 50 | "id": StringSchema("Stable id for this card (auto-generated if omitted)"), |
| 51 | "question": StringSchema("Question prompt shown above the option chips"), |
| 52 | "options": ArraySchema( |
| 53 | StringSchema("Option label"), |
| 54 | description="Tap choices for the user", |
| 55 | min_items=1, |
| 56 | ), |
| 57 | "allow_custom": BooleanSchema( |
| 58 | description="Whether the user may type a custom answer", |
| 59 | default=False, |
| 60 | ), |
| 61 | }, |
| 62 | required=["question", "options"], |
| 63 | ) |
| 64 | |
| 65 | PersistCallback = Callable[ |
| 66 | [str, str, str, str, list[dict[str, Any]], str, str], |
| 67 | Awaitable[None], |
| 68 | ] |
| 69 | |
| 70 | |
| 71 | def is_story_confirm_option(label: str) -> bool: |
| 72 | return (label or "").strip() == STORY_CONFIRM_OPTION |
| 73 | |
| 74 | |
| 75 | def is_reference_image_edit_option(label: str) -> bool: |
| 76 | return (label or "").strip() in _REFERENCE_IMAGE_EDIT_ALIASES |
| 77 | |
| 78 | |
| 79 | def _option_labels(options: list[Any]) -> list[str]: |
| 80 | labels: list[str] = [] |
| 81 | for option in options: |
| 82 | if isinstance(option, dict): |
| 83 | labels.append(str(option.get("label") or "").strip()) |
| 84 | else: |
| 85 | labels.append(str(option).strip()) |
| 86 | return [label for label in labels if label] |
| 87 | |
| 88 | |
| 89 | def is_story_direction_card( |
| 90 | *, |
| 91 | card_id: str, |
| 92 | question: str, |
| 93 | option_labels: list[str], |
| 94 | ) -> bool: |
| 95 | """True for confirm-story / story-premise cards that must offer reference-image edit.""" |
| 96 | if any(marker in question for marker in _NOT_STORY_DIRECTION_QUESTION_MARKERS): |
| 97 | return False |
| 98 | if _UPLOAD_GATE_OPTIONS.intersection(option_labels): |
| 99 | return False |
| 100 | if card_id in _STORY_DIRECTION_CARD_IDS: |
| 101 | return True |
| 102 | if any(marker in question for marker in _STORY_DIRECTION_QUESTION_MARKERS): |
| 103 | return True |
| 104 | return bool(_STORY_DIRECTION_OPTION_MARKERS.intersection(option_labels)) |
| 105 | |
| 106 | |
| 107 | def ensure_reference_image_edit_option(card: dict[str, Any]) -> dict[str, Any]: |
| 108 | """Append the canonical edit-reference option when a story-direction card omitted it.""" |
| 109 | options = card.get("options") |
| 110 | if not isinstance(options, list): |
| 111 | return card |
| 112 | labels = _option_labels(options) |
| 113 | if any(label in _REFERENCE_IMAGE_EDIT_ALIASES for label in labels): |
| 114 | return card |
| 115 | if not is_story_direction_card( |
| 116 | card_id=str(card.get("id") or ""), |
| 117 | question=str(card.get("question") or ""), |
| 118 | option_labels=labels, |
| 119 | ): |
| 120 | return card |
| 121 | card["options"] = [*options, {"label": REFERENCE_IMAGE_EDIT_OPTION}] |
| 122 | return card |
| 123 | |
| 124 | |
| 125 | def normalize_question_cards(raw_questions: list[Any]) -> list[dict[str, Any]] | str: |
| 126 | """Normalize and validate question payloads for WebUI + session storage.""" |
| 127 | if not isinstance(raw_questions, list) or not raw_questions: |
| 128 | return "Error: questions must be a non-empty list" |
| 129 | |
| 130 | normalized: list[dict[str, Any]] = [] |
| 131 | for index, item in enumerate(raw_questions): |
| 132 | if not isinstance(item, dict): |
| 133 | return f"Error: questions[{index}] must be an object" |
| 134 | |
| 135 | question = item.get("question") |
| 136 | if not isinstance(question, str) or not question.strip(): |
| 137 | return f"Error: questions[{index}].question must be a non-empty string" |
| 138 | |
| 139 | options_raw = item.get("options") |
| 140 | if not isinstance(options_raw, list) or not options_raw: |
| 141 | return f"Error: questions[{index}].options must be a non-empty list" |
| 142 | |
| 143 | options: list[dict[str, str]] = [] |
| 144 | for opt_index, opt in enumerate(options_raw): |
| 145 | if isinstance(opt, str) and opt.strip(): |
| 146 | options.append({"label": opt.strip()}) |
| 147 | continue |
| 148 | if isinstance(opt, dict) and isinstance(opt.get("label"), str) and opt["label"].strip(): |
| 149 | options.append({"label": opt["label"].strip()}) |
| 150 | continue |
| 151 | return f"Error: questions[{index}].options[{opt_index}] must be a string or {{label}}" |
| 152 | |
| 153 | card_id = item.get("id") |
| 154 | if not isinstance(card_id, str) or not card_id.strip(): |
| 155 | card_id = f"q-{index}" |
| 156 | |
| 157 | allow_custom = item.get("allow_custom", item.get("allowCustom", False)) is True |
| 158 | |
| 159 | normalized.append( |
| 160 | ensure_reference_image_edit_option( |
| 161 | { |
| 162 | "id": card_id.strip(), |
| 163 | "question": question.strip(), |
| 164 | "options": options, |
| 165 | "allow_custom": allow_custom, |
| 166 | "status": "pending", |
| 167 | "answered": None, |
| 168 | } |
| 169 | ) |
| 170 | ) |
| 171 | |
| 172 | return normalized |
| 173 | |
| 174 | |
| 175 | @tool_parameters( |
| 176 | tool_parameters_schema( |
| 177 | content=StringSchema( |
| 178 | "Intro text shown above the cards. Do not repeat option labels here." |
| 179 | ), |
| 180 | questions=ArraySchema( |
| 181 | _QUESTION_SCHEMA, |
| 182 | description="Structured question cards for the user to tap", |
| 183 | min_items=1, |
| 184 | ), |
| 185 | required=["content", "questions"], |
| 186 | ) |
| 187 | ) |
| 188 | class AskUserTool(Tool): |
| 189 | """Present structured multiple-choice cards in the WebUI.""" |
| 190 | |
| 191 | def __init__( |
| 192 | self, |
| 193 | send_callback: Callable[[OutboundMessage], Awaitable[None]] | None = None, |
| 194 | persist_callback: PersistCallback | None = None, |
| 195 | default_channel: str = "", |
| 196 | default_chat_id: str = "", |
| 197 | default_session_key: str = "", |
| 198 | ): |
| 199 | self._send_callback = send_callback |
| 200 | self._persist_callback = persist_callback |
| 201 | self._default_channel: ContextVar[str] = ContextVar( |
| 202 | "ask_user_default_channel", default=default_channel |
| 203 | ) |
| 204 | self._default_chat_id: ContextVar[str] = ContextVar( |
| 205 | "ask_user_default_chat_id", default=default_chat_id |
| 206 | ) |
| 207 | self._default_session_key: ContextVar[str] = ContextVar( |
| 208 | "ask_user_default_session_key", default=default_session_key |
| 209 | ) |
| 210 | self._tool_call_id: ContextVar[str] = ContextVar("ask_user_tool_call_id", default="") |
| 211 | self._sent_in_turn_var: ContextVar[bool] = ContextVar( |
| 212 | "ask_user_sent_in_turn", default=False |
| 213 | ) |
| 214 | |
| 215 | def set_context( |
| 216 | self, |
| 217 | channel: str, |
| 218 | chat_id: str, |
| 219 | *, |
| 220 | session_key: str | None = None, |
| 221 | ) -> None: |
| 222 | self._default_channel.set(channel) |
| 223 | self._default_chat_id.set(chat_id) |
| 224 | if session_key: |
| 225 | self._default_session_key.set(session_key) |
| 226 | |
| 227 | def set_tool_call_id(self, tool_call_id: str) -> None: |
| 228 | self._tool_call_id.set(tool_call_id.strip()) |
| 229 | |
| 230 | def set_send_callback(self, callback: Callable[[OutboundMessage], Awaitable[None]]) -> None: |
| 231 | self._send_callback = callback |
| 232 | |
| 233 | def set_persist_callback(self, callback: PersistCallback) -> None: |
| 234 | self._persist_callback = callback |
| 235 | |
| 236 | def start_turn(self) -> None: |
| 237 | self._sent_in_turn = False |
| 238 | |
| 239 | @property |
| 240 | def _sent_in_turn(self) -> bool: |
| 241 | return self._sent_in_turn_var.get() |
| 242 | |
| 243 | @_sent_in_turn.setter |
| 244 | def _sent_in_turn(self, value: bool) -> None: |
| 245 | self._sent_in_turn_var.set(value) |
| 246 | |
| 247 | @property |
| 248 | def name(self) -> str: |
| 249 | return "ask_user" |
| 250 | |
| 251 | @property |
| 252 | def description(self) -> str: |
| 253 | return ( |
| 254 | "Ask the user one or more multiple-choice questions using interactive " |
| 255 | "cards in the WebUI. Use this instead of listing options as plain text. " |
| 256 | "Put the intro in `content` and each question in `questions` with " |
| 257 | "`options` labels. The user's tap is persisted and survives page refresh." |
| 258 | ) |
| 259 | |
| 260 | async def execute( |
| 261 | self, |
| 262 | content: str, |
| 263 | questions: list[Any] | None = None, |
| 264 | channel: str | None = None, |
| 265 | chat_id: str | None = None, |
| 266 | **kwargs: Any, |
| 267 | ) -> str: |
| 268 | from nanobot.utils.helpers import strip_think |
| 269 | |
| 270 | content = strip_think(content or "") |
| 271 | if not content.strip(): |
| 272 | return "Error: content must be a non-empty string" |
| 273 | |
| 274 | normalized = normalize_question_cards(questions or []) |
| 275 | if isinstance(normalized, str): |
| 276 | return normalized |
| 277 | |
| 278 | channel = channel or self._default_channel.get() |
| 279 | chat_id = chat_id or self._default_chat_id.get() |
| 280 | session_key = self._default_session_key.get() |
| 281 | |
| 282 | if not channel or not chat_id: |
| 283 | return "Error: No target channel/chat specified" |
| 284 | if not self._send_callback: |
| 285 | return "Error: ask_user delivery not configured" |
| 286 | |
| 287 | batch_id = str(uuid.uuid4()) |
| 288 | metadata: dict[str, Any] = { |
| 289 | "questions": normalized, |
| 290 | "question_batch_id": batch_id, |
| 291 | } |
| 292 | |
| 293 | try: |
| 294 | if self._persist_callback and session_key: |
| 295 | tool_call_id = self._tool_call_id.get().strip() |
| 296 | if not tool_call_id: |
| 297 | tool_call_id = f"call_{batch_id}" |
| 298 | await self._persist_callback( |
| 299 | session_key, |
| 300 | tool_call_id, |
| 301 | batch_id, |
| 302 | content.strip(), |
| 303 | normalized, |
| 304 | channel, |
| 305 | chat_id, |
| 306 | ) |
| 307 | await self._send_callback( |
| 308 | OutboundMessage( |
| 309 | channel=channel, |
| 310 | chat_id=chat_id, |
| 311 | content="", |
| 312 | metadata=metadata, |
| 313 | ) |
| 314 | ) |
| 315 | if channel == self._default_channel.get() and chat_id == self._default_chat_id.get(): |
| 316 | self._sent_in_turn = True |
| 317 | return ( |
| 318 | f"Question cards sent to {channel}:{chat_id} " |
| 319 | f"({len(normalized)} card(s), batch={batch_id})" |
| 320 | ) |
| 321 | except Exception as e: |
| 322 | return f"Error sending question cards: {e}" |
| 323 | |
| 324 |