# -*- coding: utf-8 -*- import os import random import re from datetime import datetime from pathlib import Path from typing import Dict, Union, Optional from pydantic import BaseModel, field_validator from curl_cffi import CurlError from curl_cffi.requests import AsyncSession from requests.exceptions import HTTPError, RequestException # Ensure RequestException is imported from rich.console import Console console = Console() # Instantiate console for logging class Image(BaseModel): """ Represents a single image object returned from Gemini. Attributes: url (str): URL of the image. title (str): Title of the image (default: "[Image]"). alt (str): Optional description of the image. proxy (str | dict | None): Proxy used when saving the image. impersonate (str): Browser profile for curl_cffi to impersonate. """ url: str title: str = "[Image]" alt: str = "" proxy: Optional[Union[str, Dict[str, str]]] = None impersonate: str = "chrome110" def __str__(self): return f"{self.title}({self.url}) - {self.alt}" def __repr__(self): short_url = self.url if len(self.url) <= 50 else self.url[:20] + "..." + self.url[-20:] short_alt = self.alt[:30] + "..." if len(self.alt) > 30 else self.alt return f"Image(title='{self.title}', url='{short_url}', alt='{short_alt}')" async def save( self, path: str = "downloaded_images", filename: Optional[str] = None, cookies: Optional[dict] = None, verbose: bool = False, skip_invalid_filename: bool = True, ) -> Optional[str]: """ Save the image to disk using curl_cffi. Parameters: path: str, optional Directory to save the image (default "downloaded_images"). filename: str, optional Filename to use; if not provided, inferred from URL. cookies: dict, optional Cookies used for the image request. verbose: bool, optional If True, outputs status messages (default False). skip_invalid_filename: bool, optional If True, skips saving if the filename is invalid. Returns: Absolute path of the saved image if successful; None if skipped. Raises: HTTPError if the network request fails. RequestException/CurlError for other network errors. IOError if file writing fails. """ # Generate filename from URL if not provided if not filename: try: from urllib.parse import urlparse, unquote parsed_url = urlparse(self.url) base_filename = os.path.basename(unquote(parsed_url.path)) # Remove invalid characters for filenames safe_filename = re.sub(r'[<>:"/\|?*]', '_', base_filename) if safe_filename and len(safe_filename) > 0: filename = safe_filename else: filename = f"image_{random.randint(1000, 9999)}.jpg" except Exception: filename = f"image_{random.randint(1000, 9999)}.jpg" # Validate filename length try: _ = Path(filename) max_len = 255 if len(filename) > max_len: name, ext = os.path.splitext(filename) filename = name[:max_len - len(ext) - 1] + ext except (OSError, ValueError): if verbose: console.log(f"[yellow]Invalid filename generated: {filename}[/yellow]") if skip_invalid_filename: if verbose: console.log("[yellow]Skipping save due to invalid filename.[/yellow]") return None filename = f"image_{random.randint(1000, 9999)}.jpg" if verbose: console.log(f"[yellow]Using fallback filename: {filename}[/yellow]") # Prepare proxy dictionary for curl_cffi proxies_dict = None if isinstance(self.proxy, str): proxies_dict = {"http": self.proxy, "https": self.proxy} elif isinstance(self.proxy, dict): proxies_dict = self.proxy try: # Use AsyncSession from curl_cffi async with AsyncSession( cookies=cookies, proxies=proxies_dict, impersonate=self.impersonate # follow_redirects is handled automatically by curl_cffi ) as client: if verbose: console.log(f"Attempting to download image from: {self.url}") response = await client.get(self.url) response.raise_for_status() # Check content type content_type = response.headers.get("content-type", "").lower() if "image" not in content_type and verbose: console.log(f"[yellow]Warning: Content type is '{content_type}', not an image. Saving anyway.[/yellow]") # Create directory and save file dest_path = Path(path) dest_path.mkdir(parents=True, exist_ok=True) dest = dest_path / filename # Write image data to file dest.write_bytes(response.content) if verbose: console.log(f"Image saved successfully as {dest.resolve()}") return str(dest.resolve()) except HTTPError as e: console.log(f"[red]Error downloading image {self.url}: {e.response.status_code} {e}[/red]") raise except (RequestException, CurlError) as e: console.log(f"[red]Network error downloading image {self.url}: {e}[/red]") raise except IOError as e: console.log(f"[red]Error writing image file to {dest}: {e}[/red]") raise except Exception as e: console.log(f"[red]An unexpected error occurred during image save: {e}[/red]") raise class WebImage(Image): """ Represents an image retrieved from web search results. Returned when asking Gemini to "SEND an image of [something]". """ pass class GeneratedImage(Image): """ Represents an image generated by Google's AI image generator (e.g., ImageFX). Attributes: cookies (dict[str, str]): Cookies required for accessing the generated image URL, typically from the GeminiClient/Chatbot instance. """ cookies: Dict[str, str] # Updated validator for Pydantic V2 @field_validator("cookies") @classmethod def validate_cookies(cls, v: Dict[str, str]) -> Dict[str, str]: """Ensures cookies are provided for generated images.""" if not v or not isinstance(v, dict): raise ValueError("GeneratedImage requires a dictionary of cookies from the client.") return v async def save(self, **kwargs) -> Optional[str]: """ Save the generated image to disk. Parameters: filename: str, optional Filename to use. If not provided, a default name including a timestamp and part of the URL is used. Generated images are often in .png or .jpg format. Additional arguments are passed to Image.save. Returns: Absolute path of the saved image if successful, None if skipped. """ if "filename" not in kwargs: ext = ".jpg" if ".jpg" in self.url.lower() else ".png" url_part = self.url.split('/')[-1][:10] kwargs["filename"] = f"{datetime.now().strftime('%Y%m%d%H%M%S')}_{url_part}{ext}" # Pass the required cookies and other args (like impersonate) to the parent save method return await super().save(cookies=self.cookies, **kwargs)