| 1 | import json |
| 2 | import logging |
| 3 | import os |
| 4 | import threading |
| 5 | import uuid |
| 6 | from datetime import datetime |
| 7 | from typing import List |
| 8 | |
| 9 | from fastapi import APIRouter |
| 10 | from fastapi.concurrency import run_in_threadpool |
| 11 | |
| 12 | from api.schemas.sandbox import ( |
| 13 | SandboxI2IRequest, |
| 14 | SandboxLLMRequest, |
| 15 | SandboxT2IRequest, |
| 16 | SandboxVLMRequest, |
| 17 | SandboxVideoRequest, |
| 18 | ) |
| 19 | from config import settings |
| 20 | |
| 21 | router = APIRouter(tags=["Sandbox"]) |
| 22 | logger = logging.getLogger(__name__) |
| 23 | |
| 24 | |
| 25 | SANDBOX_DIR = os.path.join(settings.CODE_DIR, "result", "sandbox") |
| 26 | SANDBOX_HISTORY_FILE = os.path.join(SANDBOX_DIR, "history.json") |
| 27 | SANDBOX_ACTIVE_TASKS: dict[str, dict] = {} |
| 28 | SANDBOX_LOCK = threading.RLock() |
| 29 | |
| 30 | # 确保目录存在 |
| 31 | os.makedirs(SANDBOX_DIR, exist_ok=True) |
| 32 | |
| 33 | |
| 34 | def _load_history() -> List[dict]: |
| 35 | """加载历史记录""" |
| 36 | with SANDBOX_LOCK: |
| 37 | if os.path.exists(SANDBOX_HISTORY_FILE): |
| 38 | try: |
| 39 | with open(SANDBOX_HISTORY_FILE, 'r', encoding='utf-8') as f: |
| 40 | return json.load(f) |
| 41 | except Exception: |
| 42 | logger.warning("Failed to load sandbox history: %s", SANDBOX_HISTORY_FILE, exc_info=True) |
| 43 | return [] |
| 44 | return [] |
| 45 | |
| 46 | |
| 47 | def _save_history(history: List[dict]): |
| 48 | """保存历史记录""" |
| 49 | with SANDBOX_LOCK: |
| 50 | tmp_path = f"{SANDBOX_HISTORY_FILE}.{uuid.uuid4().hex}.tmp" |
| 51 | with open(tmp_path, 'w', encoding='utf-8') as f: |
| 52 | json.dump(history, f, ensure_ascii=False, indent=2) |
| 53 | os.replace(tmp_path, SANDBOX_HISTORY_FILE) |
| 54 | |
| 55 | |
| 56 | def _normalize_path(path: str) -> str: |
| 57 | """将绝对路径转换为相对路径格式 result/...""" |
| 58 | if not path: |
| 59 | return path |
| 60 | # 如果已经是相对路径，直接返回 |
| 61 | if not path.startswith('/'): |
| 62 | # 确保以 result/ 开头 |
| 63 | if not path.startswith('result/'): |
| 64 | return f"result/{path}" |
| 65 | return path |
| 66 | # 绝对路径，提取相对于 CODE_DIR 的部分 |
| 67 | code_dir = settings.CODE_DIR |
| 68 | if path.startswith(code_dir): |
| 69 | relative = path[len(code_dir):].lstrip('/') |
| 70 | # 直接返回 result/... 格式，因为 /code/ 会映射到 CODE_DIR |
| 71 | return relative |
| 72 | # 其他绝对路径，尝试提取文件名 |
| 73 | return path.split('/')[-1] |
| 74 | |
| 75 | |
| 76 | def _convert_output_paths(output_data: dict) -> dict: |
| 77 | """转换 output 中的路径为相对路径格式""" |
| 78 | if not output_data: |
| 79 | return output_data |
| 80 | converted = output_data.copy() |
| 81 | # 转换 images |
| 82 | if 'images' in converted and isinstance(converted['images'], list): |
| 83 | converted['images'] = [_normalize_path(img) for img in converted['images']] |
| 84 | # 转换 video_path |
| 85 | if 'video_path' in converted and converted['video_path']: |
| 86 | converted['video_path'] = _normalize_path(converted['video_path']) |
| 87 | # 转换 input 中的 reference_image |
| 88 | if 'reference_image' in converted.get('input', {}): |
| 89 | input_copy = converted['input'].copy() |
| 90 | input_copy['reference_image'] = _normalize_path(input_copy['reference_image']) |
| 91 | converted['input'] = input_copy |
| 92 | return converted |
| 93 | |
| 94 | |
| 95 | def _converted_result_list(paths: List[str] | None) -> List[str]: |
| 96 | """Return generated image paths in the same format used by sandbox history.""" |
| 97 | converted = _convert_output_paths({"images": paths or []}) |
| 98 | return converted.get("images", []) |
| 99 | |
| 100 | |
| 101 | def _converted_video_path(path: str | None) -> str: |
| 102 | """Return generated video path in the same format used by sandbox history.""" |
| 103 | converted = _convert_output_paths({"video_path": path or ""}) |
| 104 | return converted.get("video_path", "") |
| 105 | |
| 106 | |
| 107 | def _add_record( |
| 108 | tool: str, |
| 109 | model: str, |
| 110 | input_data: dict, |
| 111 | output_data: dict, |
| 112 | files: List[str] = None, |
| 113 | record_id: str | None = None, |
| 114 | ) -> str: |
| 115 | """添加历史记录""" |
| 116 | with SANDBOX_LOCK: |
| 117 | record_id = record_id or str(uuid.uuid4().hex[:8]) |
| 118 | # 转换路径为相对路径格式 |
| 119 | output_data = _convert_output_paths(output_data) |
| 120 | record = { |
| 121 | "id": record_id, |
| 122 | "tool": tool, |
| 123 | "model": model, |
| 124 | "input": input_data, |
| 125 | "output": output_data, |
| 126 | "files": files or [], |
| 127 | "created_at": datetime.now().isoformat(), |
| 128 | } |
| 129 | history = _load_history() |
| 130 | history.insert(0, record) # 最新记录放在最前面 |
| 131 | _save_history(history) |
| 132 | return record_id |
| 133 | |
| 134 | |
| 135 | def _start_active_task(tool: str, model: str, input_data: dict) -> str: |
| 136 | with SANDBOX_LOCK: |
| 137 | task_id = str(uuid.uuid4().hex[:8]) |
| 138 | SANDBOX_ACTIVE_TASKS[task_id] = { |
| 139 | "id": task_id, |
| 140 | "tool": tool, |
| 141 | "model": model, |
| 142 | "input": input_data, |
| 143 | "status": "running", |
| 144 | "progress": 1, |
| 145 | "created_at": datetime.now().isoformat(), |
| 146 | } |
| 147 | return task_id |
| 148 | |
| 149 | |
| 150 | def _finish_active_task(task_id: str) -> None: |
| 151 | with SANDBOX_LOCK: |
| 152 | SANDBOX_ACTIVE_TASKS.pop(task_id, None) |
| 153 | |
| 154 | |
| 155 | def _delete_record_files(files: List[str]): |
| 156 | """删除记录关联的文件""" |
| 157 | for f in files: |
| 158 | if f and os.path.exists(f): |
| 159 | try: |
| 160 | os.remove(f) |
| 161 | except Exception: |
| 162 | logger.warning("Failed to delete sandbox artifact: %s", f, exc_info=True) |
| 163 | pass |
| 164 | |
| 165 | |
| 166 | # 请求模型 |
| 167 | @router.get("/api/sandbox/history") |
| 168 | async def sandbox_get_history(): |
| 169 | """获取历史记录列表""" |
| 170 | history = _load_history() |
| 171 | # 返回完整信息（包括 output） |
| 172 | return { |
| 173 | "success": True, |
| 174 | "records": [ |
| 175 | { |
| 176 | "id": r["id"], |
| 177 | "tool": r["tool"], |
| 178 | "model": r["model"], |
| 179 | "input": r["input"], |
| 180 | "output": r.get("output"), |
| 181 | "created_at": r["created_at"], |
| 182 | } |
| 183 | for r in history |
| 184 | ] |
| 185 | } |
| 186 | |
| 187 | |
| 188 | @router.get("/api/sandbox/tasks") |
| 189 | async def sandbox_get_active_tasks(): |
| 190 | """获取临时工作台正在执行的任务""" |
| 191 | with SANDBOX_LOCK: |
| 192 | tasks = list(SANDBOX_ACTIVE_TASKS.values()) |
| 193 | return {"success": True, "tasks": tasks} |
| 194 | |
| 195 | |
| 196 | @router.get("/api/sandbox/history/{record_id}") |
| 197 | async def sandbox_get_record(record_id: str): |
| 198 | """获取单条历史记录详情""" |
| 199 | history = _load_history() |
| 200 | for r in history: |
| 201 | if r["id"] == record_id: |
| 202 | return {"success": True, "record": r} |
| 203 | return {"success": False, "error": "记录不存在"} |
| 204 | |
| 205 | |
| 206 | @router.delete("/api/sandbox/history/{record_id}") |
| 207 | async def sandbox_delete_record(record_id: str): |
| 208 | """删除历史记录""" |
| 209 | with SANDBOX_LOCK: |
| 210 | history = _load_history() |
| 211 | record_to_delete = None |
| 212 | new_history = [] |
| 213 | for r in history: |
| 214 | if r["id"] == record_id: |
| 215 | record_to_delete = r |
| 216 | else: |
| 217 | new_history.append(r) |
| 218 | |
| 219 | if record_to_delete is None: |
| 220 | return {"success": False, "error": "记录不存在"} |
| 221 | |
| 222 | _save_history(new_history) |
| 223 | |
| 224 | # 删除关联文件不需要占用历史锁。 |
| 225 | _delete_record_files(record_to_delete.get("files", [])) |
| 226 | logger.info("Sandbox history deleted: record_id=%s", record_id) |
| 227 | return {"success": True} |
| 228 | |
| 229 | |
| 230 | @router.post("/api/sandbox/llm") |
| 231 | async def sandbox_llm(req: SandboxLLMRequest): |
| 232 | """临时工作台 - LLM 文字生成""" |
| 233 | from models.llm_client import LLM |
| 234 | client = LLM() |
| 235 | input_data = {"prompt": req.prompt, "web_search": req.web_search} |
| 236 | task_id = _start_active_task("llm", req.model, input_data) |
| 237 | try: |
| 238 | logger.info("Sandbox LLM started: model=%s web_search=%s", req.model, req.web_search) |
| 239 | result = await run_in_threadpool( |
| 240 | client.query, |
| 241 | req.prompt, |
| 242 | model=req.model, |
| 243 | web_search=req.web_search, |
| 244 | ) |
| 245 | # ��存到历史记录 |
| 246 | record_id = _add_record( |
| 247 | tool="llm", |
| 248 | model=req.model, |
| 249 | input_data=input_data, |
| 250 | output_data={"response": result}, |
| 251 | record_id=task_id, |
| 252 | ) |
| 253 | logger.info("Sandbox LLM completed: model=%s record_id=%s", req.model, record_id) |
| 254 | return {"success": True, "result": result, "record_id": record_id} |
| 255 | except Exception as e: |
| 256 | logger.exception("Sandbox LLM failed: model=%s", req.model) |
| 257 | return {"success": False, "error": str(e)} |
| 258 | finally: |
| 259 | _finish_active_task(task_id) |
| 260 | |
| 261 | |
| 262 | @router.post("/api/sandbox/vlm") |
| 263 | async def sandbox_vlm(req: SandboxVLMRequest): |
| 264 | """临时工作台 - VLM 图片理解""" |
| 265 | from models.vlm_client import VLM |
| 266 | client = VLM() |
| 267 | input_data = {"prompt": req.prompt, "images": req.images} |
| 268 | task_id = _start_active_task("vlm", req.model, input_data) |
| 269 | try: |
| 270 | logger.info("Sandbox VLM started: model=%s images=%d", req.model, len(req.images or [])) |
| 271 | result = await run_in_threadpool( |
| 272 | client.query, |
| 273 | req.prompt, |
| 274 | image_paths=req.images, |
| 275 | model=req.model, |
| 276 | ) |
| 277 | # 保存到历史记录 |
| 278 | record_id = _add_record( |
| 279 | tool="vlm", |
| 280 | model=req.model, |
| 281 | input_data=input_data, |
| 282 | output_data={"response": result}, |
| 283 | record_id=task_id, |
| 284 | ) |
| 285 | logger.info("Sandbox VLM completed: model=%s record_id=%s", req.model, record_id) |
| 286 | return {"success": True, "result": result, "record_id": record_id} |
| 287 | except Exception as e: |
| 288 | logger.exception("Sandbox VLM failed: model=%s", req.model) |
| 289 | return {"success": False, "error": str(e)} |
| 290 | finally: |
| 291 | _finish_active_task(task_id) |
| 292 | |
| 293 | |
| 294 | @router.post("/api/sandbox/t2i") |
| 295 | async def sandbox_t2i(req: SandboxT2IRequest): |
| 296 | """临时工作台 - 文生图""" |
| 297 | from models.image_client import ImageClient |
| 298 | client = ImageClient() |
| 299 | input_data = {"prompt": req.prompt, "style": req.style, "ratio": req.ratio} |
| 300 | task_id = _start_active_task("t2i", req.model, input_data) |
| 301 | try: |
| 302 | logger.info("Sandbox T2I started: model=%s ratio=%s", req.model, req.ratio) |
| 303 | result = await run_in_threadpool( |
| 304 | client.generate_image, |
| 305 | req.prompt, |
| 306 | model=req.model, |
| 307 | image_paths=None, |
| 308 | video_ratio=req.ratio, |
| 309 | ) |
| 310 | # result 是图片路径列表 |
| 311 | # 保存到历史记录 |
| 312 | record_id = _add_record( |
| 313 | tool="t2i", |
| 314 | model=req.model, |
| 315 | input_data=input_data, |
| 316 | output_data={"images": result}, |
| 317 | files=result if isinstance(result, list) else [], |
| 318 | record_id=task_id, |
| 319 | ) |
| 320 | logger.info( |
| 321 | "Sandbox T2I completed: model=%s record_id=%s images=%d", |
| 322 | req.model, |
| 323 | record_id, |
| 324 | len(result) if isinstance(result, list) else 0, |
| 325 | ) |
| 326 | return { |
| 327 | "success": True, |
| 328 | "result": _converted_result_list(result if isinstance(result, list) else []), |
| 329 | "record_id": record_id, |
| 330 | } |
| 331 | except Exception as e: |
| 332 | logger.exception("Sandbox T2I failed: model=%s", req.model) |
| 333 | return {"success": False, "error": str(e)} |
| 334 | finally: |
| 335 | _finish_active_task(task_id) |
| 336 | |
| 337 | |
| 338 | @router.post("/api/sandbox/i2i") |
| 339 | async def sandbox_i2i(req: SandboxI2IRequest): |
| 340 | """临时工作台 - 图生图""" |
| 341 | from models.image_client import ImageClient |
| 342 | client = ImageClient() |
| 343 | input_data = {"prompt": req.prompt, "reference_image": req.image} |
| 344 | task_id = _start_active_task("i2i", req.model, input_data) |
| 345 | try: |
| 346 | logger.info("Sandbox I2I started: model=%s ratio=%s", req.model, req.ratio) |
| 347 | result = await run_in_threadpool( |
| 348 | client.generate_image, |
| 349 | req.prompt, |
| 350 | image_paths=[req.image], |
| 351 | model=req.model, |
| 352 | video_ratio=req.ratio, |
| 353 | ) |
| 354 | # 保存到历史记录 |
| 355 | record_id = _add_record( |
| 356 | tool="i2i", |
| 357 | model=req.model, |
| 358 | input_data=input_data, |
| 359 | output_data={"images": result}, |
| 360 | files=result if isinstance(result, list) else [], |
| 361 | record_id=task_id, |
| 362 | ) |
| 363 | logger.info( |
| 364 | "Sandbox I2I completed: model=%s record_id=%s images=%d", |
| 365 | req.model, |
| 366 | record_id, |
| 367 | len(result) if isinstance(result, list) else 0, |
| 368 | ) |
| 369 | return { |
| 370 | "success": True, |
| 371 | "result": _converted_result_list(result if isinstance(result, list) else []), |
| 372 | "record_id": record_id, |
| 373 | } |
| 374 | except Exception as e: |
| 375 | logger.exception("Sandbox I2I failed: model=%s", req.model) |
| 376 | return {"success": False, "error": str(e)} |
| 377 | finally: |
| 378 | _finish_active_task(task_id) |
| 379 | |
| 380 | |
| 381 | @router.post("/api/sandbox/video") |
| 382 | async def sandbox_video(req: SandboxVideoRequest): |
| 383 | """临时工作台 - 视频生成""" |
| 384 | from models.video_client import VideoClient |
| 385 | client = VideoClient() |
| 386 | duration = int(req.duration or 5) |
| 387 | input_data = { |
| 388 | "prompt": req.prompt, |
| 389 | "reference_image": req.image, |
| 390 | "ratio": req.ratio, |
| 391 | "resolution": req.resolution, |
| 392 | "duration": duration, |
| 393 | } |
| 394 | task_id = _start_active_task("video", req.model, input_data) |
| 395 | try: |
| 396 | # 生成唯一的保存路径 |
| 397 | save_dir = os.path.join(SANDBOX_DIR, "videos") |
| 398 | os.makedirs(save_dir, exist_ok=True) |
| 399 | save_path = os.path.join(save_dir, f"{uuid.uuid4().hex[:8]}.mp4") |
| 400 | logger.info( |
| 401 | "Sandbox video started: model=%s image=%s ratio=%s resolution=%s duration=%ss", |
| 402 | req.model, |
| 403 | bool(req.image), |
| 404 | req.ratio, |
| 405 | req.resolution, |
| 406 | duration, |
| 407 | ) |
| 408 | |
| 409 | result = await run_in_threadpool( |
| 410 | client.generate_video, |
| 411 | prompt=req.prompt, |
| 412 | image_path=req.image or "", |
| 413 | save_path=save_path, |
| 414 | model=req.model, |
| 415 | duration=duration, |
| 416 | shot_type="multi", |
| 417 | video_ratio=req.ratio or "16:9", |
| 418 | resolution=req.resolution or "720P", |
| 419 | ) |
| 420 | # 保存到历史记录 |
| 421 | record_id = _add_record( |
| 422 | tool="video", |
| 423 | model=req.model, |
| 424 | input_data=input_data, |
| 425 | output_data={"video": result, "video_path": save_path}, |
| 426 | files=[save_path], |
| 427 | record_id=task_id, |
| 428 | ) |
| 429 | logger.info("Sandbox video completed: model=%s record_id=%s video=%s", req.model, record_id, save_path) |
| 430 | return { |
| 431 | "success": True, |
| 432 | "result": result, |
| 433 | "video_path": _converted_video_path(save_path), |
| 434 | "record_id": record_id, |
| 435 | } |
| 436 | except Exception as e: |
| 437 | logger.exception("Sandbox video failed: model=%s", req.model) |
| 438 | return {"success": False, "error": str(e)} |
| 439 | finally: |
| 440 | _finish_active_task(task_id) |
| 441 |