ai调用demo

from flask import Flask, request, jsonify
from flask_docs import ApiDoc
from openai import AzureOpenAI
import os
from flask_docs import ApiDoc
from flask_siwadoc import SiwaDoc
from flasgger import Swagger

app = Flask(__name__)

siwa = SiwaDoc(app)
ApiDoc(
    app,
    title="Sample App",
    version="1.0.0",
    description="A simple app API",
)

Swagger(app)

# 初始化Azure OpenAI客户端
version = "2023-05-15"
endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
api_key = os.getenv("AZURE_OPENAI_API_KEY")
client = AzureOpenAI(azure_endpoint=endpoint, api_key=api_key, api_version=version)

# 角色预设
roles = {
    'Admin': [],
    'PythonDevops':[],
    'JavaDevops':[],
}

@app.route('/set_role', methods=['POST'])
@siwa.doc()
def set_role():
    """
    设置当前角色
    """
    global current_role
    role_name = request.json.get('role')
    if role_name in roles:
        current_role = role_name
        return jsonify({"code":100,"message": f"Role set to {current_role}"})
    else:
        return jsonify({"code":999,"message": f"错误 role"})


@app.route('/add_context', methods=['POST'])
@siwa.doc()
def add_context():
    """
    向当前角色添加上下文
    """
    global current_role
    if current_role:
        context = request.json.get('context')
        if context:
            roles[current_role].append(context)
            return jsonify({"code":100,"message": f"Context '{context}' added to {current_role}"})
        else:
            return jsonify({"code":999,"error": f"Context is required"})
    else:
        return jsonify({"code":998,"error": f"role 未设置"})

@app.route('/clear_context', methods=['POST'])
@siwa.doc()
def clear_context():
    """
    清除当前角色的上下文
    """
    global current_role
    if current_role:
        roles[current_role].clear()
        return jsonify({"code":100,"message": f"{current_role}'s context cleared"})
    else:
        return jsonify({"code":998,"error": "role 未设置"})

@app.route('/ask', methods=['POST'])
@siwa.doc()
def ask():
    """
    向AI模型提问，基于当前角色和上下文
    """
    global current_role, client
    if current_role:
        question = request.json.get('question')
        if question:
            messages = [{"role": "system", "content": f"You are now {current_role}."}]
            messages.extend([{"role": "user", "content": ctx} for ctx in roles[current_role]])
            messages.append({"role": "user", "content": question})
            
            try:
                response = client.chat.completions.create(
                    model="gpt-35-turbo-16k",
                    messages=messages
                )
                answer = response.choices[0].message.content
                roles[current_role].append(question)
                roles[current_role].append(answer)
                return jsonify({"code":100,"message": answer})
            except Exception as e:
                return jsonify({"code":888,"error": str(e)})
        else:
            return jsonify({"code":997,"error": "没有输入问题"})
    else:
        return jsonify({"code":998,"error": "role 未设置"})
@app.route('/view_context', methods=['GET'])
def aview_context():
    """查看上下文"""
    global current_role
    if current_role:
        return jsonify({"code":100,"message": f"{roles[current_role]}","more":f"{roles}"})
    else:
        return jsonify({"code":998,"error": f"role 未设置"})
if __name__ == '__main__':
    app.run(debug=True)