58 lines
2.8 KiB
Python
58 lines
2.8 KiB
Python
from dashscope import Generation
|
||
import random
|
||
from http import HTTPStatus
|
||
from enum import Enum
|
||
from .qgi import Generator as BaseGenerator
|
||
|
||
class Model(str, Enum):
|
||
# 通义千问超大规模语言模型，支持中文、英文等不同语言输入。
|
||
# 模型支持 8,000 tokens上下文，为了保证正常的使用和输出，API限定用户输入为 6,000 tokens。
|
||
QWEN_TURBO = "qwen-turbo"
|
||
|
||
# 通义千问超大规模语言模型增强版，支持中文、英文等不同语言输入。
|
||
# 模型支持 32,000 tokens上下文，为了保证正常的使用和输出，API限定用户输入为30,000 tokens。
|
||
QWEN_PLUS = "qwen-plus"
|
||
|
||
# 通义千问千亿级别超大规模语言模型，支持中文、英文等不同语言输入。
|
||
# 随着模型的升级，qwen-max将滚动更新升级，如果希望使用固定版本，请使用下面的历史快照版本。
|
||
# 当前qwen-max模型与qwen-max-0428快照版本等价，均为最新版本的qwen-max模型，也是当前通义千问2.5产品版本背后的API模型。
|
||
# 模型支持 8,000 tokens上下文，为了保证正常的使用和输出，API限定用户输入为 6,000 tokens。
|
||
QWEN_MAX = "qwen-max"
|
||
|
||
# 通义千问千亿级别超大规模语言模型，支持中文、英文等不同语言输入。
|
||
# 模型支持 30,000 tokens上下文，为了保证正常的使用和输出，API限定用户输入为 28,000 tokens。
|
||
QWEN_MAX_LONGCONTEXT = "qwen-max-longcontext"
|
||
|
||
|
||
class Generator(BaseGenerator):
|
||
def __init__(self, model:str, messages:list[dict[str,str]]|None, apikey:str, seed:int=0) -> None:
|
||
'''
|
||
参数:
|
||
- model: str - 使用的大模型名称，例如 qwen-turbo。
|
||
- messages: list[dict[str,str]] - 提供给大模型的 messages 参数
|
||
- apikey: str - 用于调用模型的 dashscope API密钥。
|
||
- seed: int - 随机种子，用于确保生成结果的可复现性，如果忽略或为0，则生成一个随机数。
|
||
'''
|
||
self.model = model
|
||
self.messages = messages
|
||
self.apikey = apikey
|
||
self.seed = seed
|
||
|
||
def _generate(self, prompt:str|None) -> tuple[HTTPStatus, str]:
|
||
# 设置随机数种子seed，如果没有设置，则随机数种子默认为1234
|
||
seed = self.seed
|
||
if seed == 0:
|
||
seed = random.randint(1, 10000)
|
||
|
||
ret = Generation.call(model=self.model,
|
||
messages=self.messages,
|
||
prompt=prompt,
|
||
api_key=self.apikey,
|
||
seed=seed,
|
||
result_format='text')
|
||
|
||
if ret.status_code != HTTPStatus.OK:
|
||
return ret.status_code, ret.message
|
||
else:
|
||
return ret.status_code, ret.output.text
|