esotericelf/image_edit_creation

1"""2Lightweight fal.ai connection test suite for the Nano Banana 2 backend.3 4Run locally (host API expected on port 8001):5    python test_connection.py6 7Health check from the Windows host:8    curl http://localhost:8001/api/v1/health9 10Run inside the running container (API listens on container port 8000):11    docker compose exec api python test_connection.py12 13Run as a one-off test container:14    docker compose --profile test run --rm connection-test15"""16 17from __future__ import annotations18 19import json20import os21import sys22import urllib.error23import urllib.request24from dataclasses import dataclass25from pathlib import Path26from typing import Any27 28try:29    import fal_client.auth as fal_auth30except ImportError:  # pragma: no cover - exercised when fal-client is missing31    fal_auth = None32 33TARGET_MODEL_ENDPOINT = "fal-ai/nano-banana-2/edit"34FAL_PLATFORM_MODELS_URL = "https://api.fal.ai/v1/models?limit=1"35DEFAULT_TIMEOUT_SECONDS = 1536HOST_API_PORT = 800137CONTAINER_API_PORT = 800038 39 40@dataclass(frozen=True)41class CheckResult:42    name: str43    passed: bool44    message: str45    details: dict[str, Any] | None = None46 47 48class Console:49    GREEN = "\033[92m"50    RED = "\033[91m"51    YELLOW = "\033[93m"52    CYAN = "\033[96m"53    BOLD = "\033[1m"54    RESET = "\033[0m"55 56    def __init__(self) -> None:57        self._color_enabled = self._supports_color()58 59    @staticmethod60    def _supports_color() -> bool:61        if os.getenv("NO_COLOR"):62            return False63        return hasattr(sys.stdout, "isatty") and bool(sys.stdout.isatty())64 65    def _paint(self, text: str, color: str) -> str:66        if not self._color_enabled:67            return text68        return f"{color}{text}{self.RESET}"69 70    def success(self, message: str) -> None:71        print(self._paint(f"[PASS] {message}", self.GREEN))72 73    def failure(self, message: str) -> None:74        print(self._paint(f"[FAIL] {message}", self.RED))75 76    def warning(self, message: str) -> None:77        print(self._paint(f"[WARN] {message}", self.YELLOW))78 79    def info(self, message: str) -> None:80        print(self._paint(f"[INFO] {message}", self.CYAN))81 82    def heading(self, message: str) -> None:83        print(self._paint(message, f"{self.BOLD}{self.CYAN}"))84 85 86console = Console()87 88 89def check_fal_key_configured() -> CheckResult:90    fal_key = os.getenv("FAL_KEY", "").strip()91    if not fal_key:92        return CheckResult(93            name="fal_key",94            passed=False,95            message=(96                "FAL_KEY is missing or empty. Set it in your shell, .env file, "97                "or docker-compose environment before running the API."98            ),99        )100 101    return CheckResult(102        name="fal_key",103        passed=True,104        message="FAL_KEY is configured.",105        details={"length": len(fal_key)},106    )107 108 109def check_fal_client_credentials() -> CheckResult:110    if fal_auth is None:111        return CheckResult(112            name="fal_client",113            passed=False,114            message="fal-client is not installed. Run: pip install -r requirements.txt",115        )116 117    try:118        credentials = fal_auth.fetch_auth_credentials()119    except fal_auth.MissingCredentialsError:120        return CheckResult(121            name="fal_client",122            passed=False,123            message="fal-client could not resolve credentials from the environment.",124        )125    except Exception as exc:126        return CheckResult(127            name="fal_client",128            passed=False,129            message=f"fal-client credential resolution failed: {exc}",130        )131 132    return CheckResult(133        name="fal_client",134        passed=True,135        message="fal-client successfully resolved API credentials.",136        details={"scheme": credentials.scheme},137    )138 139 140def probe_fal_platform_api(timeout_seconds: int = DEFAULT_TIMEOUT_SECONDS) -> CheckResult:141    fal_key = os.getenv("FAL_KEY", "").strip()142    if not fal_key:143        return CheckResult(144            name="fal_upstream",145            passed=False,146            message="Skipped upstream probe because FAL_KEY is not configured.",147        )148 149    request = urllib.request.Request(150        FAL_PLATFORM_MODELS_URL,151        headers={152            "Authorization": f"Key {fal_key}",153            "Accept": "application/json",154            "User-Agent": "nano-banana-2-connection-test/1.0",155        },156        method="GET",157    )158 159    try:160        with urllib.request.urlopen(request, timeout=timeout_seconds) as response:161            status_code = response.getcode()162            body = response.read().decode("utf-8", errors="replace")163    except urllib.error.HTTPError as exc:164        error_body = exc.read().decode("utf-8", errors="replace")165        if exc.code == 401:166            return CheckResult(167                name="fal_upstream",168                passed=False,169                message=(170                    "fal.ai rejected the API key (HTTP 401 Unauthorized). "171                    "Verify FAL_KEY and ensure the header format is 'Authorization: Key <key>'."172                ),173                details={"http_status": exc.code, "response": _safe_response_preview(error_body)},174            )175        if exc.code == 403:176            return CheckResult(177                name="fal_upstream",178                passed=False,179                message=(180                    "fal.ai denied access (HTTP 403 Forbidden). "181                    "Your key may lack the required API scope."182                ),183                details={"http_status": exc.code, "response": _safe_response_preview(error_body)},184            )185        return CheckResult(186            name="fal_upstream",187            passed=False,188            message=f"fal.ai platform API returned HTTP {exc.code}.",189            details={"http_status": exc.code, "response": _safe_response_preview(error_body)},190        )191    except urllib.error.URLError as exc:192        reason = getattr(exc, "reason", exc)193        return CheckResult(194            name="fal_upstream",195            passed=False,196            message=(197                f"Network error reaching fal.ai ({reason}). "198                "Check Docker/WSL2 networking, DNS, firewall, and outbound HTTPS access."199            ),200            details={"error": str(reason)},201        )202    except TimeoutError:203        return CheckResult(204            name="fal_upstream",205            passed=False,206            message=(207                f"fal.ai platform API timed out after {timeout_seconds}s. "208                "The service may be unreachable from this environment."209            ),210        )211    except Exception as exc:212        return CheckResult(213            name="fal_upstream",214            passed=False,215            message=f"Unexpected error during upstream probe: {exc}",216        )217 218    if status_code != 200:219        return CheckResult(220            name="fal_upstream",221            passed=False,222            message=f"fal.ai platform API returned unexpected HTTP {status_code}.",223            details={"http_status": status_code, "response": _safe_response_preview(body)},224        )225 226    return CheckResult(227        name="fal_upstream",228        passed=True,229        message="Successfully reached fal.ai and validated the API key (no inference invoked).",230        details={231            "http_status": status_code,232            "endpoint": FAL_PLATFORM_MODELS_URL,233            "target_model": TARGET_MODEL_ENDPOINT,234            "response": _safe_response_preview(body),235        },236    )237 238 239def _local_api_health_url() -> str:240    explicit = os.getenv("LOCAL_API_HEALTH_URL", "").strip()241    if explicit:242        return explicit243    if Path("/.dockerenv").exists():244        return f"http://localhost:{CONTAINER_API_PORT}/api/v1/health"245    return f"http://localhost:{HOST_API_PORT}/api/v1/health"246 247 248def check_local_api_health(timeout_seconds: int = DEFAULT_TIMEOUT_SECONDS) -> CheckResult:249    health_url = _local_api_health_url()250    request = urllib.request.Request(251        health_url,252        headers={"Accept": "application/json", "User-Agent": "nano-banana-2-connection-test/1.0"},253        method="GET",254    )255 256    try:257        with urllib.request.urlopen(request, timeout=timeout_seconds) as response:258            status_code = response.getcode()259            body = response.read().decode("utf-8", errors="replace")260    except urllib.error.HTTPError as exc:261        error_body = exc.read().decode("utf-8", errors="replace")262        return CheckResult(263            name="local_api_health",264            passed=False,265            message=f"Local API health check failed with HTTP {exc.code} at {health_url}.",266            details={"http_status": exc.code, "response": _safe_response_preview(error_body)},267        )268    except urllib.error.URLError as exc:269        reason = getattr(exc, "reason", exc)270        return CheckResult(271            name="local_api_health",272            passed=False,273            message=(274                f"Could not reach local API at {health_url} ({reason}). "275                f"Start the stack with `docker compose up` (host port {HOST_API_PORT})."276            ),277            details={"error": str(reason), "health_url": health_url},278        )279    except TimeoutError:280        return CheckResult(281            name="local_api_health",282            passed=False,283            message=f"Local API health check timed out after {timeout_seconds}s ({health_url}).",284        )285    except Exception as exc:286        return CheckResult(287            name="local_api_health",288            passed=False,289            message=f"Unexpected error during local API health check: {exc}",290        )291 292    if status_code != 200:293        return CheckResult(294            name="local_api_health",295            passed=False,296            message=f"Local API health check returned HTTP {status_code} at {health_url}.",297            details={"http_status": status_code, "response": _safe_response_preview(body)},298        )299 300    return CheckResult(301        name="local_api_health",302        passed=True,303        message=f"Local API is healthy at {health_url}.",304        details={"http_status": status_code, "response": _safe_response_preview(body)},305    )306 307 308def get_health_snapshot() -> dict[str, Any]:309    fal_key = os.getenv("FAL_KEY", "").strip()310    configured = len(fal_key) > 0311    fal_key_status = "ok" if configured else "missing"312 313    return {314        "status": "ok" if configured else "unavailable",315        "service": "nano-banana-2-api",316        "target_model": TARGET_MODEL_ENDPOINT,317        "checks": {318            "server": {"status": "ok"},319            "fal_key": {320                "status": fal_key_status,321                "message": (322                    "FAL_KEY is configured."323                    if configured324                    else (325                        "FAL_KEY is missing or empty. Set it in your shell, .env file, "326                        "or docker-compose environment before running the API."327                    )328                ),329                "details": {"length": len(fal_key)} if configured else None,330            },331        },332    }333 334 335def resolve_health_status_code(snapshot: dict[str, Any]) -> int:336    if snapshot.get("status") == "ok":337        return 200338    return 503339 340 341def run_connection_suite(timeout_seconds: int = DEFAULT_TIMEOUT_SECONDS) -> int:342    health_url = _local_api_health_url()343    console.heading("Nano Banana 2 - fal.ai Connection Test")344    console.info(f"Target model endpoint: {TARGET_MODEL_ENDPOINT}")345    console.info(f"Upstream probe URL: {FAL_PLATFORM_MODELS_URL}")346    console.info(f"Local health URL: {health_url}")347    print()348 349    checks = [350        check_fal_key_configured(),351        check_fal_client_credentials(),352        probe_fal_platform_api(timeout_seconds=timeout_seconds),353        check_local_api_health(timeout_seconds=timeout_seconds),354    ]355 356    failures = 0357    for index, result in enumerate(checks, start=1):358        console.heading(f"Step {index}/{len(checks)}: {result.name}")359        if result.passed:360            console.success(result.message)361        else:362            console.failure(result.message)363            failures += 1364 365        if result.details:366            for key, value in result.details.items():367                console.info(f"{key}: {value}")368        print()369 370        if not result.passed and result.name == "fal_key":371            console.failure("Aborting remaining checks because FAL_KEY is required.")372            break373 374    if failures:375        console.failure(f"Connection test failed ({failures} check(s) failed).")376        return 1377 378    console.success("All connection checks passed. Backend is ready to call fal.ai.")379    return 0380 381 382def _safe_response_preview(body: str, max_length: int = 300) -> Any:383    try:384        parsed = json.loads(body)385    except json.JSONDecodeError:386        preview = body[:max_length]387        return preview + ("..." if len(body) > max_length else "")388 389    text = json.dumps(parsed)390    if len(text) <= max_length:391        return parsed392    return text[:max_length] + "..."393 394 395def main() -> int:396    timeout = int(os.getenv("FAL_CONNECTION_TIMEOUT", str(DEFAULT_TIMEOUT_SECONDS)))397    return run_connection_suite(timeout_seconds=timeout)398 399 400if __name__ == "__main__":401    raise SystemExit(main())402