from qwen_agent.agents import Assistant
from qwen_agent.tools.base import BaseTool, register_tool

# Define a custom tool
@register_tool('my_image_gen')
class MyImageGen(BaseTool):
    description = 'AI painting service that generates images from text descriptions.'
    parameters = [{
        'name': 'prompt',
        'type': 'string',
        'description': 'Detailed description of the desired image content',
        'required': True
    }]

    def call(self, params: str, **kwargs) -> str:
        prompt = json5.loads(params)['prompt']
        # Your image generation logic here
        return json5.dumps({'image_url': f'https://example.com/{prompt}'})

# Configure the LLM
llm_cfg = {
    'model': 'qwen-max-latest',
    'model_type': 'qwen_dashscope',
    'generate_cfg': {'top_p': 0.8}
}

# Create an agent with tools
bot = Assistant(
    llm=llm_cfg,
    system_message='You are a helpful assistant that can generate images.',
    function_list=['my_image_gen', 'code_interpreter']
)

# Chat with the agent
messages = [{'role': 'user', 'content': 'Draw a sunset over mountains'}]
for response in bot.run(messages=messages):
    print(response)