server.py

"""
Comfortool MCP Server
Provides tools for creating autism-friendly social stories with cultural adaptation
"""
from fastmcp import FastMCP
from openai import OpenAI
import os
from dotenv import load_dotenv
import base64
import io
from PIL import Image
from pathlib import Path
import json

# Load environment variables
load_dotenv()

# Initialize FastMCP server
mcp = FastMCP("comfortool")

# Constants
TEMPLATES_DIR = Path(__file__).resolve().parent / "mcp_server" / "templates"
GENERATED_IMAGES_DIR = Path(__file__).resolve().parent / "generated_images"
GENERATED_IMAGES_DIR.mkdir(exist_ok=True)

# Vibe descriptions for image generation (exported for use in app.py)
VIBE_DESCRIPTIONS = {
    "Comic": "comic-style illustration, comic strip format with exactly 4 panels, bold outlines, dynamic poses, speech bubbles, vibrant colors. If there is text, translate it to the selected language with perfect ortography",
    "Kawaii": "kawaii-style illustration, pastel colors, cute rounded characters, big eyes.",
    "Pictorial": "pictorial illustration, painted with a brush, gentle and calm atmosphere",
    "Basic B&W": "black and white illustration, simple lines, no other colors allowed, like a coloring book.Absolutely only use black and white.",
    "Cartoon": "cartoon-style illustration, exaggerated features, bright colors, playful mood",
    "Soft Pastel": "Only use soft pastel colours, smooth textures, calming feeling"
}

# Valid OpenAI voices for TTS (exported for use in app.py)
VALID_VOICES = {"nova", "shimmer", "echo", "onyx", "fable", "alloy", "ash", "sage", "coral"}

# Initialize OpenAI client
def get_openai_client():
    api_key = os.environ.get("OPENAI_API_KEY")
    if not api_key:
        raise ValueError("OPENAI_API_KEY not found in environment variables")
    return OpenAI(api_key=api_key)


# Helper functions (not exposed as tools)
def _adapt_story_impl(
    story: str,
    culture: str = "default",
    age: str = "7",
    gender: str = "female",
    vibe: str = "Cartoon",
    comfort_character: str = "Koala"
) -> str:
    """Internal implementation of adapt_story"""
    prompt = (
        f"Adapt the following story for an autistic child of age {age}, gender {gender}, from {culture} culture. "
        f"Include '{comfort_character}' as a supportive friend, use the style '{vibe}', and make the story concrete, supportive, and easy to understand. "
        "Avoid excessive emotion and exclamation marks. "
        "Adapt language according to age, simpler sentences for smaller (ages 2 to 4) kids. "
        "The absolute max length is 640 characters. "
        "Do not mention any info related to format, like pages, panels or scenes. It is just one story. "
        "Do not translate. Return only the adapted story in English."
        f"\n\nStory:\n{story}"
    )
    
    client = get_openai_client()
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}]
    )
    
    return response.choices[0].message.content.strip()


def _translate_story_impl(
    story: str,
    language: str = "en",
    gender: str = "female"
) -> str:
    """Internal implementation of translate_story"""
    if language.lower() == "en":
        return story
    
    prompt = (
        f"Translate the following story to {language} and the corresponding regional or cultural variant if mentioned. "
        f"Use correct grammatical gender and pronouns, adapted to the selected gender of the kid ('{gender}'). "
        "If the story mentions an animal or character, translate its name to the most culturally and linguistically appropriate version for the target language and region. "
        "For example, if the story mentions 'fox', translate it as 'zorro' for Spanish from Spain. "
        "Return only the translated story, using the regional variety requested."
        f"\n\nStory:\n{story}"
    )
    
    client = get_openai_client()
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}]
    )
    
    return response.choices[0].message.content.strip()


@mcp.tool()
def adapt_story(
    story: str,
    culture: str = "default",
    age: str = "7",
    gender: str = "female",
    vibe: str = "Cartoon",
    comfort_character: str = "Koala"
) -> str:
    """
    Adapts a story for an autistic child based on culture, age, gender, and comfort character.
    
    Args:
        story: The original story text to adapt
        culture: Cultural context (e.g., 'Latino', 'Roma', 'Muslim', 'default')
        age: Child's age (e.g., '7')
        gender: Gender identity (e.g., 'boy', 'girl', 'non-binary')
        vibe: Illustration style for the story
        comfort_character: Supportive character/animal (e.g., 'Koala', 'Robot')
    
    Returns:
        The adapted story in English
    """
    return _adapt_story_impl(story, culture, age, gender, vibe, comfort_character)


@mcp.tool()
def translate_story(
    story: str,
    language: str = "en",
    gender: str = "female"
) -> str:
    """
    Translates a story to the selected language with proper grammatical gender.
    
    Args:
        story: The story text to translate
        language: Target language (e.g., 'Spanish from Spain', 'en', 'fr')
        gender: Gender for grammatical agreement (e.g., 'boy', 'girl', 'non-binary')
    
    Returns:
        The translated story
    """
    return _translate_story_impl(story, language, gender)


# Helper function to list available scenarios
@mcp.tool()
def list_scenarios() -> list[str]:
    """
    Lists all available scenario templates.
    
    Returns:
        List of scenario names
    """
    return [p.stem.replace("_", " ") for p in TEMPLATES_DIR.glob("*.json")]