from __future__ import annotations import asyncio import json import os import uuid from concurrent.futures import ThreadPoolExecutor from pathlib import Path from typing import Any from PIL import Image from shared.deepy import vision as deepy_vision from shared.deepy import video_tools as deepy_video_tools from shared.deepy.engine import ( assistant_steering_interrupt_due, begin_assistant_pause, begin_assistant_action, begin_assistant_thought, checkpoint_assistant_thought, checkpoint_assistant_turn, clear_assistant_pause, clear_assistant_steering, finish_assistant_action, finish_assistant_thought, finish_assistant_turn, mark_assistant_paused, mark_assistant_turn_message, rollback_assistant_turn, interrupt_assistant_for_steering, wait_for_assistant_resume, ) from shared.gradio import assistant_chat from .base import BackendEvent from .config import ENGINE_DISABLED, ENGINE_LABELS, is_remote_engine, resolve_role_engine from .images import temporary_image_paths from .registry import create_backend from .usage import build_remote_usage_stats def _send(send_cmd, payload) -> None: if payload is not None: send_cmd("chat_output", payload) def _visual_query(server_config: dict[str, Any], media_record, question: str, frame_no: int | None = None, max_image_edge: int | None = None, file_access_policy=None) -> dict[str, Any]: try: asyncio.get_running_loop() except RuntimeError: return _visual_query_without_running_loop(server_config, media_record, question, frame_no, max_image_edge, file_access_policy) with ThreadPoolExecutor(max_workers=1, thread_name_prefix="wangp-remote-vision") as executor: return executor.submit(_visual_query_without_running_loop, server_config, media_record, question, frame_no, max_image_edge, file_access_policy).result() def _visual_query_without_running_loop(server_config: dict[str, Any], media_record, question: str, frame_no: int | None = None, max_image_edge: int | None = None, file_access_policy=None) -> dict[str, Any]: engine = resolve_role_engine(server_config, "visual_inspector") if engine == ENGINE_DISABLED: return {"status": "error", "question": question, "answer": "", "error": "Visual Inspector is disabled in Configuration."} if not is_remote_engine(engine): return {"status": "error", "question": question, "answer": "", "error": "A remote Deepy engine requires a remote Visual Inspector. Choose Auto or Same as Deepy."} records = list(media_record) if isinstance(media_record, list) else [media_record] video_max_pixels = None if max_image_edge is None else int(max_image_edge) * int(max_image_edge) max_image_edge = deepy_vision.VISION_REMOTE_MAX_IMAGE_EDGE if max_image_edge is None else int(max_image_edge) images, inspected = [None] * len(records), [] video_inputs: dict[str, list[tuple[int, int, list[int] | None]]] = {} for input_index, record in enumerate(records): path = str(record.get("path", "") or "").strip() public_path = file_access_policy.virtualize_path(path) if file_access_policy is not None else "" if file_access_policy is not None and file_access_policy.virtualized and not public_path.startswith("@"): public_path = str(record.get("media_id", "") or public_path) if not os.path.isfile(path): raise FileNotFoundError(f"Media file not found: {public_path or path}") media_type = str(record.get("media_type", "") or "").lower() bbox = record.get("bbox", None) resolved_frame = None time_seconds = record.get("time_seconds", None) if media_type == "video": requested_frame = record.get("frame_no", frame_no) resolved_frame = deepy_video_tools.resolve_video_frame_no(path, frame_no=requested_frame, time_seconds=time_seconds) if requested_frame is not None or time_seconds is not None else 0 video_inputs.setdefault(path, []).append((input_index, resolved_frame, bbox)) else: with Image.open(path) as source: images[input_index] = deepy_vision.prepare_inspection_image(source, max_edge=max_image_edge, bbox=bbox) label = file_access_policy.virtualize_result(str(record.get("label", "") or "")) if file_access_policy is not None else record.get("label", "") if file_access_policy is not None and file_access_policy.virtualized: for physical_path in {path, str(Path(path).resolve()), Path(path).as_posix()}: label = str(label).replace(physical_path, public_path) inspected.append({"media_id": record.get("media_id", ""), "media_type": media_type, "label": label, "frame_no": resolved_frame, "time_seconds": time_seconds if media_type == "video" else None, "bbox": bbox, **({"path": public_path} if public_path else {})}) for path, indexed_frames in video_inputs.items(): bboxes = [item[2] for item in indexed_frames] decode_kwargs = {**({"max_pixels": video_max_pixels} if video_max_pixels is not None else {"max_edge": max_image_edge}), **({"bboxes": bboxes} if any(bbox is not None for bbox in bboxes) else {})} decoded_images = deepy_vision.decode_inspection_video_frames(path, [item[1] for item in indexed_frames], **decode_kwargs) for (input_index, _resolved_frame, _bbox), decoded_image in zip(indexed_frames, decoded_images): images[input_index] = decoded_image visual_labels = [] for index, item in enumerate(inspected): source_label = str(item.get("label", "") or os.path.basename(str(records[index].get("path", "")))).strip() public_path = str(item.get("path", "") or "").strip() if public_path and public_path.casefold() != source_label.casefold(): source_label = f"{source_label} ({public_path})" bbox_label = "" if item["bbox"] is None else f", bbox {item['bbox']}" if item["media_type"] == "video": time_label = "" if item["time_seconds"] is None else f" at {float(item['time_seconds']):.3f} seconds" visual_labels.append(f"Visual {index + 1}: video {source_label}, frame {item['frame_no']}{time_label}{bbox_label}.") else: visual_labels.append(f"Visual {index + 1}: image {source_label}{bbox_label}.") labeled_question = "\n".join([*visual_labels, "", str(question or "").strip()]) backend = create_backend(engine, server_config) try: with temporary_image_paths(images) as paths: answer = backend.one_shot(labeled_question, system_prompt=deepy_vision.VISION_QA_SYSTEM_PROMPT, images=paths, max_output_tokens=deepy_vision.VISION_ANSWER_MAX_NEW_TOKENS) finally: backend.close() result = {"status": "done", "media_ids": [item["media_id"] for item in inspected], "media": inspected, "question": question, "answer": answer, "error": ""} if len(inspected) == 1: result.update(inspected[0]) return result def run_remote_deepy_turn(server_config: dict[str, Any], session, text: str, system_prompt: str, toolbox, send_cmd) -> None: engine = resolve_role_engine(server_config, "deepy") backends = session.remote_backends backend = backends.get(engine) if backend is None: backend = create_backend(engine, server_config, toolbox=toolbox) backends[engine] = backend assistant_id = assistant_chat.create_assistant_turn(session) mark_assistant_turn_message(session, assistant_id) assistant_badge = str(session.current_turn.get("assistant_badge", "") or "") if isinstance(session.current_turn, dict) else "" if assistant_badge: _send(send_cmd, assistant_chat.set_message_badge(session, assistant_id, assistant_badge)) user_message = assistant_chat.build_user_model_message(session, text, toolbox=toolbox) session.messages.append(user_message) checkpoint_assistant_turn(session) answer_parts: list[str] = [] answer_segment_parts: list[str] = [] answer_block_id = "" commentary_messages: dict[str, dict[str, Any]] = {} reasoning_parts: list[str] = [] active_tool: dict[str, str] = {} reasoning_active = False reasoning_source = "" reasoning_block_id = "" reasoning_block_no = 0 status_phase = "" engine_label = ENGINE_LABELS.get(engine, str(engine or "Remote LLM").strip() or "Remote LLM") def set_remote_status(phase: str, text: str, kind: str) -> None: nonlocal status_phase if status_phase == phase: return status_phase = phase _send(send_cmd, assistant_chat.build_status_event(text, kind=kind, session=session)) def wait_while_paused() -> bool: nonlocal status_phase if not session.pause_requested: return not session.interrupt_requested if not begin_assistant_pause(session): return not session.interrupt_requested if not mark_assistant_paused(session): return not session.interrupt_requested _send(send_cmd, assistant_chat.build_status_event("Deepy is paused.", kind="paused", session=session)) if not wait_for_assistant_resume(session): return False _send(send_cmd, assistant_chat.build_status_event(f"Resuming {engine_label}...", kind="resuming", session=session)) status_phase = "" return True def finish_reasoning() -> None: nonlocal reasoning_active if reasoning_active: completed_reasoning = "".join(reasoning_parts).strip() if reasoning_block_id and reasoning_parts: _reasoning_id, payload = assistant_chat.upsert_reasoning_block(session, assistant_id, reasoning_block_id, completed_reasoning, streaming=False) _send(send_cmd, payload) finish_assistant_thought(session) reasoning_active = False if completed_reasoning: checkpoint_assistant_thought(session, f"\n{completed_reasoning}\n") def finish_answer_segment() -> None: nonlocal answer_block_id if answer_block_id and answer_segment_parts: answer_block_id, payload = assistant_chat.upsert_assistant_content_block(session, assistant_id, answer_block_id, "".join(answer_segment_parts), streaming=False) _send(send_cmd, payload) answer_block_id = "" answer_segment_parts.clear() def on_event(event: BackendEvent) -> None: nonlocal answer_block_id, reasoning_active, reasoning_source, reasoning_block_id, reasoning_block_no, reasoning_parts if not wait_while_paused(): return if event.kind in {"commentary_delta", "commentary_replace"} and event.text: finish_reasoning() finish_answer_segment() data = event.data if isinstance(event.data, dict) else {} item_id = str(data.get("item_id", "") or "").strip() or "commentary_stream" commentary = commentary_messages.get(item_id) if commentary is None: commentary = {"block_id": "", "parts": [event.text]} commentary_messages[item_id] = commentary else: if event.kind == "commentary_replace": commentary["parts"][:] = [event.text] else: commentary["parts"].append(event.text) commentary["block_id"], payload = assistant_chat.upsert_assistant_content_block(session, assistant_id, commentary["block_id"], "".join(commentary["parts"])) _send(send_cmd, payload) set_remote_status("thinking", f"{engine_label} is thinking...", "thinking") elif event.kind == "commentary_remove": data = event.data if isinstance(event.data, dict) else {} item_id = str(data.get("item_id", "") or "").strip() commentary = commentary_messages.pop(item_id, None) if commentary and commentary["block_id"]: _send(send_cmd, assistant_chat.remove_assistant_content_block(session, assistant_id, commentary["block_id"])) elif event.kind == "commentary_promote" and event.text: finish_reasoning() finish_answer_segment() data = event.data if isinstance(event.data, dict) else {} item_id = str(data.get("item_id", "") or "").strip() or "commentary_stream" commentary = commentary_messages.pop(item_id, None) if commentary is None: _block_id, payload = assistant_chat.upsert_assistant_content_block(session, assistant_id, None, event.text, streaming=False) _send(send_cmd, payload) else: commentary["parts"][:] = [event.text] commentary["block_id"], payload = assistant_chat.upsert_assistant_content_block(session, assistant_id, commentary["block_id"], event.text, streaming=False) _send(send_cmd, payload) answer_parts.append(event.text) set_remote_status("responding", f"{engine_label} is responding...", "status") elif event.kind == "text_delta" and event.text: if reasoning_active: finish_reasoning() if session.steering_pending: interrupt_assistant_for_steering(session) return set_remote_status("responding", f"{engine_label} is responding...", "status") answer_parts.append(event.text) answer_segment_parts.append(event.text) answer_block_id, payload = assistant_chat.upsert_assistant_content_block(session, assistant_id, answer_block_id, "".join(answer_segment_parts)) _send(send_cmd, payload) elif event.kind == "text_replace" and event.text: finish_reasoning() set_remote_status("responding", f"{engine_label} is responding...", "status") answer_parts[:] = [event.text] answer_segment_parts[:] = [event.text] answer_block_id, payload = assistant_chat.upsert_assistant_content_block(session, assistant_id, answer_block_id, event.text) _send(send_cmd, payload) elif event.kind == "reasoning_delta" and event.text: finish_answer_segment() data = event.data if isinstance(event.data, dict) else {} item_id = str(data.get("item_id", "") or "").strip() summary_index = data.get("summary_index") source = f"{item_id}:{summary_index}" if item_id or summary_index is not None else "stream" event_text = event.text if not reasoning_active: begin_assistant_thought(session) reasoning_active = True reasoning_block_no += 1 reasoning_block_id = f"remote_reasoning_{reasoning_block_no}" reasoning_parts = [] elif source != reasoning_source and reasoning_parts: reasoning_parts.append("\n\n") event_text = event.text.lstrip() else: event_text = event.text reasoning_source = source set_remote_status("thinking", f"{engine_label} is thinking...", "thinking") reasoning_parts.append(event_text) _reasoning_id, payload = assistant_chat.upsert_reasoning_block(session, assistant_id, reasoning_block_id, "".join(reasoning_parts)) _send(send_cmd, payload) elif event.kind == "reasoning_start": finish_answer_segment() set_remote_status("thinking", f"{engine_label} is thinking...", "thinking") elif event.kind == "tool_request_start": finish_reasoning() finish_answer_segment() set_remote_status("tool-request", f"{engine_label} is preparing a tool request...", "tool") elif event.kind == "tool_request_error": finish_reasoning() finish_answer_segment() data = event.data if isinstance(event.data, dict) else {} tool_name = str(data.get("name", "") or "").removeprefix("mcp__wangp__") or "remote_tool" arguments = dict(data.get("input", {}) or {}) tool_label = f"{toolbox.get_tool_display_name(tool_name)} Request" ui_tool_id, payload = assistant_chat.add_tool_call(session, assistant_id, tool_name, arguments, tool_label=tool_label) _send(send_cmd, payload) result = {"status": "error", "tool": tool_name, "error": event.text or "The remote provider rejected this tool request before execution."} _send(send_cmd, assistant_chat.complete_tool_call(session, assistant_id, ui_tool_id, result)) checkpoint_assistant_turn(session) set_remote_status("waiting", f"Waiting for {engine_label}...", "loading") elif event.kind == "usage": stats = build_remote_usage_stats(event.data) if stats is not None: session.remote_usage_stats = stats _send(send_cmd, assistant_chat.build_stats_event(stats)) elif event.kind == "compaction": finish_reasoning() finish_answer_segment() data = event.data if isinstance(event.data, dict) else {} before_tokens = data.get("pre_tokens", data.get("preTokens", None)) try: before_label = f" It occurred near {int(before_tokens):,} tokens." if int(before_tokens or 0) > 0 else "" except (TypeError, ValueError): before_label = "" detail = "The remote LLM compacted earlier conversation context. Its internal summary is not exposed to WanGP." + before_label _block_id, payload = assistant_chat.add_context_summary(session, assistant_id, detail) _send(send_cmd, payload) checkpoint_assistant_turn(session) set_remote_status("compaction", f"{engine_label} compacted its context...", "loading") def tool_progress(status=None, status_text=None, result=None): if active_tool.get("id"): safe_result = toolbox.file_access_policy.virtualize_result(result if result is not None else {}) _send(send_cmd, assistant_chat.update_tool_call(session, assistant_id, active_tool["id"], status=status, status_text=status_text, result=safe_result)) if status_text: set_remote_status(f"tool:{active_tool['id']}:{status}:{status_text}", f"{active_tool['label']}: {status_text}", "tool") def call_tool(name: str, arguments: dict[str, Any]) -> dict[str, Any]: nonlocal reasoning_active if not wait_while_paused(): return {"status": "interrupted", "tool": name, "cancelled": True, "error": "The user interrupted before this action started."} finish_answer_segment() if reasoning_active: finish_reasoning() validation_error = toolbox.validate_tool_call(name, arguments) if validation_error: return {"status": "error", "tool": name, "error": validation_error} if not begin_assistant_action(session): return {"status": "interrupted", "tool": name, "cancelled": True, "error": "The user interrupted before this action started."} tool_call_id = f"remote_{uuid.uuid4().hex}" tool_label = toolbox.get_tool_transcript_label(name, arguments) ui_tool_id, payload = assistant_chat.add_tool_call(session, assistant_id, name, arguments, tool_label=tool_label) active_tool.update({"id": ui_tool_id, "label": tool_label}) _send(send_cmd, payload) set_remote_status(f"tool:{ui_tool_id}", f"{tool_label}...", "tool") session.messages.append({"role": "assistant", "content": "", "tool_calls": [{"id": tool_call_id, "type": "function", "function": {"name": name, "arguments": json.dumps(arguments, ensure_ascii=False)}}]}) try: result = toolbox.call(name, arguments) except Exception as exc: result = toolbox.file_access_policy.virtualize_result({"status": "error", "tool": name, "error": str(exc)}) session.messages.append({"role": "tool", "tool_call_id": tool_call_id, "content": json.dumps(result, ensure_ascii=False)}) _send(send_cmd, assistant_chat.complete_tool_call(session, assistant_id, ui_tool_id, result)) checkpoint_assistant_turn(session) active_tool.clear() finish_assistant_action(session) wait_while_paused() if not session.interrupt_requested: set_remote_status("waiting", f"Waiting for {engine_label}...", "loading") return result toolbox.bind_runtime_tools(vision_query_callback=lambda record, question, frame=None, max_image_edge=None: _visual_query(server_config, record, question, frame, max_image_edge, toolbox.file_access_policy), tool_progress_callback=tool_progress, vision_is_remote=True) execution_instructions_getter = getattr(toolbox, "get_remote_execution_instructions", None) execution_instructions = str(execution_instructions_getter() or "").strip() if callable(execution_instructions_getter) else "" if execution_instructions: system_prompt = f"{system_prompt}\n\n{execution_instructions}".strip() try: set_remote_status("waiting", f"Waiting for {engine_label}...", "loading") answer = backend.run_turn(user_message.get("model_content", user_message["content"]), system_prompt=system_prompt, tools=toolbox.get_tool_schemas(), images=[], on_event=on_event, call_tool=call_tool, should_stop=lambda: bool(session.interrupt_requested or assistant_steering_interrupt_due(session))) if answer and not answer_parts: answer_parts.append(answer) _answer_block_id, payload = assistant_chat.upsert_assistant_content_block(session, assistant_id, None, answer, streaming=False) _send(send_cmd, payload) finish_reasoning() finish_answer_segment() for commentary in commentary_messages.values(): if commentary["block_id"] and commentary["parts"]: commentary["block_id"], payload = assistant_chat.upsert_assistant_content_block(session, assistant_id, commentary["block_id"], "".join(commentary["parts"]), streaming=False) _send(send_cmd, payload) if not session.interrupt_requested: final_answer = "".join(answer_parts).strip() if final_answer: session.messages.append({"role": "assistant", "content": final_answer}) checkpoint_assistant_turn(session) _send(send_cmd, assistant_chat.linkify_message_download_references(session, assistant_id, getattr(toolbox, "file_access_policy", None))) except Exception as exc: if not bool(getattr(exc, "preserve_backend", False)): if backends.get(engine) is backend: backends.pop(engine, None) backend.close() raise finally: finish_reasoning() with session.turn_lock: if session.interrupt_requested: backend.interrupt() rollback_assistant_turn(session) finish_assistant_turn(session) clear_assistant_steering(session) clear_assistant_pause(session) _send(send_cmd, assistant_chat.build_status_event(None, visible=False, session=session))