#!/usr/bin/env python3
#coding:utf-8

__author__ = 'xmxoxo<xmxoxo@qq.com>'

import logging
from ..baselibs.llm_task import LLM_TASK

def test_llm_task():

    model = "qwen-turbo" # qwen3-32b qwq-32b qwen-turbo
    print(f"model:{model}")
    
    # 定义提示词
    prompt_template = """
    # 请根据用户的关键词生成一首七言律诗
    # 输出格式：请使用JSON格式输出,格式为：{"result":"生成结果"}。
    # 关键词：
    $keyword$
    """

    task_config = {
        "base_url": "http://192.168.15.111:3000/v1",    # one-api
        "api_key": "sk-pLLG2ucf61sKFjMxA0Fd11E88c734427A078Bc554e516e26",
        "model": model,
        "prompt_template": prompt_template
    }
    # 构建对象实例
    llm = LLM_TASK(**task_config)
    
    # 调用
    parm = {"$keyword$":" 郑成功 秋天 中美贸易战 国际形势"}
    print(f"parm:{parm}")
    result = llm.predict(parm)
    print(result)
    assert type(result) == dict

if __name__ == '__main__':
    logging.basicConfig(level=logging.DEBUG) # INFO  DEBUG
    test_llm_task()

