mcp_client.py

"""
MCPクライアントのサンプルコード
ClaudeCodeからこのMCPサーバーを利用する際の参考実装
"""

import requests
import base64
from PIL import Image
import io

class ImageGenMCPClient:
    def __init__(self, base_url: str):
        """
        Args:
            base_url: MCPサーバーのベースURL (例: "https://your-space.hf.space")
        """
        self.base_url = base_url.rstrip('/')
    
    def list_tools(self):
        """利用可能なツールのリストを取得"""
        response = requests.post(f"{self.base_url}/mcp/list_tools")
        return response.json()
    
    def generate_image(self, prompt: str) -> tuple[Image.Image, str, str]:
        """
        画像を生成する
        
        Args:
            prompt: 生成したい画像の説明
            
        Returns:
            生成された画像、メッセージ、画像URL
        """
        payload = {
            "name": "generate_image",
            "arguments": {
                "prompt": prompt
            }
        }
        
        response = requests.post(f"{self.base_url}/mcp/call_tool", json=payload)
        result = response.json()
        
        if result.get("success"):
            # 画像URLから画像を取得
            image_url = result["image_url"]
            img_response = requests.get(image_url)
            image = Image.open(io.BytesIO(img_response.content))
            return image, result["message"], image_url
        else:
            return None, result["message"], None
    
    def health_check(self):
        """ヘルスチェック"""
        response = requests.get(f"{self.base_url}/health")
        return response.json()

# 使用例
if __name__ == "__main__":
    # MCPサーバーのURL
    client = ImageGenMCPClient("https://your-space.hf.space")
    
    # ヘルスチェック
    print("Health check:", client.health_check())
    
    # ツールリストの取得
    tools = client.list_tools()
    print("Available tools:", tools)
    
    # 画像生成
    image, message, image_url = client.generate_image("美しい夕日の風景")
    if image:
        print(f"Success: {message}")
        print(f"Image URL: {image_url}")
        image.save("generated_image.png")
    else:
        print(f"Error: {message}")