`#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Simple tool to test LLM structured output support.
"""

import os
from typing import List, Optional
from pydantic import BaseModel, Field
from langchain_openai import ChatOpenAI
import time

# --- 1. Define the structured output schema ---
class TestProfile(BaseModel):
    """Schema for the extracted user profile."""
    name: str = Field(description="Full name of the user")
    age: int = Field(description="Age of the user")
    email: Optional[str] = Field(default=None, description="Email address")
    interests: List[str] = Field(description="List of interests or hobbies")


# --- 2. Define test cases ---
TEST_PROMPTS = [
    "My name is Alex Johnson, I'm 28 years old. I enjoy hiking and photography.",
    "The user is Taylor Smith, age 35, contact at [email protected]. Likes reading sci-fi and gaming.",
    "Chris Lee. 42. Interests include jazz music and woodworking.",
]


# --- 3. Core test function ---
def test_structured_output(model_name: str, api_key: str, base_url: str):
    """
    Tests a single model's ability to output structured data.
    """
    print(f"\n🔍 Testing model: {model_name}")
    print("-" * 40)

    try:
        # Create the LLM client
        llm = ChatOpenAI(
            model=model_name,
            openai_api_key=api_key,
            openai_api_base=base_url,
            temperature=0.1,
        )

        # Bind the structured output schema
        structured_llm = llm.with_structured_output(TestProfile)

        for i, prompt in enumerate(TEST_PROMPTS, 1):
            print(f"\nPrompt {i}: '{prompt[:50]}...'")

            try:
                start = time.time()
                # Invoke the model
                result: TestProfile = structured_llm.invoke([
                    ("system", "Extract the user profile information from the following text."),
                    ("user", prompt)
                ])
                elapsed = time.time() - start

                print(f"   ✅ SUCCESS ({elapsed:.2f}s)")
                print(f"      -> Name: {result.name}, Age: {result.age}")
                print(f"      -> Email: {result.email}")
                print(f"      -> Interests: {', '.join(result.interests)}")

            except Exception as e:
                print(f"   ❌ FAILED")
                print(f"      -> Error: {type(e).__name__}: {e}")

    except Exception as e:
        print(f"   ⚠️  Failed to initialize model client.")
        print(f"      -> Error: {e}")


# --- 4. Main execution block ---
if __name__ == "__main__":
    # Configuration (Set your own environment variables or replace here)
    API_KEY = os.getenv("OPENAI_API_KEY", "your_api_key_here")
    BASE_URL = os.getenv("OPENAI_BASE_URL", "http://your.api.base.url/v1")
    MODEL_TO_TEST = "glm-4-plus"  # Change this to your target model

    print("🚀 Starting structured output compatibility test...")
    test_structured_output(MODEL_TO_TEST, API_KEY, BASE_URL)
    print("\nTest complete.")`