ultimate-gemini-mcp 3.0.7__py3-none-any.whl

src/services/gemini_client.py

@@ -0,0 +1,304 @@
+"""
+Gemini API client for Gemini 3 Pro Image generation.
+Uses the official Google GenAI SDK.
+"""
+
+import asyncio
+import base64
+import io
+import logging
+from functools import partial
+from typing import Any
+
+from google import genai
+from google.genai import types
+from PIL import Image
+
+from ..config.constants import GEMINI_MODELS
+from ..core.exceptions import (
+    APIError,
+    AuthenticationError,
+    ContentPolicyError,
+    RateLimitError,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class GeminiClient:
+    """Client for Gemini 3 Pro Image API using official Google GenAI SDK."""
+
+    def __init__(self, api_key: str, timeout: int = 60):
+        """
+        Initialize Gemini client.
+
+        Args:
+            api_key: Gemini API key
+            timeout: Request timeout in seconds
+        """
+        self.api_key = api_key
+        self.timeout = timeout
+        self.client = genai.Client(api_key=api_key)
+
+    async def generate_image(
+        self,
+        prompt: str,
+        *,
+        model: str = "gemini-3-pro-image-preview",
+        reference_images: list[str] | None = None,
+        aspect_ratio: str | None = None,
+        image_size: str = "2K",
+        response_modalities: list[str] | None = None,
+        enable_google_search: bool = False,
+        **kwargs: Any,
+    ) -> dict[str, Any]:
+        """
+        Generate or edit an image using Gemini 3 Pro Image.
+
+        Args:
+            prompt: Text prompt for image generation or editing instruction
+            model: Model to use (default: gemini-3-pro-image-preview)
+            reference_images: List of base64-encoded reference images (up to 14)
+            aspect_ratio: Desired aspect ratio (optional)
+            image_size: Image resolution (1K, 2K, 4K - default: 2K)
+            response_modalities: Response types (TEXT, IMAGE - default: ["TEXT", "IMAGE"])
+            enable_google_search: Enable Google Search grounding for real-time data
+            **kwargs: Additional parameters
+
+        Returns:
+            Dict with 'images' key containing list of base64-encoded image data,
+            'thoughts' key for thinking process, and 'text' key for text responses
+
+        Raises:
+            APIError: If the API request fails
+        """
+        model_id = GEMINI_MODELS.get(model, model)
+
+        try:
+            # Build contents list with reference images and prompt
+            contents: list[Any] = []
+
+            # Add reference images if provided (up to 14)
+            if reference_images:
+                for ref_image_b64 in reference_images[:14]:  # Limit to max 14
+                    # Decode base64 to bytes for PIL Image
+                    image_bytes = base64.b64decode(ref_image_b64)
+                    image = Image.open(io.BytesIO(image_bytes))
+                    contents.append(image)
+
+            # Add text prompt
+            contents.append(prompt)
+
+            # Build configuration
+            if response_modalities is None:
+                response_modalities = ["TEXT", "IMAGE"]
+
+            # Build image config (SDK 1.52+ supports both aspect_ratio and image_size)
+            image_config = types.ImageConfig(
+                aspect_ratio=aspect_ratio if aspect_ratio else None,
+                image_size=image_size if image_size else None,
+            )
+
+            # Build generation config
+            config_args: dict[str, Any] = {
+                "response_modalities": response_modalities,
+                "image_config": image_config,
+            }
+
+            # Add Google Search grounding if enabled
+            if enable_google_search:
+                config_args["tools"] = [{"google_search": {}}]
+
+            config = types.GenerateContentConfig(**config_args)
+
+            logger.info(f"Generating image with model: {model_id}")
+            logger.info(f"Contents: {len(contents)} items")
+            logger.info(f"Config: {config}")
+            logger.info(f"Aspect ratio: {aspect_ratio}, Image size: {image_size}")
+
+            # Generate content using official SDK (run in executor since it's synchronous)
+            loop = asyncio.get_event_loop()
+            response = await loop.run_in_executor(
+                None,
+                partial(
+                    self.client.models.generate_content,
+                    model=model_id,
+                    contents=contents,
+                    config=config,
+                ),
+            )
+
+            # Extract images, thoughts, and text from response
+            extraction_result = self._extract_content_from_response(response)
+            images = extraction_result["images"]
+            thoughts = extraction_result["thoughts"]
+            text_parts = extraction_result["text"]
+
+            if not images and "IMAGE" in response_modalities:
+                logger.error(
+                    f"No images extracted from response. Response has {len(response.parts)} parts"
+                )
+                logger.error(f"Thoughts extracted: {len(thoughts)}, Text parts: {len(text_parts)}")
+                logger.error(f"Response_modalities: {response_modalities}")
+                for idx, part in enumerate(response.parts):
+                    logger.error(
+                        f"  Part {idx}: has_inline_data={hasattr(part, 'inline_data')}, has_text={hasattr(part, 'text')}, thought={getattr(part, 'thought', None)}, thought_sig={hasattr(part, 'thought_signature')}"
+                    )
+                raise APIError("No image data found in Gemini API response")
+
+            result = {
+                "images": images,
+                "text": text_parts,
+                "thoughts": thoughts,
+                "model": model,
+            }
+
+            # Include grounding metadata if Google Search was used
+            if enable_google_search and hasattr(response, "grounding_metadata"):
+                result["grounding_metadata"] = response.grounding_metadata
+
+            return result
+
+        except Exception as e:
+            logger.error(f"Gemini API request failed: {e}")
+            self._handle_exception(e)
+            raise APIError(f"Gemini API request failed: {e}") from e
+
+    async def generate_text(
+        self,
+        prompt: str,
+        *,
+        model: str = "gemini-flash-latest",
+        system_instruction: str | None = None,
+    ) -> str:
+        """
+        Generate text using Gemini (for prompt enhancement).
+
+        Args:
+            prompt: Text prompt
+            model: Model to use
+            system_instruction: Optional system instruction
+
+        Returns:
+            Generated text response
+        """
+        model_id = GEMINI_MODELS.get(model, model)
+
+        try:
+            # Build config with proper types instead of using **kwargs
+            config = (
+                types.GenerateContentConfig(system_instruction=system_instruction)
+                if system_instruction
+                else None
+            )
+
+            # Run in executor since genai SDK is synchronous
+            loop = asyncio.get_event_loop()
+            response = await loop.run_in_executor(
+                None,
+                partial(
+                    self.client.models.generate_content,
+                    model=model_id,
+                    contents=prompt,
+                    config=config,
+                ),
+            )
+
+            # Extract text from response
+            return response.text or ""
+
+        except Exception as e:
+            logger.error(f"Gemini text generation failed: {e}")
+            raise APIError(f"Gemini text generation failed: {e}") from e
+
+    def _extract_content_from_response(self, response: Any) -> dict[str, Any]:
+        """
+        Extract images, text, and thoughts from Gemini SDK response.
+
+        The genai SDK automatically handles thought signatures, so we just
+        need to extract the content.
+
+        Returns dict with keys:
+        - images: List of base64-encoded image data
+        - text: List of text strings
+        - thoughts: List of thought objects with images and text
+        """
+        images: list[str] = []
+        text_parts: list[str] = []
+        thoughts: list[dict[str, Any]] = []
+
+        try:
+            logger.info(f"Response has {len(response.parts)} parts")
+            # Iterate through all parts in the response
+            for idx, part in enumerate(response.parts):
+                logger.info(
+                    f"Part {idx}: type={type(part)}, has_inline_data={hasattr(part, 'inline_data')}, has_text={hasattr(part, 'text')}, has_thought={hasattr(part, 'thought')}, has_thought_sig={hasattr(part, 'thought_signature')}"
+                )
+                # Check if this is a thought (thinking process)
+                is_thought = getattr(part, "thought", False)
+
+                # Extract image data using SDK's as_image() method
+                if hasattr(part, "inline_data"):
+                    try:
+                        logger.info(f"Part {idx} has inline_data, attempting to extract image...")
+                        image = part.as_image()
+                        if image:
+                            logger.info(f"Successfully got PIL image: {image.size}")
+                            # Convert PIL Image to base64
+                            buffer = io.BytesIO()
+                            # Save as PNG - use positional argument instead of keyword
+                            image.save(buffer, "PNG")
+                            image_b64 = base64.b64encode(buffer.getvalue()).decode()
+
+                            if is_thought:
+                                logger.info(f"Adding to thoughts (is_thought={is_thought})")
+                                thoughts.append(
+                                    {"type": "image", "data": image_b64, "index": len(thoughts)}
+                                )
+                            else:
+                                logger.info(f"Adding to images (is_thought={is_thought})")
+                                images.append(image_b64)
+                        else:
+                            logger.warning(f"Part {idx}: as_image() returned None")
+                    except Exception as e:
+                        logger.error(f"Could not extract image from part {idx}: {e}", exc_info=True)
+
+                # Extract text
+                if hasattr(part, "text") and part.text:
+                    if is_thought:
+                        thoughts.append({"type": "text", "data": part.text, "index": len(thoughts)})
+                    else:
+                        text_parts.append(part.text)
+
+        except Exception as e:
+            logger.error(f"Error extracting content from response: {e}", exc_info=True)
+
+        logger.info(
+            f"Extraction complete: {len(images)} images, {len(text_parts)} text parts, {len(thoughts)} thoughts"
+        )
+        return {
+            "images": images,
+            "text": text_parts,
+            "thoughts": thoughts,
+        }
+
+    def _handle_exception(self, error: Exception) -> None:
+        """Handle exceptions from genai SDK."""
+        error_msg = str(error)
+
+        logger.error(f"API request failed: {error_msg}")
+
+        # Try to determine error type from message
+        if "authentication" in error_msg.lower() or "api key" in error_msg.lower():
+            raise AuthenticationError("Authentication failed. Please check your Gemini API key.")
+        elif "rate limit" in error_msg.lower() or "quota" in error_msg.lower():
+            raise RateLimitError("Rate limit exceeded. Please try again later.")
+        elif "safety" in error_msg.lower() or "blocked" in error_msg.lower():
+            raise ContentPolicyError(
+                "Content was blocked by safety filters. Please modify your prompt."
+            )
+
+    async def close(self) -> None:
+        """Close the Gemini client (genai SDK handles cleanup automatically)."""
+        # genai SDK doesn't require explicit cleanup
+        pass

src/services/image_service.py

@@ -0,0 +1,174 @@
+"""
+Image service for Gemini 3 Pro Image API.
+Provides interface for image generation using Gemini 3 Pro Image.
+"""
+
+import base64
+import logging
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+from ..config.constants import GEMINI_MODELS
+from ..core import sanitize_filename
+from ..core.exceptions import ImageProcessingError
+from .gemini_client import GeminiClient
+from .prompt_enhancer import PromptEnhancer
+
+logger = logging.getLogger(__name__)
+
+
+class ImageResult:
+    """Container for generated image data and metadata."""
+
+    def __init__(
+        self,
+        image_data: str,
+        prompt: str,
+        model: str,
+        index: int = 0,
+        metadata: dict[str, Any] | None = None,
+    ):
+        self.image_data = image_data  # Base64-encoded
+        self.prompt = prompt
+        self.model = model
+        self.index = index
+        self.metadata = metadata or {}
+        self.timestamp = datetime.now()
+
+    def save(self, output_dir: Path, filename: str | None = None) -> Path:
+        """Save image to disk."""
+        if filename is None:
+            filename = self._generate_filename()
+
+        output_path = output_dir / filename
+
+        try:
+            # Decode base64 and save
+            image_bytes = base64.b64decode(self.image_data)
+            output_path.write_bytes(image_bytes)
+            logger.info(f"Saved image to {output_path}")
+            return output_path
+        except Exception as e:
+            raise ImageProcessingError(f"Failed to save image: {e}") from e
+
+    def _generate_filename(self) -> str:
+        """Generate clean, short filename."""
+        timestamp = self.timestamp.strftime("%Y%m%d_%H%M%S")
+        # Shorten model name
+        model_short = self.model.replace("gemini-3-pro-image-preview", "gemini3").replace(
+            "imagen-4-", "img4-"
+        )
+        # Sanitize and shorten prompt (max 30 chars)
+        prompt_snippet = sanitize_filename(self.prompt[:30])
+        index_str = f"_{self.index + 1}" if self.index > 0 else ""
+        return f"{model_short}_{timestamp}_{prompt_snippet}{index_str}.png"
+
+    def get_size(self) -> int:
+        """Get image size in bytes."""
+        return len(base64.b64decode(self.image_data))
+
+
+class ImageService:
+    """Service for image generation using Gemini 3 Pro Image."""
+
+    def __init__(self, api_key: str, *, enable_enhancement: bool = True, timeout: int = 60):
+        """
+        Initialize image service.
+
+        Args:
+            api_key: API key for Gemini API
+            enable_enhancement: Enable automatic prompt enhancement
+            timeout: Request timeout in seconds
+        """
+        self.api_key = api_key
+        self.enable_enhancement = enable_enhancement
+        self.timeout = timeout
+
+        # Initialize Gemini client
+        self.gemini_client = GeminiClient(api_key, timeout)
+        self.prompt_enhancer: PromptEnhancer | None = None
+
+        if enable_enhancement:
+            # Prompt enhancer uses the same Gemini client
+            self.prompt_enhancer = PromptEnhancer(self.gemini_client)
+
+    async def generate(
+        self, prompt: str, *, model: str | None = None, enhance_prompt: bool = True, **kwargs: Any
+    ) -> list[ImageResult]:
+        """
+        Generate images using Gemini 3 Pro Image API.
+
+        Args:
+            prompt: Text prompt for image generation
+            model: Model to use (default: gemini-3-pro-image-preview)
+            enhance_prompt: Whether to enhance the prompt
+            **kwargs: Additional parameters (aspect_ratio, reference_images, etc.)
+
+        Returns:
+            List of ImageResult objects
+        """
+        # Use Gemini 3 Pro Image
+        if model is None:
+            model = "gemini-3-pro-image-preview"
+
+        if model not in GEMINI_MODELS:
+            raise ValueError(f"Unknown model: {model}. Only Gemini 3 Pro Image is supported.")
+
+        # Enhance prompt if enabled
+        original_prompt = prompt
+        enhancement_context = self._build_enhancement_context(kwargs)
+
+        if enhance_prompt and self.enable_enhancement and self.prompt_enhancer:
+            try:
+                result = await self.prompt_enhancer.enhance_prompt(
+                    prompt, context=enhancement_context
+                )
+                prompt = result["enhanced_prompt"]
+                logger.info(f"Prompt enhanced: {len(original_prompt)} -> {len(prompt)} chars")
+            except Exception as e:
+                logger.warning(f"Prompt enhancement failed: {e}")
+
+        # Generate images using Gemini API
+        return await self._generate_with_gemini(prompt, model, original_prompt, kwargs)
+
+    async def _generate_with_gemini(
+        self, prompt: str, model: str, original_prompt: str, params: dict[str, Any]
+    ) -> list[ImageResult]:
+        """Generate images using Gemini API."""
+        response = await self.gemini_client.generate_image(prompt=prompt, model=model, **params)
+
+        images = response["images"]
+        results = []
+
+        for i, image_data in enumerate(images):
+            result = ImageResult(
+                image_data=image_data,
+                prompt=original_prompt,
+                model=model,
+                index=i,
+                metadata={"enhanced_prompt": prompt, "api": "gemini", **params},
+            )
+            results.append(result)
+
+        return results
+
+    def _build_enhancement_context(self, params: dict[str, Any]) -> dict[str, Any]:
+        """Build context for prompt enhancement."""
+        context: dict[str, Any] = {}
+
+        if "reference_images" in params and params["reference_images"]:
+            context["has_reference_images"] = True
+            context["num_reference_images"] = len(params["reference_images"])
+
+        if "aspect_ratio" in params:
+            context["aspect_ratio"] = params["aspect_ratio"]
+
+        if params.get("enable_google_search"):
+            context["use_google_search"] = True
+
+        return context
+
+    async def close(self) -> None:
+        """Close Gemini client."""
+        await self.gemini_client.close()

src/services/prompt_enhancer.py

@@ -0,0 +1,137 @@
+"""
+Prompt enhancement service using Gemini Flash.
+Automatically optimizes prompts for better image generation results.
+"""
+
+import logging
+from typing import Any
+
+from .gemini_client import GeminiClient
+
+logger = logging.getLogger(__name__)
+
+
+PROMPT_ENHANCEMENT_SYSTEM_INSTRUCTION = """You are an expert prompt engineer for AI image generation models. Your task is to enhance user prompts to produce the best possible results.
+
+Follow these guidelines:
+1. Preserve the user's core intent and subject matter
+2. Add specific, professional details about:
+   - Composition (framing, perspective, angle)
+   - Lighting (type, quality, direction, mood)
+   - Materials and textures
+   - Atmosphere and mood
+   - Artistic style (if appropriate)
+3. Use photographic and cinematic terminology when relevant
+4. Be hyper-specific rather than generic
+5. For portraits: describe features, expressions, clothing
+6. For scenes: describe environment, weather, time of day
+7. Keep prompts concise but detailed (aim for 100-300 words)
+8. NEVER use hex color values (like #FF0000). Always describe colors using natural language (e.g., "dark red", "neon blue", "warm amber", "deep crimson")
+9. Output ONLY the enhanced prompt, no explanations"""
+
+
+class PromptEnhancer:
+    """Service for enhancing image generation prompts."""
+
+    def __init__(self, gemini_client: GeminiClient):
+        """
+        Initialize prompt enhancer.
+
+        Args:
+            gemini_client: Gemini client for text generation
+        """
+        self.gemini_client = gemini_client
+
+    async def enhance_prompt(
+        self,
+        original_prompt: str,
+        *,
+        context: dict[str, Any] | None = None,
+    ) -> dict[str, str]:
+        """
+        Enhance a prompt for better image generation.
+
+        Args:
+            original_prompt: Original user prompt
+            context: Optional context (features, image type, etc.)
+
+        Returns:
+            Dict with 'enhanced_prompt' and 'original_prompt'
+        """
+        # Build enhancement instruction
+        instruction = self._build_enhancement_instruction(original_prompt, context)
+
+        try:
+            enhanced = await self.gemini_client.generate_text(
+                prompt=instruction,
+                system_instruction=PROMPT_ENHANCEMENT_SYSTEM_INSTRUCTION,
+                model="gemini-flash-latest",
+            )
+
+            # Clean up the enhanced prompt
+            enhanced = enhanced.strip()
+
+            logger.info(f"Enhanced prompt: {len(original_prompt)} -> {len(enhanced)} chars")
+
+            return {
+                "original_prompt": original_prompt,
+                "enhanced_prompt": enhanced,
+            }
+
+        except Exception as e:
+            logger.warning(f"Prompt enhancement failed, using original: {e}")
+            return {
+                "original_prompt": original_prompt,
+                "enhanced_prompt": original_prompt,
+            }
+
+    def _build_enhancement_instruction(self, prompt: str, context: dict[str, Any] | None) -> str:
+        """Build the instruction for prompt enhancement."""
+        instruction_parts = [f"Enhance this image generation prompt:\n\n{prompt}"]
+
+        if context:
+            # Add context hints
+            if context.get("is_editing"):
+                instruction_parts.append("\nContext: This is for image editing/modification")
+
+            if context.get("maintain_character_consistency"):
+                instruction_parts.append(
+                    "\nIMPORTANT: Describe the character with specific, consistent features "
+                    "for use across multiple generations"
+                )
+
+            if context.get("blend_images"):
+                instruction_parts.append(
+                    "\nContext: Multiple images will be blended. Describe how elements "
+                    "should be composed naturally together"
+                )
+
+            if context.get("use_world_knowledge"):
+                instruction_parts.append(
+                    "\nContext: Include accurate real-world details for historical figures, "
+                    "landmarks, or factual scenarios"
+                )
+
+            if context.get("aspect_ratio"):
+                ratio = context["aspect_ratio"]
+                if ratio in ["16:9", "21:9"]:
+                    instruction_parts.append("\nFormat: Wide landscape composition")
+                elif ratio in ["9:16", "2:3", "3:4"]:
+                    instruction_parts.append("\nFormat: Vertical/portrait composition")
+
+        return "\n".join(instruction_parts)
+
+
+async def create_prompt_enhancer(api_key: str, timeout: int = 30) -> PromptEnhancer:
+    """
+    Factory function to create prompt enhancer.
+
+    Args:
+        api_key: Gemini API key
+        timeout: Request timeout
+
+    Returns:
+        PromptEnhancer instance
+    """
+    gemini_client = GeminiClient(api_key=api_key, timeout=timeout)
+    return PromptEnhancer(gemini_client)

src/tools/__init__.py

@@ -0,0 +1,11 @@
+"""Tools module for Ultimate Gemini MCP."""
+
+from .batch_generate import batch_generate_images, register_batch_generate_tool
+from .generate_image import generate_image_tool, register_generate_image_tool
+
+__all__ = [
+    "generate_image_tool",
+    "register_generate_image_tool",
+    "batch_generate_images",
+    "register_batch_generate_tool",
+]