MLflowのOpenAIフレーバーにおけるAPIキー管理方法の変更

こちらを読んでいたら以下のような記述が。

https://mlflow.org/docs/latest/python_api/openai/index.html

> Specifying secrets for model serving with `MLFLOW_OPENAI_SECRET_SCOPE` is deprecated. Use [secrets-based environment variables](https://docs.databricks.com/ja/machine-learning/model-serving/store-env-variable-model-serving.html) instead.

これまでは、OpenAIのAPIキーはDatabricksシークレットのスコープ名を環境変数[`MLFLOW_OPENAI_SECRET_SCOPE`に指定](https://qiita.com/taka_yayoi/items/a058484e6c0abbfbc476)する形で引き渡していましたが、これが非推奨になるとのこと。代わりにシークレットベースの環境変数を指定しろとのこと。なので、早速試してみます。

```py
%pip install tiktoken
%pip install openai
```

```py
import os
import mlflow
import openai

model_name = "takaakiyayoi_catalog.rag_chatbot.open_ai"

mlflow.set_registry_uri("databricks-uc")

with mlflow.start_run():
    model_info = mlflow.openai.log_model(
        model="gpt-3.5-turbo",
        task=openai.ChatCompletion,
        messages=[{"role": "system", "content": "You are an MLflow expert"}],
        artifact_path="model",
        registered_model_name = model_name,
    )
```

MLflowでトラッキングするとともに、Unity Catalogに登録します。
![Screenshot 2023-12-19 at 21.39.52.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/d7a2f006-00ff-5c12-927c-73a158280beb.png)

REST APIを使ってモデルサービングエンドポイントを作成しますが、この際に`environment_vars`にシークレットを参照する環境変数を設定します。以下の例ではスコープ名が`demo-token-takaaki.yayoi`でキー名が`openai`のシークレットにOpenAIのAPIキーを格納しています。

```py
import requests

secret_scope_name = "demo-token-takaaki.yayoi"
secret_key_name = "openai"
databricks_host = "<Databricksワークスペースのホスト名>"

endpoint_name = "my_serving_endpoint_taka"
model_name = model_name
model_version = 1
databricks_api_token = dbutils.notebook.entry_point.getDbutils().notebook().getContext().apiToken().get()

# Create your endpoint
data = {
    "name": endpoint_name,
    "config": {
        "served_models": [
            {
                "model_name": model_name,
                "model_version": model_version,
                "workload_size": "Small",
                "scale_to_zero_enabled": True,
                "environment_vars": {
                    "OPENAI_API_KEY": f"{{{{secrets/{secret_scope_name}/{secret_key_name}}}}}"
                }
            }
        ]
    }
}

headers = {
    "Context-Type": "text/json",
    "Authorization": f"Bearer {databricks_api_token}"
}

response = requests.post(
    url=f"https://{databricks_host}/api/2.0/serving-endpoints",
    json=data,
    headers=headers
)
print("Response status:", response.status_code)
print("Response text:", response.text, "\n")
```

しばらくするとエンドポイントが稼働し始めます。
![Screenshot 2023-12-19 at 21.42.30.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/14b9a6a7-df5e-0d82-61c3-664e094dd55a.png)

動作確認のために(いつものように)streamlitで画面を作ります。

```py:/Workspace/Users/takaaki.yayoi@databricks.com/20231219_openai/streamlit.py
import streamlit as st 
import numpy as np 
from PIL import Image
import base64
import io

import os
import requests
import numpy as np
import pandas as pd

import json

from databricks.sdk.runtime import dbutils

st.header('OpenAI Chatbot on Databrikcs')

def score_model(prompt):
  # 1. パーソナルアクセストークンを設定してください
  token = dbutils.secrets.get("demo-token-takaaki.yayoi", "token")

  # 2. モデルエンドポイントのURLを設定してください
  url = '<モデルサービングエンドポイントのURL>'
  headers = {'Authorization': f'Bearer {token}'}
  
  data_json_str = f"""
  {{"inputs":[
{{"role": ["user"], "content": ["{prompt}"]}}
]
}}
"""

  data_json = json.loads(data_json_str)
  response = requests.request(method='POST', headers=headers, url=url, json=data_json)

  if response.status_code != 200:
    raise Exception(f'Request failed with status {response.status_code}, {response.text}')
  return response.json()

prompt = st.text_input("プロンプト")

if prompt != "":
  response = score_model(prompt)
  st.write(response['predictions'][0])
```

元のノートブックに戻ってstreamlitをインストールします。

```py
%pip install streamlit watchdog
dbutils.library.restartPython()
```

```py
from dbruntime.databricks_repl_context import get_context

def front_url(port):
    """
    フロントエンドを実行するための URL を返す

    Returns
    -------
    proxy_url : str
        フロントエンドのURL
    """
    ctx = get_context()
    proxy_url = f"https://{ctx.browserHostName}/driver-proxy/o/{ctx.workspaceId}/{ctx.clusterId}/{port}/"

    return proxy_url

PORT = 1501

# Driver ProxyのURLを表示
print(front_url(PORT))

# 利便性のためにリンクをHTML表示
displayHTML(f"<a href='{front_url(PORT)}' target='_blank' rel='noopener noreferrer'>別ウインドウで開く</a>")
```

streamlitを起動します。

```py
streamlit_file = "/Workspace/Users/takaaki.yayoi@databricks.com/20231219_openai/streamlit.py"
!streamlit run {streamlit_file} --server.port {PORT}
```

上のセルを実行した際に表示されるリンクをクリックするとGUIにアクセスできます。

ということで動作確認もできました！
![Screenshot 2023-12-19 at 21.45.36.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/588ceeec-1ee5-61a4-f144-ae97cc036e82.png)


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

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


### Databricks無料トライアル

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