import asyncio
from nl2sql_tool.model.protocal import NL2SQLRequest
from nl2sql_tool.tool.nl2sql import NL2SQLAgent

async def main():
    agent = NL2SQLAgent()
    
    request = NL2SQLRequest(
        request_id="test-123",
        query="查询不同国家的销售总额",
        current_date_info="当前时间信息：2025-09-12,星期五",
        table_id_list=["sales"],
        column_info=[
            {
                "modelCode": "sales",
                "modelName": "超市销售明细数据",
                "schemaList": [
                    {
                        "columnId": "COUNTRY",
                        "columnName": "国家",
                        "columnComment": "国家",
                        "dataType": "VARCHAR",
                        "fewShot": "中国,美国,日本",
                    },
                    {
                        "columnId": "SALES",
                        "columnName": "销售额",
                        "columnComment": "销售额",
                        "dataType": "DECIMAL",
                        "fewShot": "1000.00,2000.50",
                    }
                ],
            }
        ],
        dialect="mysql"
    )
    
    result = await agent.run(request)
    print(result)

if __name__ == "__main__":
    asyncio.run(main())