有JWT认证的API调用大模型

大模型开发/技术交流

大模型推理
AI加速器线上加速营
API

9月12日4517看过

本文基于Python语言的Flask框架封装自定义的大语言模型推理的API，该API实现了JWT的用户认证，用户数据保存在内置的sqllite数据库中。涉及技术有：轻量数据库sqlite、向量数据库VDB、AppBuilder SDK、千帆SDK、流式输出、JWT认证等。

下面代码涉及到的python模块请使用pip install -r requirements.txt进行安装。

  
  
   
   
    appbuilder_sdk==0.9.3
   
   
   
   
    Flask==3.0.3
   
   
   
   
    PyJWT==2.9.0
   
   
   
   
    qianfan==0.4.8
   
   
   
   
    redis==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 row
   
   
   
   
            
   
   
   
   
    if __name__ == '__main__':
   
   
   
   
        # 初始化数据库。
   
   
   
   
        initDB()

二、JWT公共类，使用了上面定义的mysqlite模块，查询用户是否有权限执行后续操作。文件名：auth_util.py代码如下：

  
  
   
   
    import jwt
   
   
   
   
    import datetime
   
   
   
   
    from flask import request,jsonify
   
   
   
   
    # 查询python自带的sqllite数据库，需要先初始化数据库，执行：python mysqlite.py
   
   
   
   
    from mysqlite import queryDB
   
   
   
   
    
   
   
   
   
    # 秘钥，用于签名和验证Token
   
   
   
   
    SECRET_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!'}), 401
   
   
   
   
            
   
   
   
   
            try:
   
   
   
   
                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:
   
   
   
   
                    # 判断用户是否有某种权限逻辑，如果查询不到用户，设置为0
   
   
   
   
                    kwargs['hasPermission'] = 0 #转化为0
   
   
   
   
                else:
   
   
   
   
                    # 判断用户是否有某种权限逻辑
   
   
   
   
                    kwargs['hasPermission'] = user_row[3] #True转化为1
   
   
   
   
                    
   
   
   
   
            except jwt.ExpiredSignatureError:
   
   
   
   
                return jsonify({'message': 'Token has expired!'}), 401
   
   
   
   
            except jwt.InvalidTokenError:
   
   
   
   
                return jsonify({'message': 'Token is invalid!'}), 401
   
   
   
   
            
   
   
   
   
            return 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_util
   
   
   
   
    
   
   
   
   
    app = 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!'}), 400
   
   
   
   
        token = auth_util.generate_token(user_info)
   
   
   
   
        return token
   
   
   
   
    
   
   
   
   
    
   
   
   
   
    # 受保护的接口，需要Token验证
   
   
   
   
    @app.route('/protected_api', methods=['GET'])
   
   
   
   
    @auth_util.token_required
   
   
   
   
    def protected(user_info,hasPermission):
   
   
   
   
        return jsonify({'message': 'This is a protected endpoint!', 'user_info': user_info,"hasPermission":hasPermission}), 200
   
   
   
   
    
   
   
   
   
    
   
   
   
   
    import os
   
   
   
   
    import 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%A5
   
   
   
   
    os.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_id
   
   
   
   
        api_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_required
   
   
   
   
    def 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"] = msg
   
   
   
   
                        format_msg["score"] = score
   
   
   
   
    
   
   
   
   
                        yield 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_required
   
   
   
   
    def 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 os
   
   
   
   
    import qianfan
   
   
   
   
    
   
   
   
   
    # 通过环境变量初始化认证信息
   
   
   
   
    # 推荐】使用安全认证AK/SK鉴权
   
   
   
   
    # 替换下列示例中参数，安全认证Access Key替换your_iam_ak，Secret Key替换your_iam_sk，如何获取请查看https://cloud.baidu.com/doc/Reference/s/9jwvz2egb
   
   
   
   
    os.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_required
   
   
   
   
    def 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。

http://localhost:5000/login

入参为json格式（可以输入错误的用户名密码尝试无权访问情况）：

  
  
   
   
    {
   
   
   
   
        "user_info": {
   
   
   
   
            "username": "guo",
   
   
   
   
            "password": "guopwd",
   
   
   
   
            "level":"architect"
   
   
   
   
        }
   
   
   
   
    }

  
  
   
   
    {
   
   
   
   
        "name": "guo",
   
   
   
   
        "token": "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJleHAiOjE3MjYxMzYwODksInVzZXJfaW5mbyI6eyJ1c2VybmFtZSI6Imd1byIsInBhc3N3b3JkIjoiZ3VvcHdkIiwibGV2ZWwiOiJhcmNoaXRlY3QifX0.zHkibFhPU9qxPx0OL_fUmeWiXZ2WFb84Hua9RhmNDQk"
   
   
   
   
    }

第二步，测试与大语言模型的对话，

http://localhost:5000/chat_with_ABSDK?user_input=你有什么高兴的事？

注意：在header中的token设置，我这是使用Postman进行测试。