347 lines
11 KiB
Python
347 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import base64
|
|
import math
|
|
import re
|
|
from io import BytesIO
|
|
from typing import Any
|
|
|
|
from PIL import Image
|
|
|
|
DEFAULT_IMAGE_SIZE = (1024, 1024)
|
|
IMAGE_INPUT_TOKEN_MODEL = "gpt-5.4-mini"
|
|
|
|
PATCH_SIZE = 32
|
|
TILE_SIZE = 512
|
|
TILE_HIGH_SHORT_SIDE = 768
|
|
|
|
PATCH_1536_MODELS = (
|
|
"gpt-5.4-mini",
|
|
"gpt-5.4-nano",
|
|
"gpt-5-mini",
|
|
"gpt-5-nano",
|
|
"gpt-5.2",
|
|
"gpt-5.3-codex",
|
|
"gpt-5-codex-mini",
|
|
"gpt-5.1-codex-mini",
|
|
"gpt-5.2-codex",
|
|
"gpt-5.2-chat-latest",
|
|
"o4-mini",
|
|
"gpt-4.1-mini",
|
|
"gpt-4.1-nano",
|
|
)
|
|
|
|
PATCH_MULTIPLIERS = {
|
|
"gpt-5.4-mini": 1.62,
|
|
"gpt-5.4-nano": 2.46,
|
|
"gpt-5-mini": 1.62,
|
|
"gpt-5-nano": 2.46,
|
|
"gpt-4.1-mini": 1.62,
|
|
"gpt-4.1-nano": 2.46,
|
|
"o4-mini": 1.72,
|
|
}
|
|
|
|
|
|
def _model_name(model: str) -> str:
|
|
return str(model or "").strip().lower()
|
|
|
|
|
|
def image_size_from_bytes(data: bytes) -> tuple[int, int] | None:
|
|
if not data:
|
|
return None
|
|
try:
|
|
with Image.open(BytesIO(data)) as image:
|
|
width, height = image.size
|
|
except Exception:
|
|
return None
|
|
if width <= 0 or height <= 0:
|
|
return None
|
|
return int(width), int(height)
|
|
|
|
|
|
def _decode_data_url(value: str) -> bytes:
|
|
text = str(value or "").strip()
|
|
payload = text.split(",", 1)[1] if text.startswith("data:") and "," in text else text
|
|
return base64.b64decode(payload)
|
|
|
|
|
|
def image_size_from_data_url(value: str) -> tuple[int, int] | None:
|
|
try:
|
|
return image_size_from_bytes(_decode_data_url(value))
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def parse_image_size(size: object, default: tuple[int, int] = DEFAULT_IMAGE_SIZE) -> tuple[int, int]:
|
|
if isinstance(size, (tuple, list)) and len(size) >= 2:
|
|
try:
|
|
width = int(size[0])
|
|
height = int(size[1])
|
|
if width > 0 and height > 0:
|
|
return width, height
|
|
except (TypeError, ValueError):
|
|
pass
|
|
match = re.search(r"(\d{2,5})\D+(\d{2,5})", str(size or ""))
|
|
if not match:
|
|
return default
|
|
width, height = int(match.group(1)), int(match.group(2))
|
|
return (width, height) if width > 0 and height > 0 else default
|
|
|
|
|
|
def _patch_count(width: float, height: float) -> int:
|
|
return math.ceil(width / PATCH_SIZE) * math.ceil(height / PATCH_SIZE)
|
|
|
|
|
|
def _patch_multiplier(model: str) -> float:
|
|
name = _model_name(model)
|
|
for prefix, multiplier in PATCH_MULTIPLIERS.items():
|
|
if name.startswith(prefix):
|
|
return multiplier
|
|
return 1.0
|
|
|
|
|
|
def _patch_limits(model: str, detail: str) -> tuple[int, int] | None:
|
|
name = _model_name(model)
|
|
if any(name.startswith(prefix) for prefix in PATCH_1536_MODELS):
|
|
return 1536, 2048
|
|
if name.startswith("gpt-5.5"):
|
|
return (10000, 6000) if detail in {"auto", "original"} else (2500, 2048)
|
|
if name.startswith("gpt-5.4"):
|
|
return (10000, 6000) if detail == "original" else (2500, 2048)
|
|
return None
|
|
|
|
|
|
def _patch_tokens(width: int, height: int, model: str, detail: str) -> int:
|
|
multiplier = _patch_multiplier(model)
|
|
if detail == "low":
|
|
return math.ceil(256 * multiplier)
|
|
|
|
limits = _patch_limits(model, detail)
|
|
if limits is None:
|
|
return 0
|
|
patch_budget, max_dimension = limits
|
|
scale = min(1.0, max_dimension / max(width, height))
|
|
resized_width = width * scale
|
|
resized_height = height * scale
|
|
|
|
if _patch_count(resized_width, resized_height) > patch_budget:
|
|
shrink_factor = math.sqrt((PATCH_SIZE * PATCH_SIZE * patch_budget) / (resized_width * resized_height))
|
|
width_units = resized_width * shrink_factor / PATCH_SIZE
|
|
height_units = resized_height * shrink_factor / PATCH_SIZE
|
|
adjusted_shrink_factor = shrink_factor * min(
|
|
math.floor(width_units) / width_units if width_units else 1,
|
|
math.floor(height_units) / height_units if height_units else 1,
|
|
)
|
|
resized_width *= adjusted_shrink_factor
|
|
resized_height *= adjusted_shrink_factor
|
|
|
|
tokens = min(_patch_count(max(1, resized_width), max(1, resized_height)), patch_budget)
|
|
return math.ceil(tokens * multiplier)
|
|
|
|
|
|
def _tile_rates(model: str) -> tuple[int, int]:
|
|
name = _model_name(model)
|
|
if name in {"gpt-5", "gpt-5-chat-latest"}:
|
|
return 70, 140
|
|
if name.startswith("gpt-4o-mini"):
|
|
return 2833, 5667
|
|
if name.startswith(("o1", "o1-pro", "o3")):
|
|
return 75, 150
|
|
if name.startswith("computer-use-preview"):
|
|
return 65, 129
|
|
return 85, 170
|
|
|
|
|
|
def _tile_tokens(width: int, height: int, model: str, detail: str) -> int:
|
|
base_tokens, tile_tokens = _tile_rates(model)
|
|
if detail == "low":
|
|
return base_tokens
|
|
|
|
scale = min(1.0, 2048 / max(width, height))
|
|
resized_width = width * scale
|
|
resized_height = height * scale
|
|
short_side = min(resized_width, resized_height)
|
|
if short_side > 0:
|
|
scale = TILE_HIGH_SHORT_SIDE / short_side
|
|
resized_width *= scale
|
|
resized_height *= scale
|
|
|
|
tiles = math.ceil(resized_width / TILE_SIZE) * math.ceil(resized_height / TILE_SIZE)
|
|
return base_tokens + tiles * tile_tokens
|
|
|
|
|
|
def count_image_input_tokens(
|
|
width: int,
|
|
height: int,
|
|
model: str,
|
|
detail: str = "auto",
|
|
input_fidelity: str = "low",
|
|
) -> int:
|
|
if width <= 0 or height <= 0:
|
|
return 0
|
|
detail = str(detail or "auto").strip().lower() or "auto"
|
|
return _patch_tokens(width, height, IMAGE_INPUT_TOKEN_MODEL, detail)
|
|
|
|
|
|
def _part_size(part: dict[str, Any]) -> tuple[int, int] | None:
|
|
try:
|
|
width = int(part.get("width") or 0)
|
|
height = int(part.get("height") or 0)
|
|
except (TypeError, ValueError):
|
|
width = height = 0
|
|
if width > 0 and height > 0:
|
|
return width, height
|
|
|
|
data = part.get("data")
|
|
if isinstance(data, (bytes, bytearray)):
|
|
return image_size_from_bytes(bytes(data))
|
|
|
|
image_url = part.get("image_url")
|
|
if isinstance(image_url, dict):
|
|
image_url = image_url.get("url") or image_url.get("image_url")
|
|
if isinstance(image_url, str) and image_url.startswith("data:"):
|
|
return image_size_from_data_url(image_url)
|
|
|
|
source = part.get("source")
|
|
if isinstance(source, dict) and str(source.get("type") or "") == "base64":
|
|
try:
|
|
return image_size_from_bytes(base64.b64decode(str(source.get("data") or "")))
|
|
except Exception:
|
|
return None
|
|
return None
|
|
|
|
|
|
def count_image_content_tokens(content: object, model: str, default_detail: str = "auto") -> int:
|
|
if not isinstance(content, list):
|
|
return 0
|
|
total = 0
|
|
for part in content:
|
|
if not isinstance(part, dict):
|
|
continue
|
|
part_type = str(part.get("type") or "").strip()
|
|
if part_type not in {"image", "image_url", "input_image"} and not part.get("source"):
|
|
continue
|
|
size = _part_size(part)
|
|
if not size:
|
|
continue
|
|
total += count_image_input_tokens(
|
|
size[0],
|
|
size[1],
|
|
model,
|
|
str(part.get("detail") or default_detail or "auto"),
|
|
str(part.get("input_fidelity") or part.get("inputFidelity") or "low"),
|
|
)
|
|
return total
|
|
|
|
|
|
def count_image_inputs_tokens(images: object, model: str, default_detail: str = "auto") -> int:
|
|
if not images:
|
|
return 0
|
|
total = 0
|
|
entries = images if isinstance(images, list) else [images]
|
|
for image in entries:
|
|
data = image[0] if isinstance(image, tuple) and image else image
|
|
if not isinstance(data, (bytes, bytearray)):
|
|
continue
|
|
size = image_size_from_bytes(bytes(data))
|
|
if size:
|
|
total += count_image_input_tokens(size[0], size[1], model, default_detail)
|
|
return total
|
|
|
|
|
|
def count_generated_image_tokens(width: int, height: int, quality: str = "auto") -> int:
|
|
patches = _patch_count(width, height)
|
|
quality = str(quality or "auto").strip().lower()
|
|
if quality == "low":
|
|
return math.ceil(patches * 17 / 64)
|
|
if quality in {"high", "hd"}:
|
|
return math.ceil(patches * 65 / 16)
|
|
return math.ceil(patches * 33 / 32)
|
|
|
|
|
|
def count_image_output_tokens(size: object = None, quality: str = "auto", count: int = 1) -> int:
|
|
width, height = parse_image_size(size)
|
|
return max(0, int(count or 0)) * count_generated_image_tokens(width, height, quality)
|
|
|
|
|
|
def count_image_output_items_tokens(
|
|
items: object,
|
|
size: object = None,
|
|
quality: str = "auto",
|
|
) -> int:
|
|
if not isinstance(items, list) or not items:
|
|
return 0
|
|
fallback_size = parse_image_size(size)
|
|
total = 0
|
|
for item in items:
|
|
image_size = None
|
|
if isinstance(item, dict):
|
|
b64_json = str(item.get("b64_json") or "").strip()
|
|
if b64_json:
|
|
try:
|
|
image_size = image_size_from_bytes(base64.b64decode(b64_json))
|
|
except Exception:
|
|
image_size = None
|
|
width, height = image_size or fallback_size
|
|
total += count_generated_image_tokens(width, height, quality)
|
|
return total
|
|
|
|
|
|
def token_usage(
|
|
input_text_tokens: int = 0,
|
|
input_image_tokens: int = 0,
|
|
output_text_tokens: int = 0,
|
|
output_image_tokens: int = 0,
|
|
) -> dict[str, Any]:
|
|
input_tokens = max(0, int(input_text_tokens or 0)) + max(0, int(input_image_tokens or 0))
|
|
output_tokens = max(0, int(output_text_tokens or 0)) + max(0, int(output_image_tokens or 0))
|
|
return {
|
|
"input_tokens": input_tokens,
|
|
"output_tokens": output_tokens,
|
|
"total_tokens": input_tokens + output_tokens,
|
|
"input_tokens_details": {
|
|
"text_tokens": max(0, int(input_text_tokens or 0)),
|
|
"image_tokens": max(0, int(input_image_tokens or 0)),
|
|
"cached_tokens": 0,
|
|
},
|
|
"output_tokens_details": {
|
|
"text_tokens": max(0, int(output_text_tokens or 0)),
|
|
"image_tokens": max(0, int(output_image_tokens or 0)),
|
|
"reasoning_tokens": 0,
|
|
},
|
|
}
|
|
|
|
|
|
def image_usage(
|
|
input_text_tokens: int = 0,
|
|
input_image_tokens: int = 0,
|
|
output_tokens: int = 0,
|
|
) -> dict[str, Any]:
|
|
return token_usage(
|
|
input_text_tokens=input_text_tokens,
|
|
input_image_tokens=input_image_tokens,
|
|
output_image_tokens=output_tokens,
|
|
)
|
|
|
|
|
|
def chat_usage_from_image_usage(usage: dict[str, Any]) -> dict[str, Any]:
|
|
input_tokens = int(usage.get("input_tokens") or 0)
|
|
output_tokens = int(usage.get("output_tokens") or 0)
|
|
input_details = usage.get("input_tokens_details") if isinstance(usage.get("input_tokens_details"), dict) else {}
|
|
output_details = usage.get("output_tokens_details") if isinstance(usage.get("output_tokens_details"), dict) else {}
|
|
return {
|
|
"prompt_tokens": input_tokens,
|
|
"completion_tokens": output_tokens,
|
|
"total_tokens": input_tokens + output_tokens,
|
|
"prompt_tokens_details": {
|
|
"text_tokens": int(input_details.get("text_tokens") or 0),
|
|
"image_tokens": int(input_details.get("image_tokens") or 0),
|
|
"cached_tokens": int(input_details.get("cached_tokens") or 0),
|
|
},
|
|
"completion_tokens_details": {
|
|
"text_tokens": int(output_details.get("text_tokens") or 0),
|
|
"image_tokens": int(output_details.get("image_tokens") or 0),
|
|
"reasoning_tokens": int(output_details.get("reasoning_tokens") or 0),
|
|
},
|
|
}
|