117 lines
4.2 KiB
Python
117 lines
4.2 KiB
Python
"""Storage for images uploaded into the AI Assistant chat.
|
|
|
|
Images live under {user}/PromptChain/chat-uploads/, keyed by content hash —
|
|
NOT base64-embedded in the workflow JSON (that would bloat every saved
|
|
workflow/PNG). Chat history carries only the hash; the server rehydrates
|
|
pixels from disk just-in-time for the turn the image was uploaded on
|
|
(caption-once — older turns never re-send the bytes to the model).
|
|
"""
|
|
import base64
|
|
import binascii
|
|
import hashlib
|
|
import logging
|
|
from pathlib import Path
|
|
|
|
import folder_paths
|
|
|
|
logger = logging.getLogger("promptchain.chat_uploads")
|
|
|
|
# media_type → file extension for the formats vision models accept.
|
|
_EXT_BY_MEDIA = {
|
|
"image/png": "png",
|
|
"image/jpeg": "jpg",
|
|
"image/webp": "webp",
|
|
"image/gif": "gif",
|
|
}
|
|
_MEDIA_BY_EXT = {v: k for k, v in _EXT_BY_MEDIA.items()}
|
|
|
|
# Reject oversized uploads before they hit disk / a model. 20 MB decoded is
|
|
# well past any sane reference image and keeps a pasted clipboard blob from
|
|
# wedging the request.
|
|
_MAX_DECODED_BYTES = 20 * 1024 * 1024
|
|
|
|
|
|
def get_uploads_dir() -> Path:
|
|
d = Path(folder_paths.get_user_directory()) / "PromptChain" / "chat-uploads"
|
|
d.mkdir(parents=True, exist_ok=True)
|
|
return d
|
|
|
|
|
|
def _strip_data_url(b64: str) -> tuple[str, str | None]:
|
|
"""Accept either a bare base64 string or a full data: URL. Returns
|
|
(data, media_type-or-None)."""
|
|
if b64.startswith("data:"):
|
|
header, _, payload = b64.partition(",")
|
|
media = header[5:].split(";", 1)[0].strip() or None
|
|
return payload, media
|
|
return b64, None
|
|
|
|
|
|
def save_upload(b64: str, media_type: str | None = None) -> dict | None:
|
|
"""Decode a base64 image, write it under the uploads dir keyed by content
|
|
hash, and return {hash, media_type, ext, bytes}. Returns None on a
|
|
malformed / disallowed / oversized payload."""
|
|
data, url_media = _strip_data_url(b64 or "")
|
|
media_type = (media_type or url_media or "image/png").lower()
|
|
if media_type not in _EXT_BY_MEDIA:
|
|
logger.warning("rejected upload: unsupported media_type %r", media_type)
|
|
return None
|
|
try:
|
|
raw = base64.b64decode(data, validate=True)
|
|
except (binascii.Error, ValueError):
|
|
logger.warning("rejected upload: invalid base64")
|
|
return None
|
|
if not raw or len(raw) > _MAX_DECODED_BYTES:
|
|
logger.warning("rejected upload: empty or over %d bytes", _MAX_DECODED_BYTES)
|
|
return None
|
|
|
|
h = hashlib.sha256(raw).hexdigest()
|
|
ext = _EXT_BY_MEDIA[media_type]
|
|
path = get_uploads_dir() / f"{h}.{ext}"
|
|
if not path.exists(): # content-addressed: identical bytes reuse the file
|
|
path.write_bytes(raw)
|
|
return {"hash": h, "media_type": media_type, "ext": ext, "bytes": len(raw)}
|
|
|
|
|
|
def resolve_upload_path(image_hash: str) -> Path | None:
|
|
"""Path for a stored upload by hash, or None if absent. Probes known
|
|
extensions since the hash alone doesn't carry the format."""
|
|
if not image_hash:
|
|
return None
|
|
d = get_uploads_dir()
|
|
for ext in _EXT_BY_MEDIA.values():
|
|
p = d / f"{image_hash}.{ext}"
|
|
if p.is_file():
|
|
return p
|
|
return None
|
|
|
|
|
|
def load_image_data(image_hash: str) -> dict | None:
|
|
"""Load a stored upload as the {data, media_type} shape that
|
|
ai_api._call_provider_complete expects for its vision calls (distinct
|
|
from the Anthropic block shape of load_image_block)."""
|
|
path = resolve_upload_path(image_hash)
|
|
if not path:
|
|
return None
|
|
media = _MEDIA_BY_EXT.get(path.suffix.lstrip("."), "image/png")
|
|
return {
|
|
"data": base64.b64encode(path.read_bytes()).decode("ascii"),
|
|
"media_type": media,
|
|
}
|
|
|
|
|
|
def load_image_block(image_hash: str) -> dict | None:
|
|
"""Load a stored upload as an Anthropic-canonical image block:
|
|
{type:"image", source:{type:"base64", media_type, data}}
|
|
Claude consumes this shape directly; _to_openai_messages converts it to
|
|
OpenAI image_url shape for Ollama / OpenAI-compat providers."""
|
|
path = resolve_upload_path(image_hash)
|
|
if not path:
|
|
return None
|
|
media = _MEDIA_BY_EXT.get(path.suffix.lstrip("."), "image/png")
|
|
data = base64.b64encode(path.read_bytes()).decode("ascii")
|
|
return {
|
|
"type": "image",
|
|
"source": {"type": "base64", "media_type": media, "data": data},
|
|
}
|