58 lines
2.0 KiB
Python
58 lines
2.0 KiB
Python
"""
|
|
MCP Server implementation for Imagen4
|
|
"""
|
|
|
|
import logging
|
|
from typing import Dict, Any, List
|
|
|
|
from mcp.server import Server
|
|
from mcp.types import Tool, TextContent, ImageContent
|
|
|
|
from src.connector import Config
|
|
from src.server.models import MCPToolDefinitions
|
|
from src.server.handlers import ToolHandlers, sanitize_args_for_logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class Imagen4MCPServer:
|
|
"""Imagen4 MCP server class"""
|
|
|
|
def __init__(self, config: Config):
|
|
"""Initialize server"""
|
|
self.config = config
|
|
self.server = Server("imagen4-mcp-server")
|
|
self.handlers = ToolHandlers(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")
|
|
return MCPToolDefinitions.get_all_tools()
|
|
|
|
@self.server.call_tool()
|
|
async def handle_call_tool(name: str, arguments: Dict[str, Any]) -> List[TextContent | ImageContent]:
|
|
"""Handle tool calls"""
|
|
# Log tool call safely without exposing sensitive data
|
|
safe_args = sanitize_args_for_logging(arguments)
|
|
logger.info(f"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
|