DatabricksのドライバープロキシーでOpenAI APIのストリーミングを行う

なかなかニッチだと思いつつ、少し進捗が出たのでまとめます。

# 要件

APIでOpenAIのコンプリーションを取得する際、デフォルトではすべての生成が終わってから返却されます。プロンプトによっては数十秒待つことに。これを改善するために、コンプリーションを**ストリーム**することができます。

OpenAI APIを直接呼び出すのであればシンプルです。こちらの記事に書いているように引数`stream`を`True`に設定するだけです。

https://qiita.com/taka_yayoi/items/a389cafdd0544504e1a6

ただ、実際にはAPIの呼び出し前や呼び出し後に処理を追加したいケースが多く、その場合には前処理、後処理を含むロジックを一連のパイプラインとしてデプロイする必要が出てきます。

そのため、ここではDatabricksのドライバープロキシーを用いてFlaskに処理のパイプラインをデプロイして、クライアントから呼び出す構成を組みます。その際に生成結果をストリーミングするようにします。

https://qiita.com/taka_yayoi/items/f7538636e88bc8973769

# 実装

## サーバー

```py
import os
os.environ["OPENAI_API_KEY"] = dbutils.secrets.get("demo-token-takaaki.yayoi", "openai_api_key")
```

```py
import openai

def chat_gpt_helper(prompt):
    """
    この関数はOpenAIのcompletions APIを用いてGpt4モデルのレスポンスを返却します
    """
    try:
        resp = ''
        openai.api_key = os.getenv('OPENAI_API_KEY')
        for chunk in openai.ChatCompletion.create(
            model="gpt-4",
            messages=[{
                "role": "user",
                "content":prompt
            }],
            stream=True,
        ):
            content = chunk["choices"][0].get("delta", {}).get("content")
            if content is not None:
                    yield content

    except Exception as e:
        print(e)
        return str(e)
```

Flask側でもMIMEタイプ`text/event-stream`でストリームを搬出しています。

```py
from flask import Flask, jsonify, request, Response, stream_with_context

app = Flask("openai-streaming")

@app.route('/', methods=['POST'])
def stream_chat_gpt():
        """
        ChatGPTのレスポンスをストリーミングします
        """
        prompt = request.get_json(force = True).get('prompt','')
        return Response(stream_with_context(chat_gpt_helper(prompt)),
                         mimetype='text/event-stream')
```

こちらはサーバーのURLやポート番号を確認するためのものです。

```py
from dbruntime.databricks_repl_context import get_context
ctx = get_context()

port = "7777"
driver_proxy_api = f"https://{ctx.browserHostName}/driver-proxy-api/o/0/{ctx.clusterId}/{port}"

print(f"""
driver_proxy_api = '{driver_proxy_api}'
cluster_id = '{ctx.clusterId}'
port = {port}
""")
```

```
driver_proxy_api = 'https://xxxxxx.databricks.com/driver-proxy-api/o/0/0713-082554-nqq9eafd/7777'
cluster_id = '0713-082554-nqq9eafd'
port = 7777
```

以下のコマンドでドライバープロキシーを起動します。明示的に停止しない限り実行され続けるので注意してください。

```py
app.run(host="0.0.0.0", port=port, debug=True, use_reloader=False, threaded=True)
```

## クライアント

今回は同じDatabricksワークスペースの同じクラスターに別のノートブックをアタッチします。

クライアント側も`text/event-stream`を受け付けるようにしています。

```py:client
import requests
import json

def get_stream(data):

    headers = {"Accept": "text/event-stream"}
    output_text = []

    s = requests.Session()
    with s.post("http://127.0.0.1:7777/", headers=headers, json=data, stream=True) as resp:
        for line in resp.iter_content(decode_unicode=True):
            if line:
                print(line, end="", flush=True)
                output_text.append(line)
        return "".join(output_text)
```

プロンプトを入力してみます。

```py
data = {"prompt": "猫を讃える歌を歌ってください"}
get_stream(data)
```

<iframe width="560" height="315" src="https://www.youtube.com/embed/Sg-5lszVZ_s?si=zljiBtR8V7-1ETUl" title="YouTube video player" frameborder="0" allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture; web-share" allowfullscreen></iframe>

動きました！

# 参考資料

色々参考にさせていただいています。

- [openai\-cookbook/examples/How\_to\_stream\_completions\.ipynb at main · openai/openai\-cookbook](https://github.com/openai/openai-cookbook/blob/main/examples/How_to_stream_completions.ipynb)
- [Streaming\-ChatGPT\-Responses\-in\-ReactJS/server/app\.py at main · anshumankmr/Streaming\-ChatGPT\-Responses\-in\-ReactJS](https://github.com/anshumankmr/Streaming-ChatGPT-Responses-in-ReactJS/blob/main/server/app.py)
- [Stream Responses from OpenAI API with Python: A Step\-by\-Step Guide — CodingTheSmartWay](https://www.codingthesmartway.com/stream-responses-from-openai-api-with-python/)
- [ChatGPTをぬるぬるにする🐌Server\-Sent Eventsの基礎知識](https://zenn.dev/chot/articles/a089c203adad74)
- [pythonでSSE \(server side events\) の例を作ってみて遊んでみた \- podhmo's diary](https://pod.hatenablog.com/entry/2021/05/26/200403)
- [Unable to send OpenAI stream response via Flask API \- Stack Overflow](https://stackoverflow.com/questions/76311866/unable-to-send-openai-stream-response-via-flask-api)
- [Streamlitを使ってChatGPTのようなチャットアプリを簡単に作る](https://zenn.dev/nishijima13/articles/3b1a50b8728261)
- [【python】【リアルタイム通信】FlaskでSSE\(Server\-Sent Events\)を実装する](https://rihitosan.com/blog/flask_sse.html)
- [FastAPIを使ってストリーム対応のLLM API RESTサーバを作る on Databricks \- Qiita](https://qiita.com/isanakamishiro2/items/c55123343e95fd2769a6#%E3%82%AF%E3%83%A9%E3%82%A4%E3%82%A2%E3%83%B3%E3%83%88%E3%81%8B%E3%82%89%E3%81%AE%E6%8E%A5%E7%B6%9A)
- [StreamlitでChatGPTのStreamを表示する \- Qiita](https://qiita.com/suzuki_sh/items/64d84c417cba43cd6351)

### Databricksクイックスタートガイド

[Databricksクイックスタートガイド](https://www.amazon.co.jp/dp/B09V1YXFVQ/)


### Databricks無料トライアル

[Databricks無料トライアル](https://databricks.com/jp/try-databricks)