import os

os.environ['QIANFAN_ENABLE_STRESS_TEST'] = "true"

from qianfan.dataset import Dataset

# 步骤一，接口调用认证
# 【推荐】使用安全认证AK/SK鉴权，通过环境变量初始化认证信息
# 替换下列示例中参数，安全认证Access Key替换your_iam_ak，Secret Key替换your_iam_sk
os.environ["QIANFAN_ACCESS_KEY"] = "百度智能云IAM的AccessKey"
os.environ["QIANFAN_SECRET_KEY"] = "百度智能云IAM的SecretKey"

# 通过千帆ModelBuilder应用认证
# os.environ["QIANFAN_AK"] = "千帆ModelBuilder控制台-应用接入-创建的应用AK"
# os.environ["QIANFAN_SK"] = "千帆ModelBuilder控制台-应用接入-创建的应用SK"

# 步骤二，初始化一个数据集对象
ds = Dataset.load(data_file="./stress_test_data.jsonl")

# 步骤三，设置压测相关参数

"""
* workers：启动10个进程来模拟并发请求。
* users：总模拟并发100。
* spawn_rate：每秒新增5个，一直增加到100个以后，不再新增， 后续从第11s开始，全部都按照100并发发压。
* model/ endpoint: 预置服务名称/ 定制服务Endpoint，注意两个参数只能填写其一。
* runtime：15m为15分钟。表示该次发压的时限为15分钟，超过该时限，脚本会自动停止，压测结束。
"""
ds.stress_test(
    users=100,
    workers=10,
    spawn_rate=5,
    # model="ERNIE-Speed-8K", # 此处填写预置服务名称，与endpoint二选一
    endpoint="YourEndpoint", # 此处填写您的定制服务Endpoint
    model_type="ChatCompletion", # 可选参数
    runtime="15m", # 可选参数
    hyperparameters={"temperature":0.9} # 可选参数
)

import os

os.environ['QIANFAN_ENABLE_STRESS_TEST'] = "true"

from qianfan.dataset import Dataset

# 步骤一，接口调用认证
#【推荐】使用安全认证AK/SK鉴权，通过环境变量初始化认证信息
# 替换下列示例中参数，安全认证Access Key替换your_iam_ak，Secret Key替换your_iam_sk
os.environ["QIANFAN_ACCESS_KEY"] = "百度智能云IAM的AccessKey"
os.environ["QIANFAN_SECRET_KEY"] = "百度智能云IAM的SecretKey"

# 通过千帆ModelBuilder应用认证
# os.environ["QIANFAN_AK"] = "千帆ModelBuilder控制台-应用接入-创建的应用AK"
# os.environ["QIANFAN_SK"] = "千帆ModelBuilder控制台-应用接入-创建的应用SK"

# 步骤二，初始化一个数据集对象
# 需要初始化一个数据集对象

ds = Dataset.load(data_file="./stress_test_data.jsonl")

# 步骤三，设置压测相关参数
"""
* workers：启动10个进程来模拟并发请求。
* origin_users：初始发压并发数为100。users数量即为并发数量。
* spawn_rate：每秒新增5个user，一直增加到100个以后，不再新增， 后续从第11s开始，全部都按照100并发发压。
* model/ endpoint: 预置服务名称/ 定制服务Endpoint，注意两个参数只能填写其一。
* first_latency_threshold: 首token时延阈值。首token返回时延超过20秒，则不会再发起下轮压测，压测在本轮停止。
* round_latency_threshold: 单轮压测时延阈值。单轮压测时延超过100秒，则不会再发起下轮压测，压测在本轮停止。
* success_rate_threshold: 单轮成功率阈值。本轮压测的请求成功率低于80%，则不会再发起下轮压测，压测在本轮停止。
* rounds：指定压测轮数。
* interval：指定压测轮之间的加压并发数。
"""
ds.multi_stress_test(
    origin_users=100,
    workers=10,
    spawn_rate=5,
    # model="ERNIE-Speed-8K", # 此处填写预置服务名称，与endpoint二选一
    endpoint="YourEndpoint", # 此处填写服务Endpoint
    model_type="ChatCompletion", # 可选参数
    runtime="15m", # 可选参数
    rounds=10, 
    interval=5,
    hyperparameters={"temperature":0.9} # 可选参数
    #first_latency_threshold = 20,
    #round_latency_threshold = 100,
    #success_rate_threshold = 80,
)

Dataset.load(data_file="./stress_test_data.jsonl")