Create images.py

gemini_client/images.py

@@ -0,0 +1,202 @@
+# -*- 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)