فهرست منبع

added lucky to test renders

Lukas Goldschmidt 1 ماه پیش
والد
کامیت
4632ec1e0b
1فایلهای تغییر یافته به همراه140 افزوده شده و 110 حذف شده
  1. 140 110
      render_test_charts.py

+ 140 - 110
render_test_charts.py

@@ -1,118 +1,148 @@
 #!/usr/bin/env python3
-"""
-Render test charts for two celebrities and save to /home/shared/.
+"""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 sys
 import os
-
-# Ensure src is on the path
-sys.path.insert(0, os.path.join(os.path.dirname(__file__), "src"))
-
-from astro_mcp.tools import calculate_natal_chart, render_natal_chart
-
-
-# ── Birth data ───────────────────────────────────────────────────────
-
-PEOPLE = {
-    "donald_trump": {
-        "birth_datetime": "1946-06-14T11:00:00-04:00",
-        "latitude": 40.70,
-        "longitude": -73.80,
-        "elevation": 10,
-    },
-    "chaka_khan": {
-        "birth_datetime": "1953-03-23T11:00:00-06:00",
-        "latitude": 41.88,
-        "longitude": -87.63,
-        "elevation": 180,
-    },
-}
-
-
-async def render_person(name: str, birth_data: dict, fmt: str, size: int) -> dict:
-    """Calculate and render a natal chart for one person.
-
-    Returns a summary dict with moon info and output path.
-    """
-    # 1. Calculate chart data
-    chart_data = await calculate_natal_chart(
-        birth_datetime=birth_data["birth_datetime"],
-        latitude=birth_data["latitude"],
-        longitude=birth_data["longitude"],
-        elevation=birth_data["elevation"],
-    )
-    if "error" in chart_data:
-        print(f"  ERROR calculating {name}: {chart_data['error']}")
-        return None
-
-    # 2. Render
-    result = await render_natal_chart(
-        birth_datetime=birth_data["birth_datetime"],
-        latitude=birth_data["latitude"],
-        longitude=birth_data["longitude"],
-        elevation=birth_data["elevation"],
-        color_mode="color",
-        size=size,
-        format=fmt,
-        title=f"Natal Chart: {name.replace('_', ' ').title()}",
-    )
-    if "error" in result:
-        print(f"  ERROR rendering {name}: {result['error']}")
-        return None
-
-    # 3. Save
-    ext = fmt
-    filename = f"{name}_natal.{ext}"
-    outdir = "/home/shared"
-    os.makedirs(outdir, exist_ok=True)
-    outpath = os.path.join(outdir, filename)
-
-    content = result["content"]
-    mode = "w" if isinstance(content, str) else "wb"
-    with open(outpath, mode) as f:
-        f.write(content)
-
-    # 4. Extract moon info
-    moon_info = None
-    for p in chart_data.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"),
+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", [])),
             }
-            break
-
-    return {
-        "name": name,
-        "output": outpath,
-        "format": result["format"],
-        "content_type": result["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_data.get("planets", [])
-        ],
-        "houses_count": len(chart_data.get("houses", [])),
-        "aspects_count": len(chart_data.get("aspects", [])),
-    }
 
 
 async def main():
-    parser = argparse.ArgumentParser(description="Render test natal charts")
+    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,
@@ -122,14 +152,14 @@ async def main():
     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/")
+    print(f"Output directory: /home/shared/astro/")
+    print(f"Persons: {', '.join(PERSON_IDS)}")
     print()
 
     results = []
-    for name, birth_data in PEOPLE.items():
-        print(f"Processing {name.replace('_', ' ').title()}...")
-        print(f"  Birth: {birth_data['birth_datetime']} @ {birth_data['latitude']}, {birth_data['longitude']}")
-        result = await render_person(name, birth_data, args.format, args.size)
+    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)")
@@ -150,17 +180,17 @@ async def main():
     for r in results:
         m = r.get("moon", {})
         if m:
-            print(f"  {r['name']:20s} | Moon in {m.get('sign','?'):12s} {m.get('degree_within_sign',0):5.1f}° | House {m.get('house','?')} | Retrograde: {m.get('retrograde',False)}")
+            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['name']:20s} | Moon data MISSING")
+            print(f"  {r['person_id']:20s} | Moon data MISSING")
     print()
 
     # Save JSON summary
-    summary_path = "/home/shared/test_charts_summary.json"
+    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())
+    asyncio.run(main())