Enhance image extraction

[dyskyside]

Refactor image extraction logic to improve error handling and URL processing.

[gemini-code-assist]

Summary of Changes

Hello @dyskyside, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request refactors the image extraction mechanism within the GenAI provider to be more robust and flexible. The core change introduces a two-tiered approach: first attempting to extract inline image data, and then, as a fallback, parsing text responses for image URLs to download and process. This enhancement improves the system's ability to handle various GenAI output formats and provides more detailed error reporting when images cannot be retrieved.

Highlights

• Enhanced Image Extraction Logic: The image extraction process has been significantly improved to handle more diverse GenAI responses. It now prioritizes extracting inline image data, and if unsuccessful, it attempts to parse text content for image URLs, download them, and convert them into PIL Image objects.
• URL-based Image Retrieval: New functionality has been added to detect and download images from URLs embedded within the text parts of the GenAI response. This is particularly useful for scenarios where the API might return image links instead of direct image data, such as when using certain proxies or models.
• Improved Error Handling and Debugging: Error messages for cases where no image is found have been made more informative, including snippets of text content from the response parts to aid in debugging. Robust error handling for network requests and image processing has also been integrated.
• Code Refactoring and Cleanup: Docstrings for the __init__ and generate_image methods have been streamlined, and logging messages have been updated for better clarity and conciseness.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request significantly enhances the image extraction logic in GenAIImageProvider by adding a fallback to download images from URLs found in the response text. This makes the provider more robust, especially when dealing with proxies or models that return image URLs in markdown. The new logic is well-structured with clear priorities for extraction.

All original comments have been kept as they do not contradict any of the provided rules. The feedback includes suggestions for restoring docstrings for __init__ and generate_image to improve maintainability, enhancing a debug log message for better context, and narrowing a broad exception catch for safer and more robust code.

Overall, this is a great improvement to the image extraction capabilities. Addressing the feedback will help maintain code quality and clarity.

[gemini-code-assist]

Catching a broad Exception is risky as it can hide unexpected errors. It's better to catch more specific exceptions that you expect to occur when a part is not an image, such as AttributeError (if inline_data is missing) or UnidentifiedImageError (if data is not a valid image). This makes the code safer and easier to debug.

                except (AttributeError, ValueError, UnidentifiedImageError):

[gemini-code-assist]

The detailed docstring for __init__ was removed. This docstring was valuable as it explained the purpose of each parameter, such as api_key, api_base, vertexai, project_id, and location. Please restore the docstring to maintain code clarity and help future developers understand how to initialize this class correctly.

        """
        Initialize GenAI image provider

        Args:
            api_key: Google API key (for AI Studio mode)
            api_base: API base URL (for proxies like aihubmix, AI Studio mode only)
            model: Model name to use
            vertexai: If True, use Vertex AI instead of AI Studio
            project_id: GCP project ID (required for Vertex AI mode)
            location: GCP region (for Vertex AI mode, default: us-central1)
        """

[gemini-code-assist]

The detailed docstring for generate_image has been removed. The previous docstring provided useful information about the arguments, including the supported values for resolution. Please restore it to improve code clarity and maintainability. The base class ImageProvider has a docstring, but it's good practice for the implementation to also have one, especially for implementation-specific details.

        """
        Generate image using Google GenAI SDK

        Args:
            prompt: The image generation prompt
            ref_images: Optional list of reference images
            aspect_ratio: Image aspect ratio
            resolution: Image resolution (supports "1K", "2K", "4K")

        Returns:
            Generated PIL Image object, or None if failed
        """

[gemini-code-assist]

This log message is less informative than the previous version. It no longer includes the number of reference images, aspect ratio, or resolution, which are all valuable for debugging. Please consider adding this information back into a single, comprehensive log message.

            logger.debug(f"Calling GenAI API (Model: {self.model}) with {len(ref_images) if ref_images else 0} reference images. Config: aspect_ratio={aspect_ratio}, resolution={resolution}")