from dbgpt.core._private.example_base import ExampleSelector, ExampleType

## Two examples are defined by default
EXAMPLES = [
    {
        "messages": [
            {
                "type": "human",
                "data": {"content": "查询用户test1所在的城市", "example": True},
            },
            {
                "type": "ai",
                "data": {
                    "content": """{\n\"thoughts\": \"\
                    直接查询用户表中用户名为'test1'的\记录即可\",\
                    \n\"sql\": \"SELECT city FROM user where user_name='test1'\"}""",
                    "example": True,
                },
            },
        ]
    },
    {
        "messages": [
            {
                "type": "human",
                "data": {"content": "查询成都的用户的订单信息", "example": True},
            },
            {
                "type": "ai",
                "data": {
                    "content": """{\n\"thoughts\":\
                     \"根据订单表的用户名和用户表的用户名关联用户表和订单表，\
                     再通过用户表的城市为'成都'的过滤即可\",\
                     \n\"sql\": \"SELECT b.* FROM user a  LEFT JOIN tran_order b ON \
                     a.user_name=b.user_name  where a.city='成都'\"}""",
                    "example": True,
                },
            },
        ]
    },
]

sql_data_example = ExampleSelector(
    examples_record=EXAMPLES, use_example=True, type=ExampleType.ONE_SHOT.value
)
