# -*- coding: utf-8 -*-
"""The main entry point of the structured output example."""
import asyncio
import json
import os
from typing import Literal

from agentscope.agent import ReActAgent
from agentscope.formatter import DashScopeChatFormatter
from agentscope.memory import InMemoryMemory
from agentscope.message import Msg
from agentscope.model import DashScopeChatModel
from agentscope.tool import Toolkit
from pydantic import BaseModel, Field


class TableModel(BaseModel):
    """A simple table model for structured output."""

    name: str = Field(description="The name of the person")
    age: int = Field(description="The age of the person", ge=0, le=120)
    intro: str = Field(description="A one-sentence introduction of the person")
    honors: list[str] = Field(
        description="A list of honors received by this person",
    )


class ChoiceModel(BaseModel):
    """A simple choice model for structured output."""

    choice: Literal["apple", "banana", "orange"] = Field(
        description="Your choice of fruit",
    )


async def main() -> None:
    """The main entry point for the structured output example."""
    toolkit = Toolkit()
    agent = ReActAgent(
        name="Friday",
        sys_prompt="You are a helpful assistant named Friday.",
        model=DashScopeChatModel(
            api_key=os.environ.get("DASHSCOPE_API_KEY"),
            model_name="qwen-max",
            stream=True,
        ),
        formatter=DashScopeChatFormatter(),
        toolkit=toolkit,
        memory=InMemoryMemory(),
    )

    query_msg_1 = Msg(
        "user",
        "Please introduce Einstein",
        "user",
    )
    res = await agent(query_msg_1, structured_model=TableModel)
    print(
        "Structured Output 1:\n"
        "```\n"
        f"{json.dumps(res.metadata, indent=4)}\n"
        "```",
    )

    query_msg_2 = Msg(
        "user",
        "Choose one of your favorite fruit",
        "user",
    )
    res = await agent(query_msg_2, structured_model=ChoiceModel)
    print(
        "Structured Output 2:\n"
        "```\n"
        f"{json.dumps(res.metadata, indent=4)}\n"
        "```",
    )


asyncio.run(main())