e36ca308b8
Allow paste/drag temporary image assets so vision turns work without a live camera frame. Co-authored-by: Cursor <cursoragent@cursor.com>
133 lines
4.0 KiB
Python
133 lines
4.0 KiB
Python
"""Short-lived uploaded images consumed by the realtime user-input protocol."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import hashlib
|
|
import hmac
|
|
import json
|
|
import time
|
|
from dataclasses import dataclass
|
|
from io import BytesIO
|
|
from uuid import uuid4
|
|
|
|
from PIL import Image, ImageOps, UnidentifiedImageError
|
|
|
|
import settings
|
|
from services.object_storage import delete_object, get_object, put_object
|
|
|
|
|
|
INPUT_ASSET_PREFIX = "conversation-inputs/"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class StoredInputImage:
|
|
token: str
|
|
width: int
|
|
height: int
|
|
size_bytes: int
|
|
|
|
|
|
def _b64encode(data: bytes) -> str:
|
|
return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")
|
|
|
|
|
|
def _b64decode(data: str) -> bytes:
|
|
padding = "=" * (-len(data) % 4)
|
|
return base64.urlsafe_b64decode(f"{data}{padding}".encode("ascii"))
|
|
|
|
|
|
def _signature(payload: str) -> str:
|
|
digest = hmac.new(
|
|
settings.AUTH_SECRET_KEY.encode("utf-8"),
|
|
payload.encode("ascii"),
|
|
hashlib.sha256,
|
|
).digest()
|
|
return _b64encode(digest)
|
|
|
|
|
|
def _normalize_image(data: bytes) -> tuple[bytes, int, int]:
|
|
try:
|
|
source = Image.open(BytesIO(data))
|
|
width, height = source.size
|
|
if width <= 0 or height <= 0:
|
|
raise ValueError("图片尺寸无效")
|
|
if width * height > settings.INPUT_IMAGE_MAX_PIXELS:
|
|
raise ValueError("图片像素尺寸过大")
|
|
source.load()
|
|
except (UnidentifiedImageError, OSError) as exc:
|
|
raise ValueError("文件不是可识别的图片") from exc
|
|
|
|
image = ImageOps.exif_transpose(source)
|
|
image.thumbnail(
|
|
(settings.INPUT_IMAGE_MAX_EDGE, settings.INPUT_IMAGE_MAX_EDGE),
|
|
Image.Resampling.LANCZOS,
|
|
)
|
|
if image.mode in {"RGBA", "LA"} or "transparency" in image.info:
|
|
rgba = image.convert("RGBA")
|
|
background = Image.new("RGB", rgba.size, "white")
|
|
background.paste(rgba, mask=rgba.getchannel("A"))
|
|
image = background
|
|
else:
|
|
image = image.convert("RGB")
|
|
|
|
buffer = BytesIO()
|
|
image.save(buffer, format="JPEG", quality=85, optimize=True)
|
|
return buffer.getvalue(), image.width, image.height
|
|
|
|
|
|
def store_input_image(data: bytes) -> StoredInputImage:
|
|
normalized, width, height = _normalize_image(data)
|
|
key = f"{INPUT_ASSET_PREFIX}{uuid4().hex}.jpg"
|
|
put_object(key, normalized, "image/jpeg")
|
|
expires_at = int(time.time()) + settings.INPUT_IMAGE_TOKEN_TTL_SECONDS
|
|
encoded_payload = _b64encode(
|
|
json.dumps(
|
|
{"key": key, "exp": expires_at},
|
|
separators=(",", ":"),
|
|
).encode("utf-8")
|
|
)
|
|
token = f"{encoded_payload}.{_signature(encoded_payload)}"
|
|
return StoredInputImage(
|
|
token=token,
|
|
width=width,
|
|
height=height,
|
|
size_bytes=len(normalized),
|
|
)
|
|
|
|
|
|
def _key_from_token(token: str) -> str:
|
|
if not token or "." not in token:
|
|
raise ValueError("图片附件 token 无效")
|
|
encoded_payload, signature = token.rsplit(".", 1)
|
|
if not hmac.compare_digest(_signature(encoded_payload), signature):
|
|
raise ValueError("图片附件 token 签名无效")
|
|
try:
|
|
payload = json.loads(_b64decode(encoded_payload))
|
|
key = str(payload.get("key") or "")
|
|
expires_at = int(payload.get("exp") or 0)
|
|
except (ValueError, TypeError, json.JSONDecodeError) as exc:
|
|
raise ValueError("图片附件 token 内容无效") from exc
|
|
if expires_at < int(time.time()):
|
|
raise ValueError("图片附件已过期，请重新添加")
|
|
if not key.startswith(INPUT_ASSET_PREFIX):
|
|
raise ValueError("图片附件存储位置无效")
|
|
return key
|
|
|
|
|
|
def consume_input_image(token: str) -> bytes:
|
|
"""Read a signed image once and remove its temporary object."""
|
|
|
|
key = _key_from_token(token)
|
|
data = get_object(key)
|
|
try:
|
|
delete_object(key)
|
|
except Exception:
|
|
# Consumption succeeded. A stale temporary object must not break the turn.
|
|
pass
|
|
return data
|
|
|
|
|
|
def discard_input_image(token: str) -> None:
|
|
delete_object(_key_from_token(token))
|