Simple Flask server that streams Azure OpenAI responses to its root endpoint.

FlaskOpenAIStreamingAPI.py

# requirements.txt:
#
# flask
# flask_restx
# openai
# python-dotenv
#
# .env:
# AZURE_OPENAI_ENDPOINT=
# AZURE_OPENAI_API_KEY=
# AZURE_OPENAI_DEPLOYMENT_NAME=

import os

from dotenv import load_dotenv
from flask import Flask, Response, request
from openai import AzureOpenAI

load_dotenv()

app = Flask(__name__)


@app.route("/", methods=["POST"])
def main():
    message = request.json["message"]

    client = AzureOpenAI(
        azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
        api_key=os.environ["AZURE_OPENAI_API_KEY"],
        api_version="2024-04-01-preview"
    )

    messages = [
        {"role": "system", "content": "You are Codsworth, a robot from the game Fallout 4."},
        {"role": "user", "content": message}
    ]

    generator = client.chat.completions.create(
        model=os.environ["AZURE_OPENAI_DEPLOYMENT_NAME"],
        messages=messages,
        stream=True
    )

    def stream(generator):
        for chunk in generator:
            if chunk.choices and chunk.choices[0].finish_reason is None:
                stream_token = chunk.choices[0].delta.content
                print(stream_token, flush=True)
                stream_token = stream_token.replace("\n", "\\n")
                yield stream_token

    response = Response(stream(generator), mimetype="text/event-stream")
    return response


if __name__ == "__main__":
    app.run()