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

src/server.py

@@ -0,0 +1,166 @@
+#!/usr/bin/env python3
+"""
+Ultimate Gemini MCP Server - Main Entry Point
+
+Unified MCP server supporting:
+- Gemini 2.5 Flash Image (with prompt enhancement and editing)
+- Imagen 3, 4, and 4-Ultra (with advanced controls)
+- Batch processing, prompt templates, and comprehensive features
+"""
+
+import logging
+import sys
+from pathlib import Path
+
+from fastmcp import FastMCP
+
+from .config import ALL_MODELS, get_settings
+from .tools import register_batch_generate_tool, register_generate_image_tool
+
+# Set up logging
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+    stream=sys.stderr,  # Important: use stderr for logging in MCP
+)
+
+logger = logging.getLogger(__name__)
+
+
+def create_app() -> FastMCP:
+    """
+    Create and configure the Ultimate Gemini MCP application.
+
+    This is the factory function used by FastMCP CLI.
+    """
+    logger.info("Initializing Ultimate Gemini MCP Server...")
+
+    try:
+        # Load settings (validates API key)
+        settings = get_settings()
+
+        logger.info(f"Output directory: {settings.output_dir}")
+        logger.info(f"Prompt enhancement: {settings.api.enable_prompt_enhancement}")
+        logger.info(f"Available models: {', '.join(ALL_MODELS.keys())}")
+
+        # Create FastMCP server
+        mcp = FastMCP(
+            "Ultimate Gemini MCP",
+            version="1.0.0",
+        )
+
+        # Register tools
+        register_generate_image_tool(mcp)
+        register_batch_generate_tool(mcp)
+
+        # Add resources
+        @mcp.resource("models://list")
+        def list_models() -> str:
+            """List all available image generation models."""
+            import json
+
+            models_info = {
+                "gemini": {
+                    "gemini-2.5-flash-image": {
+                        "name": "Gemini 2.5 Flash Image",
+                        "description": "Advanced image generation with editing and prompt enhancement",
+                        "features": [
+                            "Prompt enhancement",
+                            "Image editing",
+                            "Character consistency",
+                            "Multi-image blending",
+                            "World knowledge integration"
+                        ],
+                        "default": True
+                    }
+                },
+                "imagen": {
+                    "imagen-4": {
+                        "name": "Imagen 4",
+                        "description": "High-quality image generation with improved text rendering",
+                        "features": [
+                            "Enhanced quality",
+                            "Better text rendering",
+                            "Negative prompts",
+                            "Seed-based reproducibility",
+                            "Person generation controls",
+                            "Advanced controls"
+                        ]
+                    },
+                    "imagen-4-fast": {
+                        "name": "Imagen 4 Fast",
+                        "description": "Optimized for faster generation while maintaining good quality",
+                        "features": [
+                            "Faster generation speed",
+                            "Good quality output",
+                            "Negative prompts",
+                            "Seed-based reproducibility",
+                            "Person generation controls",
+                            "Cost-effective"
+                        ]
+                    },
+                    "imagen-4-ultra": {
+                        "name": "Imagen 4 Ultra",
+                        "description": "Highest quality with best prompt adherence",
+                        "features": [
+                            "Highest quality",
+                            "Best prompt adherence",
+                            "Professional results",
+                            "Enhanced text rendering",
+                            "Advanced controls"
+                        ]
+                    }
+                }
+            }
+
+            return json.dumps(models_info, indent=2)
+
+        @mcp.resource("settings://config")
+        def get_config() -> str:
+            """Get current server configuration."""
+            import json
+
+            config = {
+                "output_directory": str(settings.output_dir),
+                "prompt_enhancement_enabled": settings.api.enable_prompt_enhancement,
+                "batch_processing_enabled": settings.api.enable_batch_processing,
+                "default_gemini_model": settings.api.default_gemini_model,
+                "default_imagen_model": settings.api.default_imagen_model,
+                "max_batch_size": settings.api.max_batch_size,
+                "request_timeout": settings.api.request_timeout,
+                "default_aspect_ratio": settings.api.default_aspect_ratio,
+                "default_output_format": settings.api.default_output_format,
+            }
+
+            return json.dumps(config, indent=2)
+
+        logger.info("Ultimate Gemini MCP Server initialized successfully")
+        return mcp
+
+    except Exception as e:
+        logger.error(f"Failed to initialize server: {e}", exc_info=True)
+        raise
+
+
+def main() -> None:
+    """Main entry point for direct execution."""
+    try:
+        logger.info("Starting Ultimate Gemini MCP Server...")
+
+        # Create application
+        app = create_app()
+
+        # Run the server (FastMCP handles stdio transport)
+        logger.info("Server is ready and listening for MCP requests")
+        app.run()
+
+    except KeyboardInterrupt:
+        logger.info("Server shutdown requested by user")
+        sys.exit(0)
+    except Exception as e:
+        logger.error(f"Server error: {e}", exc_info=True)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

src/services/__init__.py

@@ -0,0 +1,15 @@
+"""Services module for Ultimate Gemini MCP."""
+
+from .gemini_client import GeminiClient
+from .imagen_client import ImagenClient
+from .image_service import ImageResult, ImageService
+from .prompt_enhancer import PromptEnhancer, create_prompt_enhancer
+
+__all__ = [
+    "GeminiClient",
+    "ImagenClient",
+    "ImageService",
+    "ImageResult",
+    "PromptEnhancer",
+    "create_prompt_enhancer",
+]

src/services/gemini_client.py

@@ -0,0 +1,230 @@
+"""
+Gemini API client for Gemini 2.5 Flash Image generation.
+Uses the generateContent API endpoint per Google's documentation.
+"""
+
+import base64
+import logging
+from typing import Any
+
+import httpx
+
+from ..config.constants import GEMINI_API_BASE, GEMINI_MODELS
+from ..core.exceptions import (
+    APIError,
+    AuthenticationError,
+    ContentPolicyError,
+    RateLimitError,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class GeminiClient:
+    """Client for Gemini 2.5 Flash Image API."""
+
+    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.base_url = GEMINI_API_BASE
+        self.client = httpx.AsyncClient(timeout=timeout)
+
+    async def generate_image(
+        self,
+        prompt: str,
+        *,
+        model: str = "gemini-2.5-flash-image",
+        input_image: str | None = None,
+        aspect_ratio: str | None = None,
+        **kwargs: Any,
+    ) -> dict[str, Any]:
+        """
+        Generate or edit an image using Gemini 2.5 Flash Image.
+
+        Args:
+            prompt: Text prompt for image generation or editing instruction
+            model: Model to use (default: gemini-2.5-flash-image)
+            input_image: Base64-encoded input image for editing (optional)
+            aspect_ratio: Desired aspect ratio (optional, influences output)
+            **kwargs: Additional parameters
+
+        Returns:
+            Dict with 'images' key containing list of base64-encoded image data
+
+        Raises:
+            APIError: If the API request fails
+        """
+        model_id = GEMINI_MODELS.get(model, model)
+        url = f"{self.base_url}/models/{model_id}:generateContent"
+
+        # Build request body according to doc.md
+        parts: list[dict[str, Any]] = []
+
+        # Add input image if provided (for editing)
+        if input_image:
+            parts.append({
+                "inline_data": {
+                    "mime_type": "image/png",
+                    "data": input_image
+                }
+            })
+
+        # Add text prompt (include aspect ratio hint if specified)
+        prompt_text = prompt
+        if aspect_ratio:
+            prompt_text = f"{prompt}. Aspect ratio: {aspect_ratio}"
+
+        parts.append({"text": prompt_text})
+
+        request_body = {"contents": [{"parts": parts}]}
+
+        headers = {
+            "x-goog-api-key": self.api_key,
+            "Content-Type": "application/json",
+        }
+
+        try:
+            logger.debug(f"Sending request to {url}")
+            response = await self.client.post(url, json=request_body, headers=headers)
+            response.raise_for_status()
+            data = response.json()
+
+            # Extract images from response
+            images = self._extract_images(data)
+
+            if not images:
+                raise APIError("No image data found in Gemini 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"Gemini API request failed: {e}")
+            raise APIError(f"Gemini API request failed: {e}")
+
+    async def generate_text(
+        self,
+        prompt: str,
+        *,
+        model: str = "gemini-2.0-flash",
+        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)
+        url = f"{self.base_url}/models/{model_id}:generateContent"
+
+        request_body = {
+            "contents": [{"parts": [{"text": prompt}]}]
+        }
+
+        if system_instruction:
+            request_body["system_instruction"] = {
+                "parts": [{"text": system_instruction}]
+            }
+
+        headers = {
+            "x-goog-api-key": self.api_key,
+            "Content-Type": "application/json",
+        }
+
+        try:
+            response = await self.client.post(url, json=request_body, headers=headers)
+            response.raise_for_status()
+            data = response.json()
+
+            # Extract text from response
+            text = self._extract_text(data)
+            return text
+
+        except httpx.HTTPStatusError as e:
+            self._handle_http_error(e)
+        except Exception as e:
+            logger.error(f"Gemini text generation failed: {e}")
+            raise APIError(f"Gemini text generation failed: {e}")
+
+    def _extract_images(self, response_data: dict[str, Any]) -> list[str]:
+        """Extract base64 image data from Gemini API response."""
+        images = []
+
+        try:
+            candidates = response_data.get("candidates", [])
+            for candidate in candidates:
+                content = candidate.get("content", {})
+                parts = content.get("parts", [])
+                for part in parts:
+                    if "inline_data" in part:
+                        image_data = part["inline_data"].get("data")
+                        if image_data:
+                            images.append(image_data)
+        except Exception as e:
+            logger.warning(f"Error extracting images from response: {e}")
+
+        return images
+
+    def _extract_text(self, response_data: dict[str, Any]) -> str:
+        """Extract text from Gemini API response."""
+        try:
+            candidates = response_data.get("candidates", [])
+            if not candidates:
+                return ""
+
+            content = candidates[0].get("content", {})
+            parts = content.get("parts", [])
+
+            # Concatenate all text parts
+            text_parts = [part.get("text", "") for part in parts if "text" in part]
+            return "".join(text_parts)
+
+        except Exception as e:
+            logger.warning(f"Error extracting text from response: {e}")
+            return ""
+
+    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 Gemini 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() or "BLOCKED" 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/image_service.py

@@ -0,0 +1,243 @@
+"""
+Unified image service that orchestrates Gemini and Imagen APIs.
+Provides a consistent interface for image generation regardless of the underlying model.
+"""
+
+import base64
+import logging
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+from PIL import Image
+
+from ..config.constants import GEMINI_MODELS, IMAGEN_MODELS
+from ..core import sanitize_filename
+from ..core.exceptions import ImageProcessingError
+from .gemini_client import GeminiClient
+from .imagen_client import ImagenClient
+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}")
+
+    def _generate_filename(self) -> str:
+        """Generate descriptive filename."""
+        timestamp = self.timestamp.strftime("%Y%m%d_%H%M%S")
+        # Sanitize prompt for filename
+        prompt_snippet = sanitize_filename(self.prompt[:50])
+        index_str = f"_{self.index + 1}" if self.index > 0 else ""
+        return f"{self.model}_{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:
+    """Unified service for image generation using Gemini or Imagen."""
+
+    def __init__(
+        self,
+        api_key: str,
+        *,
+        enable_enhancement: bool = True,
+        timeout: int = 60
+    ):
+        """
+        Initialize image service.
+
+        Args:
+            api_key: API key for Google AI services
+            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 clients
+        self.gemini_client = GeminiClient(api_key, timeout)
+        self.imagen_client = ImagenClient(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 the appropriate API.
+
+        Args:
+            prompt: Text prompt for image generation
+            model: Model to use (auto-detected if None)
+            enhance_prompt: Whether to enhance the prompt
+            **kwargs: Additional parameters (aspect_ratio, number_of_images, etc.)
+
+        Returns:
+            List of ImageResult objects
+        """
+        # Detect which API to use based on model
+        if model is None:
+            model = "gemini-2.5-flash-image"  # Default to Gemini
+
+        is_gemini = model in GEMINI_MODELS
+        is_imagen = model in IMAGEN_MODELS
+
+        if not is_gemini and not is_imagen:
+            raise ValueError(f"Unknown model: {model}")
+
+        # 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 appropriate API
+        if is_gemini:
+            return await self._generate_with_gemini(prompt, model, original_prompt, kwargs)
+        else:
+            return await self._generate_with_imagen(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
+
+    async def _generate_with_imagen(
+        self,
+        prompt: str,
+        model: str,
+        original_prompt: str,
+        params: dict[str, Any]
+    ) -> list[ImageResult]:
+        """Generate images using Imagen API."""
+        response = await self.imagen_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": "imagen",
+                    **params
+                }
+            )
+            results.append(result)
+
+        return results
+
+    def _build_enhancement_context(self, params: dict[str, Any]) -> dict[str, Any]:
+        """Build context for prompt enhancement."""
+        context = {}
+
+        if "input_image" in params:
+            context["is_editing"] = True
+
+        if params.get("maintainCharacterConsistency"):
+            context["maintain_character_consistency"] = True
+
+        if params.get("blendImages"):
+            context["blend_images"] = True
+
+        if params.get("useWorldKnowledge"):
+            context["use_world_knowledge"] = True
+
+        if "aspect_ratio" in params:
+            context["aspect_ratio"] = params["aspect_ratio"]
+
+        return context
+
+    async def close(self) -> None:
+        """Close all clients."""
+        await self.gemini_client.close()
+        await self.imagen_client.close()