examples/ai_testing/test_json_format.py

"""
Test JSON Prompt Format
-----------------------
Quick test to verify the new JSON format matches promt_format.text structure.
"""

import sys
import json
from pathlib import Path

# Add parent directory to path
sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from pycatan.ai.prompt_templates import PromptBuilder, get_response_schema
from pycatan.ai.config import AIConfig


def test_prompt_structure():
    """Test that generated prompts match expected structure."""
    print("\n" + "="*80)
    print("🧪 TESTING JSON PROMPT FORMAT")
    print("="*80 + "\n")
    
    # Create sample optimized game state
    sample_state = {
        "H": ["", "W12", "S5", "W4", "S8", "B6"],
        "N": [
            [],  # 0 - empty
            [[2, 5], [1, 2], None],  # Node 1
            [[1, 3], [1, 3], "W2"],  # Node 2 with wood port
        ],
        "players": {
            "a": {
                "vp": 2,
                "res": {"W": 3, "B": 1},
                "dev": {"h": ["K"], "r": []},
                "stat": []
            },
            "b": {
                "vp": 1,
                "res": {"S": 2},
                "dev": {"h": [], "r": []},
                "stat": []
            }
        },
        "bld": [
            [20, "a", "S"],
            [25, "b", "S"]
        ],
        "rds": [
            [[20, 21], "a"],
            [[25, 26], "b"]
        ],
        "meta": {
            "robber": 10,
            "phase": "MAIN_GAME",
            "curr": "a",
            "dice": 8
        }
    }
    
    # Build prompt
    builder = PromptBuilder()
    
    prompt = builder.build_prompt(
        meta_data={
            "agent_name": "a",
            "my_color": "Red",
            "role": "You are player 'a' (Red). Play strategically to win."
        },
        task_context={
            "what_just_happened": "You rolled an 8. You received 2 wood.",
            "instructions": "Choose your action from 'allowed_actions'."
        },
        game_state=sample_state,
        social_context={
            "recent_chat": [
                {"sender": "b", "content": "Anyone need sheep?"}
            ]
        },
        memory=[
            "b needs ore, I should trade",
            "Focus on longest road"
        ],
        constraints={
            "usage_instructions": "Choose one action. Use exact parameter structure.",
            "allowed_actions": [
                {
                    "action": "build_road",
                    "description": "Build a road on an edge",
                    "example_parameters": {"edge_id": 15}
                },
                {
                    "action": "end_turn",
                    "description": "End your turn",
                    "example_parameters": {}
                }
            ]
        }
    )
    
    # Verify structure matches promt_format.text
    print("✅ Testing prompt structure...\n")
    
    required_sections = ["meta_data", "task_context", "game_state", 
                        "social_context", "memory", "constraints"]
    
    for section in required_sections:
        if section in prompt:
            print(f"  ✓ {section}")
        else:
            print(f"  ❌ Missing: {section}")
    
    print("\n📋 Meta Data:")
    print(f"  • agent_name: {prompt['meta_data'].get('agent_name')}")
    print(f"  • my_color: {prompt['meta_data'].get('my_color')}")
    print(f"  • role: {prompt['meta_data'].get('role')[:50]}...")
    
    print("\n📋 Task Context:")
    print(f"  • what_just_happened: {prompt['task_context'].get('what_just_happened')}")
    
    print("\n📋 Game State:")
    game_state = prompt.get('game_state', '')
    if isinstance(game_state, str):
        lines = game_state.split('\n')
        print(f"  • Type: String with embedded JSON")
        print(f"  • Lines: {len(lines)}")
        print(f"  • Has legend: {'FORMAT GUIDE' in game_state}")
        print(f"  • Has H array: {'\"H\":' in game_state}")
        print(f"  • First 100 chars: {game_state[:100]}...")
    
    print("\n📋 Social Context:")
    print(f"  • recent_chat: {len(prompt['social_context'].get('recent_chat', []))} messages")
    
    print("\n📋 Memory:")
    print(f"  • notes_for_myself: {len(prompt['memory'].get('notes_for_myself', []))} notes")
    
    print("\n📋 Constraints:")
    print(f"  • allowed_actions: {len(prompt['constraints'].get('allowed_actions', []))} actions")
    
    # Test response schema
    print("\n" + "="*80)
    print("📝 Response Schema:")
    print("="*80)
    schema = get_response_schema()
    print(json.dumps(schema, indent=2))
    
    # Create full LLM request
    print("\n" + "="*80)
    print("📤 Complete LLM Request Structure:")
    print("="*80)
    
    llm_request = {
        "response_schema": schema,
        "system_instruction": "You are an expert Settlers of Catan player. Respond in JSON format.",
        "prompt": prompt
    }
    
    # Save to file
    output_file = Path('examples/ai_testing/my_games/test_prompt_format.json')
    output_file.parent.mkdir(exist_ok=True, parents=True)
    
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(llm_request, f, indent=2, ensure_ascii=False)
    
    print(f"\n✅ Saved complete request to: {output_file}")
    print(f"📏 File size: {output_file.stat().st_size:,} bytes")
    
    # Estimate tokens (rough)
    json_str = json.dumps(llm_request, ensure_ascii=False)
    estimated_tokens = len(json_str) // 4
    print(f"📊 Estimated tokens: ~{estimated_tokens:,}")
    
    print("\n" + "="*80)
    print("✅ ALL TESTS PASSED!")
    print("="*80)
    print("\n💡 The prompt format matches promt_format.text structure:")
    print("   • meta_data ✓")
    print("   • task_context ✓")
    print("   • game_state (with legend) ✓")
    print("   • social_context ✓")
    print("   • memory ✓")
    print("   • constraints ✓")
    print("\n💡 Complete LLM request includes:")
    print("   • response_schema (tells LLM how to respond)")
    print("   • system_instruction")
    print("   • prompt (the actual game state)")
    print("\n" + "="*80 + "\n")


if __name__ == '__main__':
    test_prompt_structure()