# run_qwen_mysql_agent.py
import os
from qwen_agent.agents import Assistant
from qwen_agent.utils.output_beautify import typewriter_print

# 1. 选定 LLM（示例用 DashScope 云端，替换为你自己的即可）
llm_cfg = {
    'model': 'qwen3-8b',

    # Use the endpoint provided by Alibaba Model Studio:
    # 'model_type': 'qwen_dashscope',
    # 'api_key': os.getenv('DASHSCOPE_API_KEY'),

    # Use a custom endpoint compatible with OpenAI API:
    'model_server': 'https://dashscope.aliyuncs.com/compatible-mode/v1',  # api_base
    'api_key': 'sk-5d90a63c1e784e8f801dee65add68867',

    # Other parameters:
    # 'generate_cfg': {
    #         # Add: When the response content is `<think>this is the thought</think>this is the answer;
    #         # Do not add: When the response has been separated by reasoning_content and content.
    #         'thought_in_content': True,
    #     },
}

# 2. 描述可用工具 —— 这里挂载刚启动的 MySQL‑MCP Server
tools = [{
  "mcpServers": {
    "mysql": {
      "command": "uv",
      "args": [
        "--directory", 
        "/home/ubuntu/.mcp",
        "run",
        "mysql_mcp_server"
      ],
      "env": {
        "MYSQL_HOST": "localhost",
        "MYSQL_PORT": "23306",
        "MYSQL_USER": "mcpuser",
        "MYSQL_PASSWORD": "StrongPass123!",
        "MYSQL_DATABASE": "magentodb"
      }
    }
  }
}]


# 3. 创建智能体
bot = Assistant(
    llm=llm_cfg,
    function_list=tools,
)

# 4. 运行示例 —— 用自然语言问数据库
messages = [{
    "role": "user",
    "content": "数据库里 catalog_product_entity_varchar 表有多少条记录？同时请展示前 5 行。"
}]

response_plain_text = ''
for responses in bot.run(messages=messages, stream=True):
    # stream=True 将逐步打印 LLM 思考与结果
    response_plain_text = typewriter_print(responses, response_plain_text)
print(responses[-1]["content"])