MLflow TracingAIエージェントに を追加する

MLflow Tracing on Ed Databricks では、生成AI エージェントの動作をモニタリングおよび分析できます。 トレースは、次の方法で有効にできます。

• 自動トレース:1行のコードでエージェントをトレースします。
• 手動トレース:エージェントをより詳細にインストゥルメントするか、自動トレースではカバーできないプロパティをトレースします。
• 自動トレースと手動トレースを組み合わせて、柔軟で便利なトレースを実現します。

自動トレース

MLflow の自動ログ 記録を使用すると、エージェントのトレースを簡単に開始できます。コードに次の行を追加するだけです。

Python

mlflow.<library>.autolog()

MLflow では、次のような多くの一般的なエージェント作成ライブラリの自動ログ記録がサポートされています。

ライブラリ	自動ログのバージョンのサポート	自動ロギングコマンド
LangChain	0.1.0~最新	mlflow.langchain.autolog()
Langgraph	0.1.1 ~ 最新	mlflow.langgraph.autolog()
OpenAI	1.0.0 ~ 最新	mlflow.openai.autolog()
OpenAIエージェントSDK	0.0.7 ~ 最新	mlflow.openai.autolog()
LlamaIndex	0.10.44 ~ 最新	mlflow.llama_index.autolog()
DSPy	2.5.17 ~ 最新	mlflow.dspy.autolog()
Amazon Bedrock	1.33.0 ~ 最新 (boto3)	mlflow.bedrock.autolog()
Anthropic	0.30.0 ~ 最新	mlflow.anthropic.autolog()
AutoGen	0.2.36 ~ 0.2.40	mlflow.autogen.autolog()
Google Gemini	1.0.0 ~ 最新	mlflow.gemini.autolog()
CrewAI	0.80.0 ~ 最新	mlflow.crewai.autolog()
LiteLLM	1.52.9 ~ 最新	mlflow.litellm.autolog()
Groq	0.13.0 ~ 最新	mlflow.groq.autolog()
Mistral	1.0.0 ~ 最新	mlflow.mistral.autolog()

サポートされているライブラリの完全なリストと詳細については、自動ログMLflowドキュメントを参照してください。

自動ログを無効にする

自動ロギング トレースは、 Databricks Runtime 15.4 ML 以降では、 LangChain、Langgraph、OpenAI、および LlamaIndex のデフォルトによって有効になります。 自動ログを無効にするには、ノートブックで次のコマンドを実行します。

Python

mlflow.<library>.autolog(log_traces=False)

手動トレース

MLflow Tracing APIsでは、エージェントをより詳細にインストゥルメント化する場合や、自動ロギングでキャプチャされないトレースを追加する場合に、手動でトレースを追加できます。

MLflow Tracing APIs 、トレースのツリー構造を管理せずにトレースを追加するためのローコード APIs です。 MLflow は、Python スタックを使用して、適切な親子スパン関係を自動的に決定します。

@mlflow.traceデコレータを使用したトレース関数

コードを手動でインストゥルメントする最も簡単な方法は、 @mlflow.trace デコレータで関数をデコレートすることです。MLflow トレース デコレーターは、デコレートされた関数のスコープで スパン を作成します。スパンには、入力、出力、遅延、および例外が記録されます。

たとえば、次のコードでは、入力引数 x と y と出力をキャプチャする my_function という名前のスパンを作成します。

Python

import mlflow

@mlflow.trace
def add(x: int, y: int) -> int:
  return x + y

スパンの名前、タイプをカスタマイズしたり、カスタム属性を追加したりできます。

Python

from mlflow.entities import SpanType

@mlflow.trace(
  # By default, the function name is used as the span name. You can override it with the `name` parameter.
  name="my_add_function",
  # Specify the span type using the `span_type` parameter.
  span_type=SpanType.TOOL,
  # Add custom attributes to the span using the `attributes` parameter. By default, MLflow only captures input and output.
  attributes={"key": "value"}
)
def add(x: int, y: int) -> int:
  return x + y

コンテキストマネージャーを使用して任意のコードブロックをトレース

関数だけでなく、任意のコードブロックの範囲を作成するには、コードブロックをラップするコンテキストマネージャーとして mlflow.start_span() を使用します。

スパンは、コンテキストが入力されたときに開始され、コンテキストが終了したときに終了します。span の入力と出力は、コンテキスト マネージャーによって生成される span オブジェクトの setter メソッドを使用して手動で提供する必要があります。詳細については、MLflowドキュメント - コンテキストハンドラーを参照してください。

Python

with mlflow.start_span(name="my_span") as span:
  span.set_inputs({"x": x, "y": y})
  result = x + y
  span.set_outputs(result)
  span.set_attribute("key", "value")

下位レベルのトレース・ライブラリ

MLflow には、トレース ツリー構造を明示的に制御するための低レベルの APIs も用意されています。 MLflow のドキュメント - 手動インストルメンテーションを参照してください。

自動ログ記録と手動トレースの組み合わせ

手動トレースと自動ロギングを組み合わせることができます。MLflow は、両方の種類のスパンを 1 つの完全なトレースにマージします。

次の例では、OpenAI の自動ログ記録と手動トレースを組み合わせています。

Python

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

client = OpenAI()

# Enable OpenAI autologging to capture LLM API calls
# (*Not necessary if you are using the Databricks Runtime 15.4 ML and above, where OpenAI autologging is enabled by default.)
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="gpt-4o-mini",
    messages=messages,
    tools=tools,
  )
  ai_msg = response.choices[0].message
  messages.append(ai_msg)

  # If the model requests tool calls, invoke the function(s) 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,
        }
      )

    # Send 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)

トレースのオーバーヘッドレイテンシ

トレースは、パフォーマンスへの影響を最小限に抑えるために非同期的に書き込まれます。ただし、トレースは、特にトレースが大きい場合に、エンドポイントの応答速度に遅延を追加します。エンドポイントをテストして影響を理解してから、本番運用にデプロイします。

次の表は、トレース サイズによる待機時間の影響を推定しています。

要求あたりのトレース サイズ	応答速度の遅延 (ミリ秒) への影響
~10キロバイト	~1ミリ秒
~1メガバイト	50 ~ 100 ミリ秒
10MB	150ミリ秒~

トラブルシューティング

トラブルシューティングと一般的な質問については、 MLflow のドキュメント: トレース攻略ガイド と MLflow のドキュメント: FAQ を参照してください。