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

src/services/imagen_client.py

@@ -0,0 +1,175 @@
+"""
+Imagen API client for Imagen 3, 4, and 4-Ultra models.
+Uses the predict API endpoint per Google's documentation.
+"""
+
+import logging
+from typing import Any
+
+import httpx
+
+from ..config.constants import IMAGEN_API_BASE, IMAGEN_MODELS
+from ..core.exceptions import (
+    APIError,
+    AuthenticationError,
+    ContentPolicyError,
+    RateLimitError,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class ImagenClient:
+    """Client for Imagen 3/4/Ultra API."""
+
+    def __init__(self, api_key: str, timeout: int = 60):
+        """
+        Initialize Imagen client.
+
+        Args:
+            api_key: Gemini/Google API key
+            timeout: Request timeout in seconds
+        """
+        self.api_key = api_key
+        self.timeout = timeout
+        self.base_url = IMAGEN_API_BASE
+        self.client = httpx.AsyncClient(timeout=timeout)
+
+    async def generate_image(
+        self,
+        prompt: str,
+        *,
+        model: str = "imagen-4-ultra",
+        number_of_images: int = 1,
+        aspect_ratio: str = "1:1",
+        output_format: str = "image/png",
+        person_generation: str = "allow_adult",
+        negative_prompt: str | None = None,
+        seed: int | None = None,
+        **kwargs: Any,
+    ) -> dict[str, Any]:
+        """
+        Generate images using Imagen API.
+
+        Args:
+            prompt: Text prompt for image generation
+            model: Imagen model to use (imagen-3, imagen-4, imagen-4-ultra)
+            number_of_images: Number of images to generate (1-4)
+            aspect_ratio: Image aspect ratio
+            output_format: Output MIME type (image/jpeg or image/png)
+            person_generation: Person generation policy
+            negative_prompt: Optional negative prompt
+            seed: Optional seed for reproducibility
+            **kwargs: Additional parameters
+
+        Returns:
+            Dict with 'images' key containing list of base64-encoded images
+
+        Raises:
+            APIError: If the API request fails
+        """
+        model_id = IMAGEN_MODELS.get(model, model)
+        url = f"{self.base_url}/{model_id}:predict"
+
+        # Build request body according to Imagen API
+        request_body: dict[str, Any] = {
+            "instances": [
+                {
+                    "prompt": prompt
+                }
+            ],
+            "parameters": {
+                "outputMimeType": output_format,
+                "sampleCount": number_of_images,
+                "personGeneration": person_generation,
+                "aspectRatio": aspect_ratio
+            }
+        }
+
+        # Add optional parameters
+        if negative_prompt:
+            request_body["instances"][0]["negativePrompt"] = negative_prompt
+
+        if seed is not None:
+            request_body["parameters"]["seed"] = seed
+
+        headers = {
+            "Content-Type": "application/json",
+        }
+
+        try:
+            logger.debug(f"Sending request to {url}")
+            # Add API key as query parameter
+            response = await self.client.post(
+                f"{url}?key={self.api_key}",
+                json=request_body,
+                headers=headers
+            )
+            response.raise_for_status()
+            data = response.json()
+
+            # Extract images from predictions
+            images = self._extract_images(data)
+
+            if not images:
+                raise APIError("No image data found in Imagen API response")
+
+            return {
+                "images": images,
+                "model": model,
+                "response": data
+            }
+
+        except httpx.HTTPStatusError as e:
+            self._handle_http_error(e)
+        except Exception as e:
+            logger.error(f"Imagen API request failed: {e}")
+            raise APIError(f"Imagen API request failed: {e}")
+
+    def _extract_images(self, response_data: dict[str, Any]) -> list[str]:
+        """Extract base64 image data from Imagen API response."""
+        images = []
+
+        try:
+            predictions = response_data.get("predictions", [])
+            for prediction in predictions:
+                # Imagen returns base64 data in bytesBase64Encoded field
+                image_data = prediction.get("bytesBase64Encoded")
+                if image_data:
+                    images.append(image_data)
+        except Exception as e:
+            logger.warning(f"Error extracting images from response: {e}")
+
+        return images
+
+    def _handle_http_error(self, error: httpx.HTTPStatusError) -> None:
+        """Handle HTTP errors and raise appropriate exceptions."""
+        status_code = error.response.status_code
+        error_text = error.response.text
+
+        logger.error(f"API request failed with status {status_code}: {error_text}")
+
+        if status_code == 401 or status_code == 403:
+            raise AuthenticationError(
+                "Authentication failed. Please check your API key.",
+                status_code=status_code
+            )
+        elif status_code == 429:
+            raise RateLimitError(
+                "Rate limit exceeded. Please try again later.",
+                status_code=status_code
+            )
+        elif status_code == 400 and "SAFETY" in error_text.upper():
+            raise ContentPolicyError(
+                "Content was blocked by safety filters. Please modify your prompt.",
+                status_code=status_code
+            )
+        else:
+            raise APIError(
+                f"API request failed with status {status_code}: {error_text}",
+                status_code=status_code
+            )
+
+    async def close(self) -> None:
+        """Close the HTTP client."""
+        await self.client.aclose()

src/services/prompt_enhancer.py

@@ -0,0 +1,140 @@
+"""
+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. 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",
+]

src/tools/batch_generate.py

@@ -0,0 +1,159 @@
+"""
+Batch image generation tool for processing multiple prompts efficiently.
+"""
+
+import asyncio
+import json
+import logging
+from typing import Any
+
+from ..config import MAX_BATCH_SIZE, get_settings
+from ..core import validate_batch_size, validate_prompts_list
+from .generate_image import generate_image_tool
+
+logger = logging.getLogger(__name__)
+
+
+async def batch_generate_images(
+    prompts: list[str],
+    model: str | None = None,
+    enhance_prompt: bool = True,
+    aspect_ratio: str = "1:1",
+    output_format: str = "png",
+    batch_size: int | None = None,
+    **shared_params: Any,
+) -> dict[str, Any]:
+    """
+    Generate multiple images from a list of prompts.
+
+    Args:
+        prompts: List of text prompts
+        model: Model to use for all images
+        enhance_prompt: Enhance all prompts
+        aspect_ratio: Aspect ratio for all images
+        output_format: Output format for all images
+        batch_size: Number of images to process in parallel (default: from config)
+        **shared_params: Additional parameters shared across all generations
+
+    Returns:
+        Dict with batch results
+    """
+    # Validate inputs
+    validate_prompts_list(prompts)
+
+    settings = get_settings()
+    if batch_size is None:
+        batch_size = settings.api.max_batch_size
+
+    validate_batch_size(batch_size, MAX_BATCH_SIZE)
+
+    # Prepare results
+    results = {
+        "success": True,
+        "total_prompts": len(prompts),
+        "batch_size": batch_size,
+        "completed": 0,
+        "failed": 0,
+        "results": []
+    }
+
+    # Process prompts in batches
+    for i in range(0, len(prompts), batch_size):
+        batch = prompts[i:i + batch_size]
+        logger.info(f"Processing batch {i // batch_size + 1}: {len(batch)} prompts")
+
+        # Create tasks for parallel processing
+        tasks = [
+            generate_image_tool(
+                prompt=prompt,
+                model=model,
+                enhance_prompt=enhance_prompt,
+                aspect_ratio=aspect_ratio,
+                output_format=output_format,
+                number_of_images=1,
+                **shared_params
+            )
+            for prompt in batch
+        ]
+
+        # Execute batch
+        batch_results = await asyncio.gather(*tasks, return_exceptions=True)
+
+        # Process results
+        for j, result in enumerate(batch_results):
+            prompt_index = i + j
+
+            if isinstance(result, Exception):
+                logger.error(f"Failed to generate image for prompt {prompt_index}: {result}")
+                results["failed"] += 1
+                results["results"].append({
+                    "prompt_index": prompt_index,
+                    "prompt": batch[j],
+                    "success": False,
+                    "error": str(result)
+                })
+            else:
+                results["completed"] += 1
+                results["results"].append({
+                    "prompt_index": prompt_index,
+                    "prompt": batch[j],
+                    **result
+                })
+
+    return results
+
+
+def register_batch_generate_tool(mcp_server: Any) -> None:
+    """Register batch_generate tool with MCP server."""
+
+    @mcp_server.tool()
+    async def batch_generate(
+        prompts: list[str],
+        model: str | None = None,
+        enhance_prompt: bool = True,
+        aspect_ratio: str = "1:1",
+        output_format: str = "png",
+        batch_size: int | None = None,
+        person_generation: str = "allow_adult",
+        negative_prompt: str | None = None,
+    ) -> str:
+        """
+        Generate multiple images from a list of prompts efficiently.
+
+        Processes prompts in parallel batches for optimal performance.
+        All images share the same generation settings.
+
+        Args:
+            prompts: List of text descriptions for image generation
+            model: Model to use for all images (default: gemini-2.5-flash-image)
+            enhance_prompt: Enhance all prompts automatically (default: True)
+            aspect_ratio: Aspect ratio for all images (default: 1:1)
+            output_format: Image format for all images (default: png)
+            batch_size: Parallel batch size (default: from config)
+            person_generation: Person policy for Imagen models (default: allow_adult)
+            negative_prompt: Negative prompt for Imagen models (optional)
+
+        Returns:
+            JSON string with batch results including individual image paths
+        """
+        try:
+            result = await batch_generate_images(
+                prompts=prompts,
+                model=model,
+                enhance_prompt=enhance_prompt,
+                aspect_ratio=aspect_ratio,
+                output_format=output_format,
+                batch_size=batch_size,
+                person_generation=person_generation,
+                negative_prompt=negative_prompt,
+            )
+
+            return json.dumps(result, indent=2)
+
+        except Exception as e:
+            logger.error(f"Batch generation error: {e}")
+            return json.dumps({
+                "success": False,
+                "error": str(e),
+                "error_type": type(e).__name__
+            }, indent=2)