3f8487f761
Co-authored-by: Alexis Tacnet <alexis@mistral.ai> Co-authored-by: Clément Drouin <clement.drouin@mistral.ai> Co-authored-by: Guillaume LE GOFF <guillaume.lgf@gmail.com> Co-authored-by: Lucas Marandat <31749711+lucasmrdt@users.noreply.github.com> Co-authored-by: Maxime Dolores <maxime.dolores@ext.mistral.ai> Co-authored-by: Pierre Rossinès <pierre.rossines@mistral.ai> Co-authored-by: Quentin <quentin.torroba@mistral.ai> Co-authored-by: Val <102326092+vdeva@users.noreply.github.com> Co-authored-by: Vincent G <10739306+VinceOPS@users.noreply.github.com> Co-authored-by: p.vezia <166131032+le-codeur-rapide@users.noreply.github.com> Co-authored-by: Hiba Chaabnia <Hiba-Chaabnia@users.noreply.github.com> Co-authored-by: Nikhil Bhima <nikhilbhima@users.noreply.github.com> Co-authored-by: Nkipohcs <Nkipohcs@users.noreply.github.com> Co-authored-by: Mistral Vibe <vibe@mistral.ai>
56 lines
1.7 KiB
Python
56 lines
1.7 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import mimetypes
|
|
from pathlib import Path
|
|
|
|
from vibe.core.types import IMAGE_EXTENSIONS, ImageAttachment
|
|
|
|
_DEFAULT_MIME_BY_EXT = {
|
|
".png": "image/png",
|
|
".jpg": "image/jpeg",
|
|
".jpeg": "image/jpeg",
|
|
".gif": "image/gif",
|
|
".webp": "image/webp",
|
|
}
|
|
|
|
|
|
class ImageSnapshotError(Exception):
|
|
pass
|
|
|
|
|
|
def snapshot_image(
|
|
source: Path, *, alias: str, session_dir: Path | None
|
|
) -> ImageAttachment:
|
|
source_abs = source.expanduser().resolve()
|
|
if not source_abs.is_file():
|
|
raise ImageSnapshotError(f"Not a file: {source_abs}")
|
|
|
|
ext = source_abs.suffix.lower()
|
|
if ext not in IMAGE_EXTENSIONS:
|
|
raise ImageSnapshotError(f"Unsupported image extension: {ext}")
|
|
|
|
mime_type = _DEFAULT_MIME_BY_EXT.get(ext) or (
|
|
mimetypes.guess_type(source_abs.name)[0] or "application/octet-stream"
|
|
)
|
|
|
|
try:
|
|
data = source_abs.read_bytes()
|
|
except OSError as e:
|
|
raise ImageSnapshotError(f"Failed to read image {source_abs}: {e}") from e
|
|
|
|
# Session logging disabled: no snapshot copy is made; the attachment
|
|
# points to the original source. Resume is not possible in this mode so
|
|
# snapshot stability is not required.
|
|
if session_dir is None:
|
|
return ImageAttachment(path=source_abs, alias=alias, mime_type=mime_type)
|
|
|
|
attachments_dir = session_dir / "attachments"
|
|
attachments_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
digest = hashlib.sha1(data, usedforsecurity=False).hexdigest()
|
|
dest = attachments_dir / f"{digest}{ext}"
|
|
if not dest.exists():
|
|
dest.write_bytes(data)
|
|
|
|
return ImageAttachment(path=dest.resolve(), alias=alias, mime_type=mime_type)
|