"""Midjourney service layer for business logic."""
import logging
from functools import wraps
from typing import List, Optional, Callable
from client import GPTNBClient
from task_handler import TaskManager
from models import (
ImagineRequest, BlendRequest, DescribeRequest, ChangeRequest,
SwapFaceRequest, ModalRequest, Dimensions
)
from config import Config, get_config
from utils import (
validate_prompt, validate_aspect_ratio, validate_base64_images,
validate_task_id, validate_image_index, format_error_message
)
from exceptions import ValidationError
logger = logging.getLogger(__name__)
def handle_service_errors(operation_name: str):
"""Decorator to handle common service errors.
Args:
operation_name: Name of the operation for error messages
"""
def decorator(func: Callable) -> Callable:
"""Decorator function."""
@wraps(func)
async def wrapper(*args, **kwargs) -> str:
"""Wrapper function with error handling."""
try:
return await func(*args, **kwargs)
except Exception as e:
logger.error(f"Error in {operation_name}: {e}")
return format_error_message(e, operation_name)
return wrapper
return decorator
class MidjourneyService:
"""High-level service for Midjourney operations."""
def __init__(self, config: Optional[Config] = None):
"""Initialize Midjourney service.
Args:
config: Configuration object (uses global config if None)
"""
self.config = config or get_config()
self.client: Optional[GPTNBClient] = None
self.task_manager: Optional[TaskManager] = None
async def __aenter__(self):
"""Async context manager entry."""
self.client = GPTNBClient(self.config)
await self.client.__aenter__()
self.task_manager = TaskManager(self.client, self.config)
return self
async def __aexit__(self, exc_type, exc_val, exc_tb):
"""Async context manager exit."""
if self.client:
await self.client.__aexit__(exc_type, exc_val, exc_tb)
async def imagine(
self,
prompt: str,
aspect_ratio: str = "1:1",
base64_images: Optional[List[str]] = None
) -> str:
"""Generate images from text prompt.
Args:
prompt: Text description of the image
aspect_ratio: Aspect ratio (e.g., "16:9", "1:1", "9:16")
base64_images: Optional reference images
Returns:
Formatted result string
Raises:
ValidationError: If input validation fails
MidjourneyMCPError: If operation fails
"""
try:
# Validate inputs
prompt = validate_prompt(prompt)
if not validate_aspect_ratio(aspect_ratio):
raise ValidationError(f"Invalid aspect ratio: {aspect_ratio}")
# Add aspect ratio to prompt if not already present
if "--ar" not in prompt:
prompt = f"{prompt} --ar {aspect_ratio}"
# Add default suffix if configured
if self.config.default_suffix and self.config.default_suffix not in prompt:
prompt = f"{prompt} {self.config.default_suffix}"
# Validate and format images if provided
validated_images = None
if base64_images:
validated_images = validate_base64_images(base64_images, min_count=0, max_count=5)
# Create request
request = ImagineRequest(
prompt=prompt,
base64Array=validated_images
)
# Submit task (don't wait for completion)
response = await self.client.submit_imagine(request)
if response.code != 1:
raise Exception(f"Task submission failed: {response.description}")
if not response.result:
raise Exception("No task ID returned from submission")
task_id = response.result
# Return immediate response with task ID
result = f"🎨 **Image Generation Started!**\n\n"
result += f"**Task ID:** {task_id}\n"
result += f"**Prompt:** {prompt}\n"
result += f"**Aspect Ratio:** {aspect_ratio}\n"
result += f"**Status:** Task submitted successfully\n\n"
result += f"💡 **Next Steps:**\n"
result += f"Use `get_task_status(task_id=\"{task_id}\")` to check current status\n\n"
result += f"⏱️ **Estimated Time:** 30-60 seconds\n"
return result
except Exception as e:
logger.error(f"Error in imagine: {e}")
return format_error_message(e, "imagine")
async def blend(
self,
base64_images: List[str],
dimensions: str = "SQUARE"
) -> str:
"""Blend multiple images together.
Args:
base64_images: List of 2-5 images to blend
dimensions: Output dimensions ("PORTRAIT", "SQUARE", "LANDSCAPE")
Returns:
Formatted result string
Raises:
ValidationError: If input validation fails
MidjourneyMCPError: If operation fails
"""
try:
# Validate inputs
validated_images = validate_base64_images(base64_images, min_count=2, max_count=5)
# Validate dimensions
try:
dim_enum = Dimensions(dimensions.upper())
except ValueError:
raise ValidationError(f"Invalid dimensions: {dimensions}. Must be PORTRAIT, SQUARE, or LANDSCAPE")
# Create request
request = BlendRequest(
base64Array=validated_images,
dimensions=dim_enum
)
# Submit task (don't wait for completion)
response = await self.client.submit_blend(request)
if response.code != 1:
raise Exception(f"Task submission failed: {response.description}")
if not response.result:
raise Exception("No task ID returned from submission")
task_id = response.result
# Return immediate response with task ID
result = f"🎨 **Image Blending Started!**\n\n"
result += f"**Task ID:** {task_id}\n"
result += f"**Images:** {len(base64_images)} images to blend\n"
result += f"**Dimensions:** {dimensions}\n"
result += f"**Status:** Task submitted successfully\n\n"
result += f"💡 **Monitor Progress:** Use `get_task_status(task_id=\"{task_id}\")`\n"
return result
except Exception as e:
logger.error(f"Error in blend: {e}")
return format_error_message(e, "blend")
async def describe(self, base64_image: str) -> str:
"""Generate text description of an image.
Args:
base64_image: Image to describe
Returns:
Formatted result string
Raises:
ValidationError: If input validation fails
MidjourneyMCPError: If operation fails
"""
try:
# Validate image
validated_images = validate_base64_images([base64_image], min_count=1, max_count=1)
# Create request
request = DescribeRequest(
base64=validated_images[0]
)
# Submit and wait for completion
task = await self.task_manager.submit_and_wait(
self.client.submit_describe, request
)
return self.task_manager.format_task_result(task)
except Exception as e:
logger.error(f"Error in describe: {e}")
return format_error_message(e, "describe")
async def change(
self,
task_id: str,
action: str,
index: Optional[int] = None
) -> str:
"""Create variations, upscales, or rerolls of existing images.
Args:
task_id: ID of the original generation task
action: Action type ("UPSCALE", "VARIATION", "REROLL")
index: Image index (1-4) for UPSCALE and VARIATION
Returns:
Formatted result string
Raises:
ValidationError: If input validation fails
MidjourneyMCPError: If operation fails
"""
try:
# Validate inputs
task_id = validate_task_id(task_id)
action = action.upper()
if action not in ["UPSCALE", "VARIATION", "REROLL"]:
raise ValidationError(f"Invalid action: {action}. Must be UPSCALE, VARIATION, or REROLL")
if action in ["UPSCALE", "VARIATION"]:
if index is None:
raise ValidationError(f"Index is required for {action} action")
index = validate_image_index(index)
# Create request
request = ChangeRequest(
taskId=task_id,
action=action,
index=index
)
# Submit and wait for completion
task = await self.task_manager.submit_and_wait(
self.client.submit_change, request
)
return self.task_manager.format_task_result(task)
except Exception as e:
logger.error(f"Error in change: {e}")
return format_error_message(e, "change")
async def modal_edit(
self,
task_id: str,
action: str,
prompt: Optional[str] = None
) -> str:
"""Perform advanced editing like zoom, pan, or inpainting.
Args:
task_id: ID of the original generation task
action: Edit action type (zoom, pan, inpaint, etc.)
prompt: Additional prompt for the edit
Returns:
Formatted result string
Raises:
ValidationError: If input validation fails
MidjourneyMCPError: If operation fails
"""
try:
# Validate inputs
task_id = validate_task_id(task_id)
if prompt:
prompt = validate_prompt(prompt)
# Note: This is a simplified implementation
# Real modal operations would need mask images and specific action parameters
logger.info(f"Performing modal edit with action: {action}")
# Create request (Note: This is a simplified version)
# In practice, modal operations require more complex parameters
request = ModalRequest(
taskId=task_id,
maskBase64="", # This would need to be provided by the user
prompt=prompt
)
# Submit and wait for completion
task = await self.task_manager.submit_and_wait(
self.client.submit_modal, request
)
return self.task_manager.format_task_result(task)
except Exception as e:
logger.error(f"Error in modal_edit: {e}")
return format_error_message(e, "modal_edit")
async def swap_face(
self,
source_image: str,
target_image: str
) -> str:
"""Swap faces between two images.
Args:
source_image: Source face image in base64 format
target_image: Target image in base64 format
Returns:
Formatted result string
Raises:
ValidationError: If input validation fails
MidjourneyMCPError: If operation fails
"""
try:
# Validate images
validated_images = validate_base64_images([source_image, target_image], min_count=2, max_count=2)
# Create request
request = SwapFaceRequest(
sourceBase64=validated_images[0],
targetBase64=validated_images[1]
)
# Submit and wait for completion
task = await self.task_manager.submit_and_wait(
self.client.submit_swap_face, request
)
return self.task_manager.format_task_result(task)
except Exception as e:
logger.error(f"Error in swap_face: {e}")
return format_error_message(e, "swap_face")
async def get_task_status(self, task_id: str) -> str:
"""Get current status of a task.
Args:
task_id: Task ID to check
Returns:
Formatted status string
Raises:
ValidationError: If input validation fails
MidjourneyMCPError: If operation fails
"""
try:
# Validate task ID
task_id = validate_task_id(task_id)
# Get task status
task = await self.task_manager.get_task_status(task_id)
return self.task_manager.format_task_result(task)
except Exception as e:
logger.error(f"Error in get_task_status: {e}")
return format_error_message(e, "get_task_status")
# Global service instance
_service_instance: Optional[MidjourneyService] = None
async def get_service() -> MidjourneyService:
"""Get or create the global service instance.
Returns:
MidjourneyService instance
"""
global _service_instance
if _service_instance is None:
_service_instance = MidjourneyService()
await _service_instance.__aenter__()
return _service_instance
async def close_service():
"""Close the global service instance."""
global _service_instance
if _service_instance is not None:
await _service_instance.__aexit__(None, None, None)
_service_instance = None
