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

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

import logging
from llm_task import LLM_TASK
# from baselibs import LLM_TASK

def test_llm_task():
    ''' 单元测试
    '''
    model = "qwen-turbo"
    print(f"model:{model}")
    # 定义提示词
    prompt_template = """
    # 请根据用户的关键词生成一首七言律诗
    # 输出格式：请使用JSON格式输出,格式为：{"result":"生成结果"}。
    # 关键词：
    $keyword$
    """
    task_config = {
        "api_key": "sk-pLLG2ucf61sKFjMxA0Fd11E88c734427A078Bc554e516e26",
        "base_url": "http://192.168.15.111:3000/v1",
        "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
    print('use token:', llm.total_tokens, llm.usages)

def test_llm_task_toon():
    ''' 单元测试，使用TOON 格式调用LLM_TASK
    '''
    model = "qwen-turbo" # qwen3-32b qwq-32b qwen-turbo
    print(f"model:{model}")
    
    # 定义提示词
    prompt_template = """
    # 请根据用户的关键词生成一首七言律诗
    # 输出格式：请使用TOON格式输出, 格式为：
    result[8]:"句子1，", "句子2。", "句子3", "句子4"
    # 关键词：
    $keyword$
    """

    task_config = {
        # "api_key": "sk-pLLG2ucf61sKFjMxA0Fd11E88c734427A078Bc554e516e26",
        "api_key": "sk-GQtwF5ag8p6m8wWf1232B8D5E17f4455A5C14e7a2d393aEe",
        "base_url": "http://192.168.15.111:3000/v1",
        "model": model,
        "prompt_template": prompt_template,
        "use_toon": True
    }
    # 构建对象实例
    llm = LLM_TASK(**task_config)
    
    # 调用
    parm = {"$keyword$":" 郑成功 秋天 中美贸易战 国际形势"}
    parm = {"$keyword$":" 郑成功 冬天 台海局势 中美国际形势"}
    parm = {"$keyword$":" 郑成功 厦门 鼓浪屿 冬天 台海局势 中美贸易战"}
    print(f"parm:{parm}")
    result = llm.predict(parm)
    print(result)
    assert type(result) == dict
    print('use token:', llm.total_tokens, llm.usages)

if __name__ == '__main__':
    logging.basicConfig(level=logging.DEBUG) # INFO DEBUG ERROR
    import fire
    fire.Fire()

