OpenAIのトレース

MLflow Tracing は、OpenAI の自動トレース機能を提供します。自動トレースを有効にする OpenAI の場合、 mlflow.openai.autolog 関数を呼び出すと、 MLflow は LLM 呼び出しのトレースをキャプチャし、それらをアクティブなエクスペリメント MLflow に記録します。

MLflow トレースでは、OpenAI 呼び出しに関する次の情報が自動的にキャプチャされます。

• プロンプトと完了応答
• 待ち時間
• モデル名
• temperature、max_tokensなどの追加のメタデータ (指定されている場合)。
• 応答で返された場合の関数呼び出し
• 例外が発生した場合

前提 条件

OpenAIで MLflow Tracing を使用するには、 MLflow とOpenAI SDKをインストールする必要があります。

Development

開発環境の場合は、Databricks の追加機能と openaiを含む完全な MLflow パッケージをインストールします。

Bash

pip install --upgrade "mlflow[databricks]>=3.1" openai

フル mlflow[databricks] パッケージには、Databricks でのローカル開発と実験のためのすべての機能が含まれています。

Production

本番運用デプロイメントの場合は、 mlflow-tracing と openaiをインストールします。

Bash

pip install --upgrade mlflow-tracing openai

mlflow-tracingパッケージは、本番運用での使用に最適化されています。

注記

OpenAI で最高のトレース エクスペリエンスを得るには、MLflow 3 を強くお勧めします。

例を実行する前に、環境を構成する必要があります。

Databricks ノートブックの外部ユーザーの場合 : Databricks 環境変数を設定します。

Bash

export DATABRICKS_HOST="https://your-workspace.cloud.databricks.com"
export DATABRICKS_TOKEN="your-personal-access-token"

Databricks ノートブック内のユーザーの場合 : これらの資格情報は自動的に設定されます。

API キー : OpenAI API キーが構成されていることを確認します。本番運用環境では、セキュアMosaic AI Gateway DatabricksAPIキー管理のために、ハードコードされた値の代わりに または シークレット を使用します。

Bash

export OPENAI_API_KEY="your-openai-api-key"

サポートされている API

MLflow は、次の OpenAI APIの自動トレースをサポートしています。

チャット完了	Embeddings	関数呼び出し	構造化出力	ストリーミング	非同期	画像	オーディオ
✅	✅	✅	✅ (※1)	✅ (※2)	✅ (※1)

(*1)ストリーミングのサポートは MLflow 2.15.0 で追加されました。

(*2)サポートされる非同期出力と構造化出力は、MLflow 2.21.0 で追加されました。

追加の のサポートをリクエストするには、 でAPIs 機能リクエスト GitHubを開いてください。

基本的な例

Python

import openai
import mlflow
import os

# Ensure your OPENAI_API_KEY is set in your environment
# os.environ["OPENAI_API_KEY"] = "your-openai-api-key" # Uncomment and set if not globally configured

# Enable auto-tracing for OpenAI
mlflow.openai.autolog()

# Set up MLflow tracking to Databricks
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/openai-tracing-demo")

openai_client = openai.OpenAI()

messages = [
    {
        "role": "user",
        "content": "What is the capital of France?",
    }
]

response = openai_client.chat.completions.create(
    model="gpt-4o-mini",
    messages=messages,
    temperature=0.1,
    max_tokens=100,
)

警告

本番運用環境では、安全な キー管理のために、ハードコードされた値の代わりにMosaic AI Gateway またはDatabricksシークレットAPI を使用します。

秘密のOpenAIの例

シークレットが保存されたら、それを MLflow トレース コードで使用します。

Python

import openai
import mlflow
import os

# Configure your secret scope and key names
secret_scope_name = "openai-secrets"
secret_key_name = "api-key"

# Retrieve the API key from Databricks secrets
os.environ["OPENAI_API_KEY"] = dbutils.secrets.get(
    scope=secret_scope_name,
    key=secret_key_name
)

# Verify the API key was loaded successfully
assert os.environ["OPENAI_API_KEY"] is not None, "API key not loaded from secrets"

# Enable auto-tracing for OpenAI
mlflow.openai.autolog()

# Set up MLflow tracking to Databricks
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/openai-tracing-demo")

# Now you can use OpenAI as usual
openai_client = openai.OpenAI()

messages = [
    {
        "role": "user",
        "content": "What is the capital of France?",
    }
]

response = openai_client.chat.completions.create(
    model="gpt-4o-mini",
    messages=messages,
    temperature=0.1,
    max_tokens=100,
)

ストリーミング

MLflow TracingはAPI OpenAI のストリーミングSDK をサポートしています。自動トレースと同じ設定では、MLflow はストリーミング応答を自動的にトレースし、連結された出力をスパン UI にレンダリングします。レスポンスストリームの実際のチャンクは、「 Event 」タブにも確認できます。

Python

import openai
import mlflow
import os

# Ensure your OPENAI_API_KEY is set in your environment
# os.environ["OPENAI_API_KEY"] = "your-openai-api-key" # Uncomment and set if not globally configured

# Enable trace logging
mlflow.openai.autolog()

# Set up MLflow tracking to Databricks if not already configured
# mlflow.set_tracking_uri("databricks")
# mlflow.set_experiment("/Shared/openai-streaming-demo")

client = openai.OpenAI()

stream = client.chat.completions.create(
    model="gpt-4o-mini",
    messages=[
        {"role": "user", "content": "How fast would a glass of water freeze on Titan?"}
    ],
    stream=True,  # Enable streaming response
)
for chunk in stream:
    print(chunk.choices[0].delta.content or "", end="")

非同期

MLflow TracingAPISDKは、MLflow 2.21.0 以降の OpenAI の非同期 をサポートしています。使用方法は同期 API と同じです。

Python

import openai
import mlflow
import os

# Ensure your OPENAI_API_KEY is set in your environment
# os.environ["OPENAI_API_KEY"] = "your-openai-api-key" # Uncomment and set if not globally configured

# Enable trace logging
mlflow.openai.autolog()

# Set up MLflow tracking to Databricks if not already configured
# mlflow.set_tracking_uri("databricks")
# mlflow.set_experiment("/Shared/openai-async-demo")

client = openai.AsyncOpenAI()

response = await client.chat.completions.create(
    model="gpt-4o-mini",
    messages=[
        {"role": "user", "content": "How fast would a glass of water freeze on Titan?"}
    ],
    # Async streaming is also supported
    # stream=True
)

関数呼び出し

MLflow Tracing は、OpenAI モデルからの関数呼び出し応答を自動的にキャプチャします。 応答内の関数命令は、トレース UI で強調表示されます。さらに、 @mlflow.trace デコレータを使用してツール関数に注釈を付けて、ツール実行のスパンを作成できます。

次の例では、OpenAI Function Calling と OpenAI の MLflow Tracing を使用して、単純な関数呼び出しエージェントを実装します。

Python

import json
from openai import OpenAI
import mlflow
from mlflow.entities import SpanType
import os

# Ensure your OPENAI_API_KEY is set in your environment
# os.environ["OPENAI_API_KEY"] = "your-openai-api-key" # Uncomment and set if not globally configured

# Set up MLflow tracking to Databricks if not already configured
# mlflow.set_tracking_uri("databricks")
# mlflow.set_experiment("/Shared/openai-function-agent-demo")

# Assuming autolog is enabled globally or called earlier
# mlflow.openai.autolog()

client = OpenAI()


# Define the tool function. Decorate it with `@mlflow.trace` to create a span for its execution.
@mlflow.trace(span_type=SpanType.TOOL)
def get_weather(city: str) -> str:
    if city == "Tokyo":
        return "sunny"
    elif city == "Paris":
        return "rainy"
    return "unknown"


tools = [
    {
        "type": "function",
        "function": {
            "name": "get_weather",
            "parameters": {
                "type": "object",
                "properties": {"city": {"type": "string"}},
            },
        },
    }
]

_tool_functions = {"get_weather": get_weather}


# Define a simple tool calling agent
@mlflow.trace(span_type=SpanType.AGENT)
def run_tool_agent(question: str):
    messages = [{"role": "user", "content": question}]

    # Invoke the model with the given question and available tools
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=messages,
        tools=tools,
    )
    ai_msg = response.choices[0].message
    messages.append(ai_msg)

    # If the model request tool call(s), invoke the function with the specified arguments
    if tool_calls := ai_msg.tool_calls:
        for tool_call in tool_calls:
            function_name = tool_call.function.name
            if tool_func := _tool_functions.get(function_name):
                args = json.loads(tool_call.function.arguments)
                tool_result = tool_func(**args)
            else:
                raise RuntimeError("An invalid tool is returned from the assistant!")

            messages.append(
                {
                    "role": "tool",
                    "tool_call_id": tool_call.id,
                    "content": tool_result,
                }
            )

        # Sent the tool results to the model and get a new response
        response = client.chat.completions.create(
            model="gpt-4o-mini", messages=messages
        )

    return response.choices[0].message.content


# Run the tool calling agent
question = "What's the weather like in Paris today?"
answer = run_tool_agent(question)

自動トレースを無効にする

OpenAI の自動トレースは、 mlflow.openai.autolog(disable=True) または mlflow.autolog(disable=True)を呼び出すことでグローバルに無効にできます。