'''
Created on 2025年3月18日

@author: foxba
'''
import os
from openai import OpenAI
from transformers.utils.versions import require_version
import time

require_version("openai>=1.5.0", "To fix: pip install openai>=1.5.0")

def callServer(input,modelName):
    port = 8000
    client = OpenAI(
        api_key="0",
        base_url="http://localhost:{}/v1".format(os.environ.get("API_PORT", port)),
    )
    messages = []
    messages.append({"role": "user", "content": input})
    result = client.chat.completions.create(messages=messages, model=modelName,max_tokens=20000)
    return result.choices[0].message.content
    

    
    
if __name__ == "__main__":
    #测试思维链
    input1 = "对于「初三女生在搀扶跌倒老奶奶后反被冤枉，但仍选择资助她千元」的新闻事件，你有什么看法？"
    modelName1="cot"
    #测试分类
    input2 = """你是一位电商评论文本的分类专家，在【分类体系】规定的类别中，【输入文本】最可能分到哪个类别？只允许输出类别名称，不允许输出其他字符。
                【分类体系】
                 ["电脑","水果","平板","衣服","酒店","洗浴","书籍","蒙牛","电器","手机"]
                 【输入文本】
                 很不错的一家5星级酒店。房间也很干净，就是电视太小，而且离床比较远，看不太清楚。服务人员态度也很好。下次一定会再住。"""
    modelName2="cat"
    #测试SPO
    input3 = """你是一位电商评论文本的分类专家，在【输入文本】中检测所有的SPO三元组，并且以【指定格式】输出。只允许输出【指定格式】规定的答案，不允许输出任何其他字符。
                【指定格式】
                [{'predicate': '编剧', 'object_type': '人物', 'subject_type': '影视作品', 'object': '神尾叶子', 'subject': '花样男子'}, {'predicate': '出版社', 'object_type': '出版社', 'subject_type': '书籍', 'object': '集英社', 'subject': '花样男子'}]
                【输入文本】
                 王炳煌（1924－1951），福建省南安县水头镇山前村人。"""
    modelName3="spo"

    a = time.time()
    result = callServer(input1,modelName1)
    b = time.time()
    print("输出长度："+str(len(result))+"字  耗时："+str(round(b-a,2))+"秒  输出速度："+str(round(len(result)/(b-a),2))+"字/秒")
    print(result)