bizwechat/sqlcode/qwenapi.py

58 lines
2.8 KiB
Python
Raw Normal View History

2025-02-17 10:34:35 +08:00
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