64 lines
2.3 KiB
Python
64 lines
2.3 KiB
Python
"""
|
|
Enhanced MCP Server implementation for Imagen4 with Preview Image Support
|
|
"""
|
|
|
|
import logging
|
|
from typing import Dict, Any, List
|
|
|
|
from mcp.server import Server
|
|
from mcp.types import Tool, TextContent
|
|
|
|
from src.connector import Config
|
|
from src.server.models import MCPToolDefinitions
|
|
from src.server.enhanced_handlers import EnhancedToolHandlers, sanitize_args_for_logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class EnhancedImagen4MCPServer:
|
|
"""Enhanced Imagen4 MCP server class with preview image support"""
|
|
|
|
def __init__(self, config: Config):
|
|
"""Initialize server"""
|
|
self.config = config
|
|
self.server = Server("imagen4-enhanced-mcp-server")
|
|
self.handlers = EnhancedToolHandlers(config)
|
|
|
|
# Register handlers
|
|
self._register_handlers()
|
|
|
|
def _register_handlers(self) -> None:
|
|
"""Register MCP handlers"""
|
|
|
|
@self.server.list_tools()
|
|
async def handle_list_tools() -> List[Tool]:
|
|
"""Return list of available tools"""
|
|
logger.info("Listing available tools (enhanced version)")
|
|
return MCPToolDefinitions.get_all_tools()
|
|
|
|
@self.server.call_tool()
|
|
async def handle_call_tool(name: str, arguments: Dict[str, Any]) -> List[TextContent]:
|
|
"""Handle tool calls with enhanced preview image support"""
|
|
# Log tool call safely without exposing sensitive data
|
|
safe_args = sanitize_args_for_logging(arguments)
|
|
logger.info(f"Enhanced tool called: {name} with arguments: {safe_args}")
|
|
|
|
if name == "generate_random_seed":
|
|
return await self.handlers.handle_generate_random_seed(arguments)
|
|
elif name == "regenerate_from_json":
|
|
return await self.handlers.handle_regenerate_from_json(arguments)
|
|
elif name == "generate_image":
|
|
return await self.handlers.handle_generate_image(arguments)
|
|
else:
|
|
raise ValueError(f"Unknown tool: {name}")
|
|
|
|
def get_server(self) -> Server:
|
|
"""Return MCP server instance"""
|
|
return self.server
|
|
|
|
|
|
# Create a factory function for easier import
|
|
def create_enhanced_server(config: Config) -> EnhancedImagen4MCPServer:
|
|
"""Factory function to create enhanced server"""
|
|
return EnhancedImagen4MCPServer(config)
|