1128 lines
40 KiB
Python
1128 lines
40 KiB
Python
from __future__ import annotations
|
|
import base64
|
|
import re
|
|
from collections.abc import Mapping, Hashable
|
|
from functools import lru_cache
|
|
from typing import (
|
|
Any,
|
|
Callable,
|
|
Literal,
|
|
Optional,
|
|
Union,
|
|
TypedDict,
|
|
TypeVar,
|
|
cast,
|
|
)
|
|
from pathlib import Path
|
|
from urllib.parse import urlparse
|
|
import mimetypes
|
|
import requests
|
|
from pydantic import BaseModel, Field
|
|
|
|
from ..core.exceptions import MultimodalError
|
|
from ..mode import Mode
|
|
|
|
F = TypeVar("F", bound=Callable[..., Any])
|
|
K = TypeVar("K", bound=Hashable)
|
|
V = TypeVar("V")
|
|
|
|
# OpenAI source: https://platform.openai.com/docs/guides/vision/what-type-of-files-can-i-upload
|
|
# Anthropic source: https://docs.anthropic.com/en/docs/build-with-claude/vision#ensuring-image-quality
|
|
VALID_MIME_TYPES = ["image/jpeg", "image/png", "image/gif", "image/webp"]
|
|
VALID_AUDIO_MIME_TYPES = [
|
|
"audio/aac",
|
|
"audio/flac",
|
|
"audio/mp3",
|
|
"audio/m4a",
|
|
"audio/mpeg",
|
|
"audio/mpga",
|
|
"audio/mp4",
|
|
"audio/opus",
|
|
"audio/pcm",
|
|
"audio/wav",
|
|
"audio/webm",
|
|
]
|
|
VALID_PDF_MIME_TYPES = ["application/pdf"]
|
|
CacheControlType = Mapping[str, str]
|
|
OptionalCacheControlType = Optional[CacheControlType]
|
|
|
|
|
|
class ImageParamsBase(TypedDict):
|
|
type: Literal["image"]
|
|
source: str
|
|
|
|
|
|
class ImageParams(ImageParamsBase, total=False):
|
|
cache_control: CacheControlType
|
|
|
|
|
|
class Image(BaseModel):
|
|
source: Union[str, Path] = Field( # noqa: UP007
|
|
description="URL, file path, or base64 data of the image"
|
|
)
|
|
media_type: str = Field(description="MIME type of the image")
|
|
data: Union[str, None] = Field( # noqa: UP007
|
|
None, description="Base64 encoded image data", repr=False
|
|
)
|
|
|
|
@classmethod
|
|
def autodetect(cls, source: str | Path) -> Image:
|
|
"""Attempt to autodetect an image from a source string or Path."""
|
|
if isinstance(source, str):
|
|
if cls.is_base64(source):
|
|
return cls.from_base64(source)
|
|
if source.startswith(("http://", "https://")):
|
|
return cls.from_url(source)
|
|
if source.startswith("gs://"):
|
|
return cls.from_gs_url(source)
|
|
# Since detecting the max length of a file universally cross-platform is difficult,
|
|
# we'll just try/catch the Path conversion and file check
|
|
try:
|
|
path = Path(source)
|
|
if path.is_file():
|
|
return cls.from_path(path)
|
|
except OSError:
|
|
pass # Fall through to raw base64 attempt
|
|
|
|
return cls.from_raw_base64(source)
|
|
|
|
if isinstance(source, Path):
|
|
return cls.from_path(source)
|
|
|
|
@classmethod
|
|
def autodetect_safely(cls, source: Union[str, Path]) -> Union[Image, str]: # noqa: UP007
|
|
"""Safely attempt to autodetect an image from a source string or path.
|
|
|
|
Args:
|
|
source (Union[str,path]): The source string or path.
|
|
Returns:
|
|
An Image if the source is detected to be a valid image, otherwise
|
|
the source itself as a string.
|
|
"""
|
|
try:
|
|
return cls.autodetect(source)
|
|
except ValueError:
|
|
return str(source)
|
|
|
|
@classmethod
|
|
def is_base64(cls, s: str) -> bool:
|
|
return bool(re.match(r"^data:image/[a-zA-Z]+;base64,", s))
|
|
|
|
@classmethod # Caching likely unnecessary
|
|
def from_base64(cls, data_uri: str) -> Image:
|
|
header, encoded = data_uri.split(",", 1)
|
|
media_type = header.split(":")[1].split(";")[0]
|
|
if media_type not in VALID_MIME_TYPES:
|
|
raise MultimodalError(
|
|
f"Unsupported image format: {media_type}. Supported formats: {', '.join(VALID_MIME_TYPES)}",
|
|
content_type="image",
|
|
)
|
|
return cls(
|
|
source=data_uri,
|
|
media_type=media_type,
|
|
data=encoded,
|
|
)
|
|
|
|
@classmethod
|
|
def from_gs_url(cls, data_uri: str, timeout: int = 30) -> Image:
|
|
"""
|
|
Create an Image instance from a Google Cloud Storage URL.
|
|
|
|
Args:
|
|
data_uri: GCS URL starting with gs://
|
|
timeout: Request timeout in seconds (default: 30)
|
|
"""
|
|
if not data_uri.startswith("gs://"):
|
|
raise ValueError("URL must start with gs://")
|
|
|
|
public_url = f"https://storage.googleapis.com/{data_uri[5:]}"
|
|
|
|
try:
|
|
response = requests.get(public_url, timeout=timeout)
|
|
response.raise_for_status()
|
|
media_type = response.headers.get("Content-Type")
|
|
if media_type not in VALID_MIME_TYPES:
|
|
raise ValueError(f"Unsupported image format: {media_type}")
|
|
|
|
data = base64.b64encode(response.content).decode("utf-8")
|
|
|
|
return cls(source=data_uri, media_type=media_type, data=data)
|
|
except requests.RequestException as e:
|
|
raise ValueError(
|
|
"Failed to access GCS image (must be publicly readable)"
|
|
) from e
|
|
|
|
@classmethod # Caching likely unnecessary
|
|
def from_raw_base64(cls, data: str) -> Image:
|
|
try:
|
|
decoded = base64.b64decode(data)
|
|
|
|
# Detect image type from file signature (magic bytes)
|
|
# This replaces imghdr which was removed in Python 3.13
|
|
img_type = None
|
|
if decoded.startswith(b"\xff\xd8\xff"):
|
|
img_type = "jpeg"
|
|
elif decoded.startswith(b"\x89PNG\r\n\x1a\n"):
|
|
img_type = "png"
|
|
elif decoded.startswith(b"GIF87a") or decoded.startswith(b"GIF89a"):
|
|
img_type = "gif"
|
|
elif decoded.startswith(b"RIFF") and decoded[8:12] == b"WEBP":
|
|
img_type = "webp"
|
|
|
|
if img_type:
|
|
media_type = f"image/{img_type}"
|
|
if media_type in VALID_MIME_TYPES:
|
|
return cls(
|
|
source=data,
|
|
media_type=media_type,
|
|
data=data,
|
|
)
|
|
raise ValueError(f"Unsupported image type: {img_type}")
|
|
except Exception as e:
|
|
raise ValueError(f"Invalid or unsupported base64 image data") from e
|
|
|
|
@classmethod
|
|
@lru_cache
|
|
def from_url(cls, url: str) -> Image:
|
|
if url.startswith("gs://"):
|
|
return cls.from_gs_url(url)
|
|
if cls.is_base64(url):
|
|
return cls.from_base64(url)
|
|
|
|
parsed_url = urlparse(url)
|
|
media_type, _ = mimetypes.guess_type(parsed_url.path)
|
|
|
|
if not media_type:
|
|
try:
|
|
response = requests.head(url, allow_redirects=True)
|
|
media_type = response.headers.get("Content-Type")
|
|
except requests.RequestException as e:
|
|
raise ValueError(f"Failed to fetch image from URL") from e
|
|
|
|
if media_type not in VALID_MIME_TYPES:
|
|
raise ValueError(f"Unsupported image format: {media_type}")
|
|
return cls(source=url, media_type=media_type, data=None)
|
|
|
|
@classmethod
|
|
@lru_cache
|
|
def from_path(cls, path: Union[str, Path]) -> Image: # noqa: UP007
|
|
path = Path(path)
|
|
if not path.is_file():
|
|
raise FileNotFoundError(f"Image file not found: {path}")
|
|
|
|
if path.stat().st_size == 0:
|
|
raise ValueError("Image file is empty")
|
|
|
|
media_type, _ = mimetypes.guess_type(str(path))
|
|
if media_type not in VALID_MIME_TYPES:
|
|
raise ValueError(f"Unsupported image format: {media_type}")
|
|
|
|
data = base64.b64encode(path.read_bytes()).decode("utf-8")
|
|
return cls(source=path, media_type=media_type, data=data)
|
|
|
|
@staticmethod
|
|
@lru_cache
|
|
def url_to_base64(url: str) -> str:
|
|
"""Cachable helper method for getting image url and encoding to base64."""
|
|
response = requests.get(url)
|
|
response.raise_for_status()
|
|
data = base64.b64encode(response.content).decode("utf-8")
|
|
return data
|
|
|
|
def to_anthropic(self) -> dict[str, Any]:
|
|
if (
|
|
isinstance(self.source, str)
|
|
and self.source.startswith(("http://", "https://"))
|
|
and not self.data
|
|
):
|
|
self.data = self.url_to_base64(self.source)
|
|
|
|
return {
|
|
"type": "image",
|
|
"source": {
|
|
"type": "base64",
|
|
"media_type": self.media_type,
|
|
"data": self.data,
|
|
},
|
|
}
|
|
|
|
def to_openai(self, mode: Mode) -> dict[str, Any]:
|
|
image_type = (
|
|
"input_image"
|
|
if mode in {Mode.RESPONSES_TOOLS, Mode.RESPONSES_TOOLS_WITH_INBUILT_TOOLS}
|
|
else "image_url"
|
|
)
|
|
if (
|
|
isinstance(self.source, str)
|
|
and self.source.startswith(("http://", "https://"))
|
|
and not self.is_base64(self.source)
|
|
):
|
|
if mode in {Mode.RESPONSES_TOOLS, Mode.RESPONSES_TOOLS_WITH_INBUILT_TOOLS}:
|
|
return {"type": "input_image", "image_url": self.source}
|
|
else:
|
|
return {"type": image_type, "image_url": {"url": self.source}}
|
|
elif self.data or self.is_base64(str(self.source)):
|
|
data = self.data or str(self.source).split(",", 1)[1]
|
|
if mode in {Mode.RESPONSES_TOOLS, Mode.RESPONSES_TOOLS_WITH_INBUILT_TOOLS}:
|
|
return {
|
|
"type": "input_image",
|
|
"image_url": f"data:{self.media_type};base64,{data}",
|
|
}
|
|
else:
|
|
return {
|
|
"type": image_type,
|
|
"image_url": {"url": f"data:{self.media_type};base64,{data}"},
|
|
}
|
|
else:
|
|
raise ValueError("Image data is missing for base64 encoding.")
|
|
|
|
def to_genai(self):
|
|
"""
|
|
Convert the Image instance to Google GenAI's API format.
|
|
"""
|
|
try:
|
|
from google.genai import types
|
|
except ImportError as err:
|
|
raise ImportError(
|
|
"google-genai package is required for GenAI integration. Install with: pip install google-genai"
|
|
) from err
|
|
|
|
# Google Cloud Storage
|
|
if isinstance(self.source, str) and self.source.startswith("gs://"):
|
|
return types.Part.from_bytes(
|
|
data=self.data, # type: ignore
|
|
mime_type=self.media_type,
|
|
)
|
|
|
|
# URL
|
|
if isinstance(self.source, str) and self.source.startswith(
|
|
("http://", "https://")
|
|
):
|
|
return types.Part.from_bytes(
|
|
data=requests.get(self.source).content,
|
|
mime_type=self.media_type,
|
|
)
|
|
|
|
if self.data or self.is_base64(str(self.source)):
|
|
data = self.data or str(self.source).split(",", 1)[1]
|
|
return types.Part.from_bytes(
|
|
data=base64.b64decode(data), mime_type=self.media_type
|
|
) # type: ignore
|
|
|
|
else:
|
|
raise ValueError("Image data is missing for base64 encoding.")
|
|
|
|
|
|
class Audio(BaseModel):
|
|
"""Represents an audio that can be loaded from a URL or file path."""
|
|
|
|
source: Union[str, Path] = Field(description="URL or file path of the audio") # noqa: UP007
|
|
data: Union[str, None] = Field( # noqa: UP007
|
|
None, description="Base64 encoded audio data", repr=False
|
|
)
|
|
media_type: str = Field(description="MIME type of the audio")
|
|
|
|
@classmethod
|
|
def autodetect(cls, source: str | Path) -> Audio:
|
|
"""Attempt to autodetect an audio from a source string or Path."""
|
|
if isinstance(source, str):
|
|
if cls.is_base64(source):
|
|
return cls.from_base64(source)
|
|
if source.startswith(("http://", "https://")):
|
|
return cls.from_url(source)
|
|
if source.startswith("gs://"):
|
|
return cls.from_gs_url(source)
|
|
# Since detecting the max length of a file universally cross-platform is difficult,
|
|
# we'll just try/catch the Path conversion and file check
|
|
try:
|
|
path = Path(source)
|
|
if path.is_file():
|
|
return cls.from_path(path)
|
|
except OSError:
|
|
pass # Fall through to error
|
|
|
|
raise ValueError("Unable to determine audio source")
|
|
|
|
if isinstance(source, Path):
|
|
return cls.from_path(source)
|
|
|
|
@classmethod
|
|
def autodetect_safely(cls, source: Union[str, Path]) -> Union[Audio, str]: # noqa: UP007
|
|
"""Safely attempt to autodetect an audio from a source string or path.
|
|
|
|
Args:
|
|
source (Union[str,path]): The source string or path.
|
|
Returns:
|
|
An Audio if the source is detected to be a valid audio, otherwise
|
|
the source itself as a string.
|
|
"""
|
|
try:
|
|
return cls.autodetect(source)
|
|
except ValueError:
|
|
return str(source)
|
|
|
|
@classmethod
|
|
def is_base64(cls, s: str) -> bool:
|
|
return bool(re.match(r"^data:audio/[a-zA-Z0-9+-]+;base64,", s))
|
|
|
|
@classmethod
|
|
def from_base64(cls, data_uri: str) -> Audio:
|
|
header, encoded = data_uri.split(",", 1)
|
|
media_type = header.split(":")[1].split(";")[0]
|
|
if media_type not in VALID_AUDIO_MIME_TYPES:
|
|
raise ValueError(f"Unsupported audio format: {media_type}")
|
|
return cls(
|
|
source=data_uri,
|
|
media_type=media_type,
|
|
data=encoded,
|
|
)
|
|
|
|
@classmethod
|
|
def from_url(cls, url: str) -> Audio:
|
|
"""Create an Audio instance from a URL."""
|
|
if url.startswith("gs://"):
|
|
return cls.from_gs_url(url)
|
|
response = requests.get(url)
|
|
content_type = response.headers.get("content-type")
|
|
assert content_type in VALID_AUDIO_MIME_TYPES, (
|
|
f"Invalid audio format. Must be one of: {', '.join(VALID_AUDIO_MIME_TYPES)}"
|
|
)
|
|
|
|
data = base64.b64encode(response.content).decode("utf-8")
|
|
return cls(source=url, data=data, media_type=content_type)
|
|
|
|
@classmethod
|
|
def from_path(cls, path: Union[str, Path]) -> Audio: # noqa: UP007
|
|
"""Create an Audio instance from a file path."""
|
|
path = Path(path)
|
|
assert path.is_file(), f"Audio file not found: {path}"
|
|
|
|
mime_type = mimetypes.guess_type(str(path))[0]
|
|
|
|
if mime_type == "audio/x-wav":
|
|
mime_type = "audio/wav"
|
|
|
|
if (
|
|
mime_type == "audio/vnd.dlna.adts"
|
|
): # <--- this is the case for aac audio files in Windows
|
|
mime_type = "audio/aac"
|
|
|
|
assert mime_type in VALID_AUDIO_MIME_TYPES, (
|
|
f"Invalid audio format. Must be one of: {', '.join(VALID_AUDIO_MIME_TYPES)}"
|
|
)
|
|
|
|
data = base64.b64encode(path.read_bytes()).decode("utf-8")
|
|
return cls(source=str(path), data=data, media_type=mime_type)
|
|
|
|
@classmethod
|
|
def from_gs_url(cls, data_uri: str, timeout: int = 30) -> Audio:
|
|
"""
|
|
Create an Audio instance from a Google Cloud Storage URL.
|
|
|
|
Args:
|
|
data_uri: GCS URL starting with gs://
|
|
timeout: Request timeout in seconds (default: 30)
|
|
"""
|
|
if not data_uri.startswith("gs://"):
|
|
raise ValueError("URL must start with gs://")
|
|
|
|
public_url = f"https://storage.googleapis.com/{data_uri[5:]}"
|
|
|
|
try:
|
|
response = requests.get(public_url, timeout=timeout)
|
|
response.raise_for_status()
|
|
media_type = response.headers.get("Content-Type")
|
|
if media_type not in VALID_AUDIO_MIME_TYPES:
|
|
raise ValueError(f"Unsupported audio format: {media_type}")
|
|
|
|
data = base64.b64encode(response.content).decode("utf-8")
|
|
|
|
return cls(source=data_uri, media_type=media_type, data=data)
|
|
except requests.RequestException as e:
|
|
raise ValueError(
|
|
"Failed to access GCS audio (must be publicly readable)"
|
|
) from e
|
|
|
|
def to_openai(self, mode: Mode) -> dict[str, Any]:
|
|
"""Convert the Audio instance to OpenAI's API format."""
|
|
if mode in {Mode.RESPONSES_TOOLS, Mode.RESPONSES_TOOLS_WITH_INBUILT_TOOLS}:
|
|
raise ValueError("OpenAI Responses doesn't support audio")
|
|
|
|
return {
|
|
"type": "input_audio",
|
|
"input_audio": {"data": self.data, "format": "wav"},
|
|
}
|
|
|
|
def to_anthropic(self) -> dict[str, Any]:
|
|
raise NotImplementedError("Anthropic is not supported yet")
|
|
|
|
def to_genai(self):
|
|
"""
|
|
Convert the Audio instance to Google GenAI's API format.
|
|
"""
|
|
try:
|
|
from google.genai import types
|
|
except ImportError as err:
|
|
raise ImportError(
|
|
"google-genai package is required for GenAI integration. Install with: pip install google-genai"
|
|
) from err
|
|
|
|
return types.Part.from_bytes(
|
|
data=base64.b64decode(self.data), # type: ignore
|
|
mime_type=self.media_type,
|
|
)
|
|
|
|
|
|
class ImageWithCacheControl(Image):
|
|
"""Image with Anthropic prompt caching support."""
|
|
|
|
cache_control: OptionalCacheControlType = Field(
|
|
None, description="Optional Anthropic cache control image"
|
|
)
|
|
|
|
@classmethod
|
|
def from_image_params(cls, image_params: ImageParams) -> Image:
|
|
source = image_params["source"]
|
|
cache_control = image_params.get("cache_control")
|
|
base_image = Image.autodetect(source)
|
|
return cls(
|
|
source=base_image.source,
|
|
media_type=base_image.media_type,
|
|
data=base_image.data,
|
|
cache_control=cache_control,
|
|
)
|
|
|
|
def to_anthropic(self) -> dict[str, Any]:
|
|
"""Override Anthropic return with cache_control."""
|
|
result = super().to_anthropic()
|
|
if self.cache_control:
|
|
result["cache_control"] = self.cache_control
|
|
return result
|
|
|
|
|
|
class PDF(BaseModel):
|
|
source: str | Path = Field(description="URL, file path, or base64 data of the PDF")
|
|
media_type: str = Field(
|
|
description="MIME type of the PDF", default="application/pdf"
|
|
)
|
|
data: str | None = Field(None, description="Base64 encoded PDF data", repr=False)
|
|
|
|
@classmethod
|
|
def autodetect(cls, source: str | Path) -> PDF:
|
|
"""Attempt to autodetect a PDF from a source string or Path.
|
|
Args:
|
|
source (Union[str,path]): The source string or path.
|
|
Returns:
|
|
A PDF if the source is detected to be a valid PDF.
|
|
Raises:
|
|
ValueError: If the source is not detected to be a valid PDF.
|
|
"""
|
|
if isinstance(source, str):
|
|
if cls.is_base64(source):
|
|
return cls.from_base64(source)
|
|
elif source.startswith(("http://", "https://")):
|
|
return cls.from_url(source)
|
|
elif source.startswith("gs://"):
|
|
return cls.from_gs_url(source)
|
|
|
|
try:
|
|
if Path(source).is_file():
|
|
return cls.from_path(source)
|
|
except FileNotFoundError as err:
|
|
raise MultimodalError(
|
|
"PDF file not found",
|
|
content_type="pdf",
|
|
file_path=str(source),
|
|
) from err
|
|
except OSError as e:
|
|
if e.errno == 63: # File name too long
|
|
raise MultimodalError(
|
|
"PDF file name too long",
|
|
content_type="pdf",
|
|
file_path=str(source),
|
|
) from e
|
|
raise MultimodalError(
|
|
"Unable to read PDF file",
|
|
content_type="pdf",
|
|
file_path=str(source),
|
|
) from e
|
|
|
|
return cls.from_raw_base64(source)
|
|
elif isinstance(source, Path):
|
|
return cls.from_path(source)
|
|
|
|
@classmethod
|
|
def autodetect_safely(cls, source: Union[str, Path]) -> Union[PDF, str]: # noqa: UP007
|
|
"""Safely attempt to autodetect a PDF from a source string or path.
|
|
|
|
Args:
|
|
source (Union[str,path]): The source string or path.
|
|
Returns:
|
|
A PDF if the source is detected to be a valid PDF, otherwise
|
|
the source itself as a string.
|
|
"""
|
|
try:
|
|
return cls.autodetect(source)
|
|
except ValueError:
|
|
return str(source)
|
|
|
|
@classmethod
|
|
def is_base64(cls, s: str) -> bool:
|
|
return bool(re.match(r"^data:application/pdf;base64,", s))
|
|
|
|
@classmethod
|
|
def from_base64(cls, data_uri: str) -> PDF:
|
|
header, encoded = data_uri.split(",", 1)
|
|
media_type = header.split(":")[1].split(";")[0]
|
|
if media_type not in VALID_PDF_MIME_TYPES:
|
|
raise ValueError(f"Unsupported PDF format: {media_type}")
|
|
return cls(
|
|
source=data_uri,
|
|
media_type=media_type,
|
|
data=encoded,
|
|
)
|
|
|
|
@classmethod
|
|
@lru_cache
|
|
def from_path(cls, path: str | Path) -> PDF:
|
|
path = Path(path)
|
|
if not path.is_file():
|
|
raise FileNotFoundError(f"PDF file not found: {path}")
|
|
|
|
if path.stat().st_size == 0:
|
|
raise ValueError("PDF file is empty")
|
|
|
|
media_type, _ = mimetypes.guess_type(str(path))
|
|
if media_type not in VALID_PDF_MIME_TYPES:
|
|
raise ValueError(f"Unsupported PDF format: {media_type}")
|
|
|
|
data = base64.b64encode(path.read_bytes()).decode("utf-8")
|
|
return cls(source=path, media_type=media_type, data=data)
|
|
|
|
@classmethod
|
|
def from_raw_base64(cls, data: str) -> PDF:
|
|
try:
|
|
decoded = base64.b64decode(data)
|
|
# Check if it's a valid PDF by looking for the PDF header
|
|
if decoded.startswith(b"%PDF-"):
|
|
return cls(
|
|
source=data,
|
|
media_type="application/pdf",
|
|
data=data,
|
|
)
|
|
raise ValueError("Invalid PDF format")
|
|
except Exception as e:
|
|
raise ValueError("Invalid or unsupported base64 PDF data") from e
|
|
|
|
@classmethod
|
|
def from_gs_url(cls, data_uri: str, timeout: int = 30) -> PDF:
|
|
"""
|
|
Create a PDF instance from a Google Cloud Storage URL.
|
|
|
|
Args:
|
|
data_uri: GCS URL starting with gs://
|
|
timeout: Request timeout in seconds (default: 30)
|
|
"""
|
|
if not data_uri.startswith("gs://"):
|
|
raise ValueError("URL must start with gs://")
|
|
|
|
public_url = f"https://storage.googleapis.com/{data_uri[5:]}"
|
|
|
|
try:
|
|
response = requests.get(public_url, timeout=timeout)
|
|
response.raise_for_status()
|
|
media_type = response.headers.get("Content-Type", "application/pdf")
|
|
if media_type not in VALID_PDF_MIME_TYPES:
|
|
raise ValueError(f"Unsupported PDF format: {media_type}")
|
|
|
|
data = base64.b64encode(response.content).decode("utf-8")
|
|
|
|
return cls(source=data_uri, media_type=media_type, data=data)
|
|
except requests.RequestException as e:
|
|
raise ValueError(
|
|
"Failed to access GCS PDF (must be publicly readable)"
|
|
) from e
|
|
|
|
@classmethod
|
|
@lru_cache
|
|
def from_url(cls, url: str) -> PDF:
|
|
if url.startswith("gs://"):
|
|
return cls.from_gs_url(url)
|
|
parsed_url = urlparse(url)
|
|
media_type, _ = mimetypes.guess_type(parsed_url.path)
|
|
|
|
if not media_type:
|
|
try:
|
|
response = requests.head(url, allow_redirects=True)
|
|
media_type = response.headers.get("Content-Type")
|
|
except requests.RequestException as e:
|
|
raise ValueError("Failed to fetch PDF from URL") from e
|
|
|
|
if media_type not in VALID_PDF_MIME_TYPES:
|
|
raise ValueError(f"Unsupported PDF format: {media_type}")
|
|
return cls(source=url, media_type=media_type, data=None)
|
|
|
|
def to_mistral(self) -> dict[str, Any]:
|
|
if (
|
|
isinstance(self.source, str)
|
|
and self.source.startswith(("http://", "https://"))
|
|
and not self.data
|
|
):
|
|
return {
|
|
"type": "document_url",
|
|
"document_url": self.source,
|
|
}
|
|
raise ValueError("Mistral only supports document URLs for now")
|
|
|
|
def to_openai(self, mode: Mode) -> dict[str, Any]:
|
|
"""Convert to OpenAI's document format."""
|
|
input_file_type = (
|
|
"input_file"
|
|
if mode in {Mode.RESPONSES_TOOLS, Mode.RESPONSES_TOOLS_WITH_INBUILT_TOOLS}
|
|
else "file"
|
|
)
|
|
|
|
if (
|
|
isinstance(self.source, str)
|
|
and self.source.startswith(("http://", "https://"))
|
|
and not self.data
|
|
):
|
|
# Fetch the file from URL and convert to base64
|
|
data = requests.get(self.source)
|
|
data = base64.b64encode(data.content).decode("utf-8")
|
|
if mode in {Mode.RESPONSES_TOOLS, Mode.RESPONSES_TOOLS_WITH_INBUILT_TOOLS}:
|
|
return {
|
|
"type": input_file_type,
|
|
"filename": self.source,
|
|
"file_data": f"data:{self.media_type};base64,{data}",
|
|
}
|
|
else:
|
|
return {
|
|
"type": input_file_type,
|
|
"file": {
|
|
"filename": self.source,
|
|
"file_data": f"data:{self.media_type};base64,{data}",
|
|
},
|
|
}
|
|
elif self.data or self.is_base64(str(self.source)):
|
|
data = self.data or str(self.source).split(",", 1)[1]
|
|
if mode in {Mode.RESPONSES_TOOLS, Mode.RESPONSES_TOOLS_WITH_INBUILT_TOOLS}:
|
|
return {
|
|
"type": input_file_type,
|
|
"filename": (
|
|
self.source
|
|
if isinstance(self.source, str)
|
|
else str(self.source)
|
|
),
|
|
"file_data": f"data:{self.media_type};base64,{data}",
|
|
}
|
|
else:
|
|
return {
|
|
"type": input_file_type,
|
|
"file": {
|
|
"filename": (
|
|
self.source
|
|
if isinstance(self.source, str)
|
|
else str(self.source)
|
|
),
|
|
"file_data": f"data:{self.media_type};base64,{data}",
|
|
},
|
|
}
|
|
else:
|
|
raise ValueError("PDF data is missing for base64 encoding.")
|
|
|
|
def to_anthropic(self) -> dict[str, Any]:
|
|
"""Convert to Anthropic's document format."""
|
|
if (
|
|
isinstance(self.source, str)
|
|
and self.source.startswith(("http://", "https://"))
|
|
and not self.data
|
|
):
|
|
return {
|
|
"type": "document",
|
|
"source": {
|
|
"type": "url",
|
|
"url": self.source,
|
|
},
|
|
}
|
|
else:
|
|
if not self.data:
|
|
self.data = requests.get(str(self.source)).content # type: ignore
|
|
self.data = base64.b64encode(self.data).decode("utf-8") # type: ignore
|
|
|
|
return {
|
|
"type": "document",
|
|
"source": {
|
|
"type": "base64",
|
|
"media_type": self.media_type,
|
|
"data": self.data,
|
|
},
|
|
}
|
|
|
|
def to_genai(self):
|
|
try:
|
|
from google.genai import types
|
|
except ImportError as err:
|
|
raise ImportError(
|
|
"google-genai package is required for GenAI integration. Install with: pip install google-genai"
|
|
) from err
|
|
|
|
if (
|
|
isinstance(self.source, str)
|
|
and self.source.startswith(("http://", "https://"))
|
|
and not self.data
|
|
):
|
|
# Fetch the file from URL and convert to base64
|
|
data = requests.get(self.source).content
|
|
data = base64.b64encode(data).decode("utf-8")
|
|
return types.Part.from_bytes(
|
|
data=base64.b64decode(data),
|
|
mime_type=self.media_type,
|
|
)
|
|
|
|
if self.data:
|
|
return types.Part.from_bytes(
|
|
data=base64.b64decode(self.data),
|
|
mime_type=self.media_type,
|
|
)
|
|
|
|
raise ValueError("Unsupported PDF format")
|
|
|
|
def to_bedrock(self, name: str | None = None) -> dict[str, Any]:
|
|
"""Convert to Bedrock's document format."""
|
|
# Determine the document name
|
|
if name is None:
|
|
if isinstance(self.source, Path):
|
|
name = self.source.name
|
|
elif isinstance(self.source, str):
|
|
# Try to extract filename from path or URL
|
|
if self.source.startswith(("http://", "https://", "gs://")):
|
|
name = Path(urlparse(self.source).path).name or "document"
|
|
else:
|
|
name = (
|
|
Path(self.source).name
|
|
if Path(self.source).exists()
|
|
else "document"
|
|
)
|
|
else:
|
|
name = "document"
|
|
|
|
# Sanitize name according to Bedrock requirements
|
|
# Only allow alphanumeric, whitespace (max one in row), hyphens, parentheses, square brackets
|
|
name = re.sub(r"[^\w\s\-\(\)\[\]]", "", name)
|
|
name = re.sub(r"\s+", " ", name) # Consolidate whitespace
|
|
name = name.strip()
|
|
|
|
# Handle S3 URIs
|
|
if isinstance(self.source, str) and self.source.startswith("s3://"):
|
|
# Parse S3 URI: s3://bucket/key
|
|
s3_match = re.match(r"s3://([^/]+)/(.*)", self.source)
|
|
if not s3_match:
|
|
raise ValueError(f"Invalid S3 URI format: {self.source}")
|
|
|
|
bucket = s3_match.group(1)
|
|
key = s3_match.group(2)
|
|
|
|
# Note: bucketOwner is optional but recommended for cross-account access
|
|
return {
|
|
"document": {
|
|
"format": "pdf",
|
|
"name": name,
|
|
"source": {
|
|
"s3Location": {
|
|
"uri": self.source
|
|
# "bucketOwner": "account-id" # Optional, can be added by user
|
|
}
|
|
},
|
|
}
|
|
}
|
|
|
|
# Handle bytes-based sources (base64 only)
|
|
if not self.data:
|
|
raise ValueError(
|
|
"PDF data is missing. Provide base64-encoded data or use an s3:// source."
|
|
)
|
|
else:
|
|
# Decode base64 data to bytes
|
|
pdf_bytes = base64.b64decode(self.data)
|
|
|
|
return {
|
|
"document": {"format": "pdf", "name": name, "source": {"bytes": pdf_bytes}}
|
|
}
|
|
|
|
|
|
class PDFWithCacheControl(PDF):
|
|
"""PDF with Anthropic prompt caching support."""
|
|
|
|
def to_anthropic(self) -> dict[str, Any]:
|
|
"""Override Anthropic return with cache_control."""
|
|
result = super().to_anthropic()
|
|
result["cache_control"] = {"type": "ephemeral"}
|
|
return result
|
|
|
|
|
|
class PDFWithGenaiFile(PDF):
|
|
@classmethod
|
|
def from_new_genai_file(
|
|
cls, file_path: str, retry_delay: int = 10, max_retries: int = 20
|
|
) -> PDFWithGenaiFile:
|
|
"""Create a new PDFWithGenaiFile from a file path."""
|
|
from google.genai.types import FileState
|
|
import time
|
|
from google.genai import Client
|
|
|
|
client = Client()
|
|
file = client.files.upload(file=file_path)
|
|
while file.state != FileState.ACTIVE:
|
|
time.sleep(retry_delay)
|
|
file = client.files.get(name=file.name) # type: ignore
|
|
if max_retries > 0:
|
|
max_retries -= 1
|
|
else:
|
|
raise Exception(
|
|
"Max retries reached. File upload has been started but is still pending"
|
|
)
|
|
|
|
return cls(source=file.uri, media_type=file.mime_type, data=None) # type: ignore
|
|
|
|
@classmethod
|
|
def from_existing_genai_file(cls, file_name: str) -> PDFWithGenaiFile:
|
|
"""Create a new PDFWithGenaiFile from a file URL."""
|
|
from google.genai import types
|
|
from google.genai.types import FileState
|
|
from google.genai import Client
|
|
|
|
client = Client()
|
|
file = client.files.get(name=file_name)
|
|
if file.source == types.FileSource.UPLOADED and file.state == FileState.ACTIVE:
|
|
return cls(
|
|
source=file.uri, # type: ignore
|
|
media_type=file.mime_type, # type: ignore
|
|
data=None,
|
|
)
|
|
else:
|
|
raise ValueError("We only support uploaded PDFs for now")
|
|
|
|
def to_genai(self):
|
|
try:
|
|
from google.genai import types
|
|
except ImportError as err:
|
|
raise ImportError(
|
|
"google-genai package is required for GenAI integration. Install with: pip install google-genai"
|
|
) from err
|
|
|
|
if (
|
|
self.source
|
|
and isinstance(self.source, str)
|
|
and "https://generativelanguage.googleapis.com/v1beta/files/" in self.source
|
|
):
|
|
return types.Part.from_uri(
|
|
file_uri=self.source,
|
|
mime_type=self.media_type,
|
|
)
|
|
|
|
return super().to_genai()
|
|
|
|
|
|
def convert_contents(
|
|
contents: Union[ # noqa: UP007
|
|
str,
|
|
dict[str, Any],
|
|
Image,
|
|
Audio,
|
|
list[Union[str, dict[str, Any], Image, Audio]], # noqa: UP007
|
|
],
|
|
mode: Mode,
|
|
) -> Union[str, list[dict[str, Any]]]: # noqa: UP007
|
|
"""Convert content items to the appropriate format based on the specified mode."""
|
|
if isinstance(contents, str):
|
|
return contents
|
|
if isinstance(contents, (Image, Audio, PDF)) or isinstance(contents, dict):
|
|
contents = [contents]
|
|
|
|
converted_contents: list[dict[str, Union[str, Image]]] = [] # noqa: UP007
|
|
text_file_type = (
|
|
"input_text"
|
|
if mode in {Mode.RESPONSES_TOOLS, Mode.RESPONSES_TOOLS_WITH_INBUILT_TOOLS}
|
|
else "text"
|
|
)
|
|
for content in contents:
|
|
if isinstance(content, str):
|
|
converted_contents.append({"type": text_file_type, "text": content})
|
|
elif isinstance(content, dict):
|
|
converted_contents.append(content)
|
|
elif isinstance(content, (Image, Audio, PDF)):
|
|
if mode in {
|
|
Mode.ANTHROPIC_JSON,
|
|
Mode.ANTHROPIC_TOOLS,
|
|
Mode.ANTHROPIC_REASONING_TOOLS,
|
|
}:
|
|
converted_contents.append(content.to_anthropic())
|
|
elif mode in {Mode.GEMINI_JSON, Mode.GEMINI_TOOLS}:
|
|
raise NotImplementedError("Gemini is not supported yet")
|
|
elif mode in {
|
|
Mode.MISTRAL_STRUCTURED_OUTPUTS,
|
|
Mode.MISTRAL_TOOLS,
|
|
} and isinstance(content, (PDF)):
|
|
converted_contents.append(content.to_mistral()) # type: ignore
|
|
else:
|
|
converted_contents.append(content.to_openai(mode))
|
|
else:
|
|
raise ValueError(f"Unsupported content type: {type(content)}")
|
|
return converted_contents
|
|
|
|
|
|
def autodetect_media(
|
|
source: str | Path | Image | Audio | PDF,
|
|
) -> Image | Audio | PDF | str:
|
|
"""Autodetect images, audio, or PDFs from a given source.
|
|
|
|
Args:
|
|
source: URL, file path, Path, or data URI to inspect.
|
|
|
|
Returns:
|
|
The detected :class:`Image`, :class:`Audio`, or :class:`PDF` instance.
|
|
If detection fails, the original source is returned.
|
|
"""
|
|
if isinstance(source, (Image, Audio, PDF)):
|
|
return source
|
|
|
|
# Normalize once for cheap checks and mimetype guess
|
|
source = str(source)
|
|
|
|
if source.startswith("data:image/"):
|
|
return Image.autodetect_safely(source)
|
|
if source.startswith("data:audio/"):
|
|
return Audio.autodetect_safely(source)
|
|
if source.startswith("data:application/pdf"):
|
|
return PDF.autodetect_safely(source)
|
|
|
|
media_type, _ = mimetypes.guess_type(source)
|
|
if media_type in VALID_MIME_TYPES:
|
|
return Image.autodetect_safely(source)
|
|
if media_type in VALID_AUDIO_MIME_TYPES:
|
|
return Audio.autodetect_safely(source)
|
|
if media_type in VALID_PDF_MIME_TYPES:
|
|
return PDF.autodetect_safely(source)
|
|
|
|
for cls in (Image, Audio, PDF):
|
|
item = cls.autodetect_safely(source) # type: ignore[arg-type]
|
|
if not isinstance(item, str):
|
|
return item
|
|
return source
|
|
|
|
|
|
def convert_messages(
|
|
messages: list[
|
|
dict[
|
|
str,
|
|
Union[ # noqa: UP007
|
|
str,
|
|
dict[str, Any],
|
|
Image,
|
|
Audio,
|
|
PDF,
|
|
list[Union[str, dict[str, Any], Image, Audio, PDF]], # noqa: UP007
|
|
],
|
|
]
|
|
],
|
|
mode: Mode,
|
|
autodetect_images: bool = False,
|
|
) -> list[dict[str, Any]]:
|
|
"""Convert messages to the appropriate format based on the specified mode."""
|
|
converted_messages = []
|
|
|
|
def is_image_params(x: Any) -> bool:
|
|
return isinstance(x, dict) and x.get("type") == "image" and "source" in x # type: ignore
|
|
|
|
for message in messages:
|
|
if "type" in message:
|
|
if message["type"] in {"audio", "image"}:
|
|
converted_messages.append(message) # type: ignore
|
|
continue
|
|
else:
|
|
raise ValueError(f"Unsupported message type: {message['type']}")
|
|
role = message["role"]
|
|
content = message["content"] or []
|
|
other_kwargs = {
|
|
k: v for k, v in message.items() if k not in ["role", "content", "type"]
|
|
}
|
|
if autodetect_images:
|
|
if isinstance(content, list):
|
|
new_content: list[str | dict[str, Any] | Image | Audio | PDF] = [] # noqa: UP007
|
|
for item in content:
|
|
if isinstance(item, str):
|
|
new_content.append(autodetect_media(item))
|
|
elif is_image_params(item):
|
|
new_content.append(
|
|
ImageWithCacheControl.from_image_params(
|
|
cast(ImageParams, item)
|
|
)
|
|
)
|
|
else:
|
|
new_content.append(item)
|
|
content = new_content
|
|
elif isinstance(content, str):
|
|
content = autodetect_media(content)
|
|
elif is_image_params(content):
|
|
content = ImageWithCacheControl.from_image_params(
|
|
cast(ImageParams, content)
|
|
)
|
|
if isinstance(content, str):
|
|
converted_messages.append( # type: ignore
|
|
{"role": role, "content": content, **other_kwargs}
|
|
)
|
|
else:
|
|
# At this point content is narrowed to non-str types accepted by convert_contents
|
|
converted_content = convert_contents(content, mode) # type: ignore
|
|
converted_messages.append( # type: ignore
|
|
{"role": role, "content": converted_content, **other_kwargs}
|
|
)
|
|
return converted_messages # type: ignore
|
|
|
|
|
|
def extract_genai_multimodal_content(
|
|
contents: list[Any],
|
|
autodetect_images: bool = True,
|
|
):
|
|
"""
|
|
Convert Typed Contents to the appropriate format for Google GenAI.
|
|
"""
|
|
from google.genai import types
|
|
|
|
result: list[Union[types.Content, types.File]] = [] # noqa: UP007
|
|
for content in contents:
|
|
# Check for Files
|
|
if isinstance(content, types.File):
|
|
result.append(content)
|
|
continue
|
|
|
|
# We only want to do the conversion for the Image type
|
|
if not isinstance(content, types.Content):
|
|
raise ValueError(
|
|
f"Unsupported content type: {type(content)}. This should only be used for the Google types"
|
|
)
|
|
# Cast to list of Parts
|
|
content = cast(types.Content, content)
|
|
converted_contents: list[types.Part] = []
|
|
|
|
if not content.parts:
|
|
raise ValueError("Content parts are empty")
|
|
|
|
# Now we need to support a few cases
|
|
for content_part in content.parts:
|
|
if content_part.text and autodetect_images:
|
|
converted_item = autodetect_media(content_part.text)
|
|
|
|
if isinstance(converted_item, (Image, Audio, PDF)):
|
|
converted_contents.append(converted_item.to_genai())
|
|
continue
|
|
|
|
converted_contents.append(content_part)
|
|
else:
|
|
converted_contents.append(content_part)
|
|
|
|
result.append(types.Content(parts=converted_contents, role=content.role))
|
|
|
|
return result
|