DeepSeek のトレース
MLflow Tracing は、OpenAI SDK の統合により、Deepseek モデルの自動トレース機能を提供します。DeepSeek は OpenAI 互換の API 形式を使用するため、 mlflow.openai.autolog()
を使用して DeepSeek モデルとのインタラクションを追跡できます。
import mlflow
mlflow.openai.autolog()
MLflow トレースは、DeepSeek 呼び出しに関する次の情報を自動的にキャプチャします。
- プロンプトと完了応答
- 待ち時間
- モデル名
temperature
、max_tokens
などの追加のメタデータ (指定されている場合)。- 応答で返された場合の関数呼び出し
- 例外が発生した場合
前提 条件
DeepSeekで MLflow Tracing を使用するには(OpenAI互換の APIを使用)、 MLflow とOpenAI SDKをインストールする必要があります。
- Development
- Production
開発環境の場合は、Databricks の追加機能と openai
を含む完全な MLflow パッケージをインストールします。
pip install --upgrade "mlflow[databricks]>=3.1" openai
フル mlflow[databricks]
パッケージには、Databricks でのローカル開発と実験のためのすべての機能が含まれています。
本番運用デプロイメントの場合は、 mlflow-tracing
と openai
をインストールします。
pip install --upgrade mlflow-tracing openai
mlflow-tracing
パッケージは、本番運用での使用に最適化されています。
最適なトレース エクスペリエンスを得るには、MLflow 3 を強くお勧めします。
例を実行する前に、環境を構成する必要があります。
Databricks ノートブックの外部ユーザーの場合 : Databricks 環境変数を設定します。
export DATABRICKS_HOST="https://your-workspace.cloud.databricks.com"
export DATABRICKS_TOKEN="your-personal-access-token"
Databricks ノートブック内のユーザーの場合 : これらの資格情報は自動的に設定されます。
API キー : DeepSeek API キーが設定されていることを確認します。本番運用で使用する場合は、ハードコードされた値の代わりに Mosaic AI Gateway または Databricks シークレットを使用します。
export DEEPSEEK_API_KEY="your-deepseek-api-key"
サポートされている APIs
MLflow は、OpenAI 統合を通じて、次の DeepSeek API の自動トレースをサポートしています。
チャット完了 | 関数呼び出し | ストリーミング | 非同期 |
---|---|---|---|
✅ | ✅ | ✅ (※1) | ✅ (※2) |
(*1)ストリーミングのサポートには、MLflow 2.15.0 以降が必要です。(*2)非同期サポートには、MLflow 2.21.0 以降が必要です。
追加のAPIのサポートをリクエストするには、GitHub で機能リクエスト を開いてください。
基本的な例
import openai
import mlflow
import os
# Ensure your DEEPSEEK_API_KEY is set in your environment
# os.environ["DEEPSEEK_API_KEY"] = "your-deepseek-api-key" # Uncomment and set if not globally configured
# Enable auto-tracing for OpenAI (works with DeepSeek)
mlflow.openai.autolog()
# Set up MLflow tracking to Databricks
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/deepseek-demo")
# Initialize the OpenAI client with DeepSeek API endpoint and your key
client = openai.OpenAI(
base_url="https://api.deepseek.com",
api_key=os.environ.get("DEEPSEEK_API_KEY") # Or directly pass your key string
)
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "What is the capital of France?"},
]
response = client.chat.completions.create(
model="deepseek-chat",
messages=messages,
temperature=0.1,
max_tokens=100,
)
上記の例では、 MLflow UI のエクスペリメントにトレースを生成する必要があります。
本番運用環境では、安全な キー管理のために、ハードコードされた値の代わりにMosaic AI Gateway またはDatabricksシークレットAPI を使用します。
ストリーミングと非同期のサポート
MLflow は、ストリーミングおよび非同期 DeepSeek APIsのトレースをサポートしています。 OpenAI SDK を使用してストリーミングと非同期呼び出しをトレースするためのコード スニペットの例については、 OpenAI トレースのドキュメント を参照してください。
高度な例: 関数呼び出しエージェント
MLflow Tracing は、OpenAI SDKを介して DeepSeek モデルからの関数呼び出し応答を自動的にキャプチャします。 応答内の関数命令は、トレース UI で強調表示されます。さらに、 @mlflow.trace
デコレータを使用してツール関数に注釈を付けて、ツール実行のスパンを作成できます。
次の例では、DeepSeek Function Calling と MLflow Tracingを使用して、単純な関数呼び出しエージェントを実装します。
import json
from openai import OpenAI
import mlflow
from mlflow.entities import SpanType
import os
# Ensure your DEEPSEEK_API_KEY is set in your environment
# os.environ["DEEPSEEK_API_KEY"] = "your-deepseek-api-key" # Uncomment and set if not globally configured
# Initialize the OpenAI client with DeepSeek API endpoint and your key
client = OpenAI(
base_url="https://api.deepseek.com",
api_key=os.environ.get("DEEPSEEK_API_KEY") # Or directly pass your key string
)
# Set up MLflow tracking to Databricks if not already configured
# mlflow.set_tracking_uri("databricks")
# mlflow.set_experiment("/Shared/deepseek-agent-demo")
# Assuming autolog is enabled globally or called earlier
# mlflow.openai.autolog()
# 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="deepseek-chat",
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="deepseek-chat", 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)
自動トレースを無効にする
DeepSeek の自動トレース (OpenAI SDK 経由) は、 mlflow.openai.autolog(disable=True)
または mlflow.autolog(disable=True)
を呼び出すことでグローバルに無効にできます。