astro-mcp/render_test_charts.py

#!/usr/bin/env python3
"""Test script for render_natal_chart_by_id via MCP client.

Calls the astro-mcp server via MCP SSE to render test charts for
celebrities stored in the database.

Usage:
    python3 render_test_charts.py [--format svg|png|jpg] [--size N]

Requires:
    - astro-mcp server running (e.g., on thinkcenter-2:7016)
"""

from __future__ import annotations

import argparse
import asyncio
import json
import os
from datetime import timedelta
from typing import Any

from mcp import ClientSession
from mcp.client.sse import sse_client

ASTRO_MCP_URL = os.environ.get(
    "ASTRO_MCP_URL", "http://192.168.0.249:7016/mcp/sse"
)

# Person nicknames in the database (as used in test_live_charts.py)
PERSON_IDS = [
    "trump",   # Donald Trump
    "chaka",   # Chaka Khan
    "lucky"
]


def _payload_from_result(result: Any) -> dict[str, Any]:
    """Extract a dict payload from an MCP CallToolResult."""
    payload = getattr(result, "structuredContent", None)
    if isinstance(payload, dict):
        return payload

    content_items = getattr(result, "content", []) or []
    for item in content_items:
        text = getattr(item, "text", None)
        if not isinstance(text, str) or not text.strip():
            continue
        try:
            decoded = json.loads(text)
        except Exception:
            continue
        if isinstance(decoded, dict):
            return decoded

    error_texts = []
    for item in content_items:
        text = getattr(item, "text", None)
        if isinstance(text, str) and text.strip():
            error_texts.append(text.strip())
    if error_texts:
        return {"error": error_texts[0]}
    return {}


async def call_astro_tool(
    session: ClientSession, tool_name: str, arguments: dict[str, Any]
) -> dict[str, Any]:
    """Call a tool on the astro-mcp server via MCP session."""
    result = await session.call_tool(tool_name, arguments)
    return _payload_from_result(result)


async def render_person_chart(person_id: str, fmt: str, size: int, color_mode: str, timeout: float = 30.0) -> dict[str, Any] | None:
    """Render a natal chart for one person by ID via MCP."""
    url = ASTRO_MCP_URL
    if not url.endswith("/mcp/sse"):
        url = url.rstrip("/") + "/mcp/sse"

    async with sse_client(url, timeout=timeout, sse_read_timeout=timeout) as streams:
        async with ClientSession(
            *streams, read_timeout_seconds=timedelta(seconds=timeout)
        ) as session:
            await session.initialize()
            
            # Get chart data first (for moon info)
            chart_args = {"person_id": person_id}
            chart_result = await call_astro_tool(session, "calculate_natal_chart_by_id", chart_args)
            if "error" in chart_result:
                print(f"  ERROR calculating {person_id}: {chart_result['error']}")
                return None
            
            # Then render
            render_args = {"person_id": person_id, "color_mode": color_mode, "size": size, "format": fmt}
            render_result = await call_astro_tool(session, "render_natal_chart_by_id", render_args)
            if "error" in render_result:
                print(f"  ERROR rendering {person_id}: {render_result['error']}")
                return None
            
            # Save to output directory
            content = render_result.get("content", "")
            ext = fmt
            filename = f"{person_id}_natal.{ext}"
            outdir = "/home/shared/astro"
            os.makedirs(outdir, exist_ok=True)
            outpath = os.path.join(outdir, filename)

            mode = "w" if isinstance(content, str) else "wb"
            with open(outpath, mode) as f:
                f.write(content)

            # Extract moon info from chart data
            moon_info = None
            for p in chart_result.get("planets", []):
                if p.get("body") == "moon":
                    moon_info = {
                        "absolute_lon": p.get("absolute_lon"),
                        "degree_within_sign": p.get("degree_within_sign"),
                        "sign": p.get("sign"),
                        "house": p.get("house"),
                        "retrograde": p.get("retrograde"),
                    }
                    break

            return {
                "person_id": person_id,
                "output": outpath,
                "format": render_result.get("format"),
                "content_type": render_result.get("content_type"),
                "size_bytes": len(content) if isinstance(content, bytes) else len(content.encode("utf-8")),
                "moon": moon_info,
                "all_planets": [
                    {"body": p.get("body"), "absolute_lon": p.get("absolute_lon"),
                     "degree_within_sign": p.get("degree_within_sign"),
                     "sign": p.get("sign"), "house": p.get("house"),
                     "retrograde": p.get("retrograde")}
                    for p in chart_result.get("planets", [])
                ],
                "houses_count": len(chart_result.get("houses", [])),
                "aspects_count": len(chart_result.get("aspects", [])),
            }


async def main():
    parser = argparse.ArgumentParser(description="Render test natal charts via MCP client")
    parser.add_argument("--format", default="svg", choices=["svg", "png", "jpg"],
                        help="Output format (default: svg)")
    parser.add_argument("--size", type=int, default=600,
                        help="Canvas size in pixels (default: 600)")
    parser.add_argument("--color-mode", default="color", choices=["color", "bw", "dark"],
                        help="Color theme (default: color)")
    args = parser.parse_args()

    print(f"Rendering test charts — format={args.format}, size={args.size}px, theme={args.color_mode}")
    print(f"Output directory: /home/shared/astro/")
    print(f"Persons: {', '.join(PERSON_IDS)}")
    print()

    results = []
    for person_id in PERSON_IDS:
        print(f"Processing {person_id}...")
        result = await render_person_chart(person_id, args.format, args.size, args.color_mode)
        if result:
            results.append(result)
            print(f"  Saved: {result['output']} ({result['size_bytes']:,} bytes)")
            if result["moon"]:
                m = result["moon"]
                print(f"  MOON: {m['sign']} {m['degree_within_sign']:.1f}° (lon={m['absolute_lon']:.2f}°, house {m['house']})")
            else:
                print(f"  MOON: not found in data!")
            print(f"  Planets: {len(result['all_planets'])}, Houses: {result['houses_count']}, Aspects: {result['aspects_count']}")
        else:
            print(f"  FAILED")
        print()

    # Summary comparison
    print("=" * 60)
    print("MOON COMPARISON")
    print("=" * 60)
    for r in results:
        m = r.get("moon", {})
        if m:
            print(f"  {r['person_id']:20s} | Moon in {m.get('sign','?'):12s} {m.get('degree_within_sign',0):5.1f}° | House {m.get('house','?')} | Retrograde: {m.get('retrograde',False)}")
        else:
            print(f"  {r['person_id']:20s} | Moon data MISSING")
    print()

    # Save JSON summary
    summary_path = "/home/shared/astro/test_charts_summary.json"
    with open(summary_path, "w") as f:
        json.dump(results, f, indent=2, default=str)
    print(f"Summary saved to: {summary_path}")


if __name__ == "__main__":
    asyncio.run(main())