96 lines
3.0 KiB
Python
96 lines
3.0 KiB
Python
"""What an agent may hand to a human: an allowed path, a ceiling, a method.
|
|
|
|
The ceiling and the method names belong to the frontend's API; the check
|
|
lives here so every frontend refuses the same way, in words the model reads.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import mimetypes
|
|
from dataclasses import dataclass, replace
|
|
from pathlib import Path
|
|
|
|
__all__ = ["AUTO", "DEFAULTS", "METHODS", "SendFileError", "SendFiles", "method_for"]
|
|
|
|
AUTO = "auto"
|
|
METHODS = ("document", "photo", "audio", "voice", "video", "video_note", "animation")
|
|
|
|
_BY_SUFFIX = {
|
|
".jpg": "photo",
|
|
".jpeg": "photo",
|
|
".png": "photo",
|
|
".webp": "photo",
|
|
".gif": "animation",
|
|
".mp3": "audio",
|
|
".flac": "audio",
|
|
".m4a": "audio",
|
|
".wav": "audio",
|
|
".aac": "audio",
|
|
".ogg": "voice",
|
|
".oga": "voice",
|
|
".opus": "voice",
|
|
".mp4": "video",
|
|
".m4v": "video",
|
|
".mov": "video",
|
|
}
|
|
_BY_TYPE = {"image": "photo", "audio": "audio", "video": "video"}
|
|
|
|
|
|
class SendFileError(Exception):
|
|
"""Why a file was not sent, worded for the model that asked."""
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class SendFiles:
|
|
"""Which files a frontend accepts from an agent and how big they may be."""
|
|
|
|
roots: tuple[Path, ...] = ()
|
|
"""Directories a path must resolve inside; nothing outside is sendable."""
|
|
max_bytes: int = 50 * 1024 * 1024
|
|
"""Upload ceiling of the API behind the frontend; the Bot API gives 50 MB."""
|
|
|
|
def with_roots(self, *roots: Path) -> SendFiles:
|
|
return replace(self, roots=tuple(dict.fromkeys((*self.roots, *roots))))
|
|
|
|
@property
|
|
def note(self) -> str:
|
|
"""Roots and ceiling as one line for the tool description."""
|
|
where = ", ".join(str(r) for r in self.roots) or "none, so nothing is sendable"
|
|
return f"paths under {where}, up to {mb(self.max_bytes)} MB"
|
|
|
|
def resolve(self, raw: str) -> Path:
|
|
path = Path(raw.strip()).expanduser().resolve()
|
|
if not any(path.is_relative_to(r.expanduser().resolve()) for r in self.roots):
|
|
msg = f"{path} is outside the allowed roots ({self.note}); copy it there"
|
|
raise SendFileError(msg)
|
|
if not path.is_file():
|
|
msg = f"{path} is not a file"
|
|
raise SendFileError(msg)
|
|
size = path.stat().st_size
|
|
if size > self.max_bytes:
|
|
msg = (
|
|
f"{path.name} is {mb(size)} MB, over the {mb(self.max_bytes)} MB limit"
|
|
)
|
|
raise SendFileError(msg)
|
|
return path
|
|
|
|
|
|
DEFAULTS = SendFiles()
|
|
|
|
|
|
def method_for(path: Path, method: str = AUTO) -> str:
|
|
if method and method != AUTO:
|
|
if method not in METHODS:
|
|
msg = f"unknown method {method!r}; one of: {', '.join((AUTO, *METHODS))}"
|
|
raise SendFileError(msg)
|
|
return method
|
|
guessed = _BY_SUFFIX.get(path.suffix.lower())
|
|
if guessed:
|
|
return guessed
|
|
mime, _ = mimetypes.guess_type(path.name)
|
|
return _BY_TYPE.get(mime.split("/")[0], "document") if mime else "document"
|
|
|
|
|
|
def mb(size: int) -> str:
|
|
return f"{round(size / 1024 / 1024, 2):g}"
|