有JWT认证的API调用大模型
大模型开发/技术交流
- 大模型推理
- AI加速器线上加速营
- API
9月12日4228看过
本文基于Python语言的Flask框架封装自定义的大语言模型推理的API,该API实现了JWT的用户认证,用户数据保存在内置的sqllite数据库中。涉及技术有:轻量数据库sqlite、向量数据库VDB、AppBuilder SDK、千帆SDK、流式输出、JWT认证等。
下面代码涉及到的python模块请使用pip install -r requirements.txt进行安装。
appbuilder_sdk==0.9.3Flask==3.0.3PyJWT==2.9.0qianfan==0.4.8redis==5.0.8
一、创建数据库并插入测试用户数据的代码,以下代码包括初始化数据和数据库查询的封装,表字段包括:用户名、密码、权限等信息。文件名:mysqlite.py。初始化数据库调用如下命令:python mysqlite.py
代码如下:
import sqlite3# 定义数据库文件名db_name = 'llm_user.db'def initDB():# 使用 with 语句创建数据库连接with sqlite3.connect(db_name) as conn:# 创建一个 Cursor 对象cursor = conn.cursor()# 执行创建表的 SQL 语句cursor.execute('''CREATE TABLE IF NOT EXISTS users (id INTEGER PRIMARY KEY AUTOINCREMENT,username TEXT NOT NULL,password TEXT NOT NULL,permission BOOL NOT NULL)''')# 执行插入数据的 SQL 语句cursor.execute('''INSERT INTO users (username, password, permission) VALUES (?, ?, ?)''', ('guo', 'guopwd', True))cursor.execute('''INSERT INTO users (username, password, permission) VALUES (?, ?, ?)''', ('visitor', 'vpwd', False))# 提交事务conn.commit()print("数据库初始化完成")def queryDB(sql):# 使用 with 语句创建数据库连接with sqlite3.connect(db_name) as conn:# 创建一个 Cursor 对象cursor = conn.cursor()# 执行查询数据的 SQL 语句cursor.execute(sql)# 获取查询结果row = cursor.fetchone()return rowif __name__ == '__main__':# 初始化数据库。initDB()
二、JWT公共类,使用了上面定义的mysqlite模块,查询用户是否有权限执行后续操作。文件名:auth_util.py代码如下:
import jwtimport datetimefrom flask import request,jsonify# 查询python自带的sqllite数据库,需要先初始化数据库,执行:python mysqlite.pyfrom mysqlite import queryDB# 秘钥,用于签名和验证TokenSECRET_KEY = 'guokey888'# 生成Token的函数def generate_token(user_info):exp_time = datetime.datetime.utcnow() + datetime.timedelta(minutes=30)payload = {'exp': exp_time,'user_info': user_info}token = jwt.encode(payload, SECRET_KEY, algorithm='HS256')return jsonify({'token':token,'name':user_info["username"]})# 验证Token的装饰器def token_required(f):def decorated(*args, **kwargs):token = request.headers.get('Authorization')if not token:return jsonify({'message': 'Token is missing!'}), 401try:data = jwt.decode(token, SECRET_KEY, algorithms=['HS256'])kwargs['user_info'] = data['user_info']sql = "select * from users where username='"+data['user_info']['username']+"' and password='"+data['user_info']['password']+"'"print (sql)user_row = queryDB(sql)print ("user info: "+ str(user_row))if not user_row:# 判断用户是否有某种权限逻辑,如果查询不到用户,设置为0kwargs['hasPermission'] = 0 #转化为0else:# 判断用户是否有某种权限逻辑kwargs['hasPermission'] = user_row[3] #True转化为1except jwt.ExpiredSignatureError:return jsonify({'message': 'Token has expired!'}), 401except jwt.InvalidTokenError:return jsonify({'message': 'Token is invalid!'}), 401return f(*args, **kwargs)return decorated
三、API接口,这个py文件,包括使用AppBuilder SDK和千帆SDK 两种方式调用大语言模型和VDB向量数据库等示例代码。注意替换api_baidu.py代码中的key为自己的值。
from flask import Flask, request, jsonify,Response# 引入自定义的token生成模块import auth_utilapp = Flask(__name__)# 登录接口,用于生成Token@app.route('/login', methods=['POST'])def login():user_info = request.json.get('user_info')if not user_info:return jsonify({'message': 'User info is required!'}), 400token = auth_util.generate_token(user_info)return token# 受保护的接口,需要Token验证@app.route('/protected_api', methods=['GET'])@auth_util.token_requireddef protected(user_info,hasPermission):return jsonify({'message': 'This is a protected endpoint!', 'user_info': user_info,"hasPermission":hasPermission}), 200import osimport appbuilder# 请前往千帆AppBuilder官网创建密钥,流程详见:https://cloud.baidu.com/doc/AppBuilder/s/Olq6grrt6#1%E3%80%81%E5%88%9B%E5%BB%BA%E5%AF%86%E9%92%A5os.environ["APPBUILDER_TOKEN"] = 'bce-v3/ALTAK-bppeq3zxoIY1Kc4C.......'######使用appbuilder SDK调用VDB向量数据库查询相似问题示例-begin#######segments = appbuilder.Message(["Guoluan同学是大二老生", "Guoluan同学在黑龙江大学就读", "Guoluan同学从3岁开始跳舞", "黑大诞生于延安","Guoluan同学的专业是计算机及应用","Guoluan同学的生日跟财神爷一天","昨天交了453.2的英语书本费","黑龙江大学今年8月24日开学"])# 初始化构建索引vector_index = appbuilder.BaiduVDBVectorStoreIndex.from_params(instance_id='vdb-bj-driaxo******',# 替换为你的VDB的instance_idapi_key='vdb$20*****',# VDD数据库的api_key,即root用户的密码,可以登录VDB数据库页面复制drop_exists=True,# 如果索引存在则删除,否则可能会查出多条相同数据)vector_index.add_segments(segments)@app.route('/chatwithVDB', endpoint='chatwithVDB', methods=['POST'])@auth_util.token_requireddef chat_with_VDB(user_info,hasPermission):if hasPermission==1:user_input = request.args.get('user_input')query = appbuilder.Message(user_input)retriever = vector_index.as_retriever()res = retriever(query,3)# print(res)def sse_output():for item in res.content:msg = item['text']score = item['score']print(f"message:{msg},score:{score}") #控制台打印看输出# 仅输出score值大于0.5的数据库返回信息,注:值越高相似语义度越强。if score >0.5:# 组装一个字典类型的json输出,如果只输出内容msg,则直接yield这个字符串format_msg = {}format_msg["msg"] = msgformat_msg["score"] = scoreyield f"data: {format_msg}\n\n" # SSE 要求每条消息后有两个换行符return Response(sse_output(), mimetype='text/event-stream')else:return "你没有权限访问大语言模型"######使用appbuilder SDK调用VDB向量数据库查询相似问题示例-end#############使用appbuilder SDK调用大模型流式返回示例-begin#######@app.route('/chat_with_ABSDK', endpoint='chat_with_ABSDK', methods=['POST'])@auth_util.token_requireddef chat_with_ABSDK(user_info,hasPermission):if hasPermission==1:user_input = request.args.get('user_input')# 定义prompt模板,template_str = "您作为一位情感交流专家,扮演{role}, 会尽你所能地展现同理心,适时提出引导性问题,确保我们的交流既富有意义又充满乐趣,以此营造一个积极、愉快的互动环境。请回答我的问题。\n\n问题:{question}。\n\n回答:"# 定义输入,调用playground组件input_msg = appbuilder.Message({"role": "我是热爱生活的年轻人,喜欢骑自行车,平时也爱学习", "question": user_input})playground = appbuilder.Playground(prompt_template=template_str, model="ERNIE Speed-AppBuilder")# 以打字机的方式,流式展示大模型回答内容output = playground(input_msg, stream=True, temperature=0.8)def sse_output():for stream_message in output.content:yield f"data: {stream_message}\n\n" # SSE 要求每条消息后有两个换行符return Response(sse_output(), mimetype='text/event-stream')else:return "你没有权限访问大语言模型"######使用appbuilder SDK调用大模型流式返回示例-end#############使用qianfan SDK调用大模型流式返回示例-begin#########调用大模型qianfan SDK比AppBuilder SDK更具有广泛性,可能是AppBuilder SDK虽然已经支持了很多模型,但需要再包装一下才能调用新模型。import osimport qianfan# 通过环境变量初始化认证信息# 推荐】使用安全认证AK/SK鉴权# 替换下列示例中参数,安全认证Access Key替换your_iam_ak,Secret Key替换your_iam_sk,如何获取请查看https://cloud.baidu.com/doc/Reference/s/9jwvz2egbos.environ["QIANFAN_ACCESS_KEY"] = "ALTAKqk85NxpKjl****"os.environ["QIANFAN_SECRET_KEY"] = "3416b7f7271b48d************7ae50"chat_comp = qianfan.ChatCompletion()@app.route('/chat_with_qianfan', endpoint='chat_with_qianfan', methods=['POST'])@auth_util.token_requireddef chat_with_qianfanSDK(user_info,hasPermission):user_input = request.args.get('user_input')resp = chat_comp.do(model="ERNIE-Speed-Pro-128K", messages=[{"role": "user","content": user_input}], stream=True)def sse_output():for r in resp:yield f"data: {r['body']['result']}\n\n" # SSE 要求每条消息后有两个换行符return Response(sse_output(), mimetype='text/event-stream')######使用qianfan SDK调用大模型流式返回示例-end#######if __name__ == '__main__':app.run(debug=True)
四、启动API接口的python文件
python api_baidu.py
五、测试
如果公网IP不能访问,Flask应用默认监听本地回环地址,这意味着它只能在本机上进行访问测试,无法从其他设备或网络访问。为了使Flask应用能够通过公网IP进行访问,需要确保在启动应用时指定正确的IP地址。在启动Flask应用时,使用
app.run(host='0.0.0.0')
或app.run(host='::')
可以确保应用监听所有可用的IPv4和IPv6地址。这样做可以让应用不仅监听本机的所有IP地址,还包括公网IP地址,从而允许外部访问。
首先访问login接口获取token。
入参为json格式(可以输入错误的用户名密码尝试无权访问情况):
{"user_info": {"username": "guo","password": "guopwd","level":"architect"}}
返回:
{"name": "guo","token": "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJleHAiOjE3MjYxMzYwODksInVzZXJfaW5mbyI6eyJ1c2VybmFtZSI6Imd1byIsInBhc3N3b3JkIjoiZ3VvcHdkIiwibGV2ZWwiOiJhcmNoaXRlY3QifX0.zHkibFhPU9qxPx0OL_fUmeWiXZ2WFb84Hua9RhmNDQk"}
第二步,测试与大语言模型的对话,
注意:在header中的token设置,我这是使用Postman进行测试。
可以看到流式输出的结果:
其他几个API测试地址如下,调用方式同上。
评论