メインコンテンツまでスキップ

AG2 のトレース

MLflow Tracing AG2 (旧 AutoGen 0.2) と統合して、マルチエージェントの会話とワークフローの統合トレースをキャプチャします。この統合により、エージェントループとツール実行が自動的にインストルメント化され、mlflow.ag2.autologを呼び出すだけです。

Python
import mlflow

mlflow.ag2.autolog()

この統合により、次の包括的な可視性が提供されます。

  • どのエージェントが異なるターンで呼び出されるか
  • エージェント間で渡されるメッセージ
  • 各エージェントによって行われるLLMとツール呼び出しは、エージェントとターンごとに編成されています
  • 待ち時間
  • 発生した場合の例外

前提 条件​

AG2 で MLflow Tracing を使用するには、 MLflow と関連する AG2 (AutoGen) パッケージをインストールする必要があります。

開発環境の場合は、Databricks エクストラと AutoGen を含む完全な MLflow パッケージをインストールします。

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

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

注記

MLflow 3 は、AG2 で最適なトレース エクスペリエンスを実現するために推奨されます。

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

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

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

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

基本的な例​

Python
import os
from typing import Annotated, Literal

from autogen import ConversableAgent
import mlflow

# Enable auto-tracing for AG2
mlflow.ag2.autolog()

# Track to Databricks (optional if already configured)
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/ag2-tracing-demo")

# Define a simple multi-agent workflow using AG2 (AutoGen 0.2)
config_list = [
    {
        "model": "gpt-4o-mini",
        # Requires OPENAI_API_KEY in env for this example
        "api_key": os.environ.get("OPENAI_API_KEY"),
    }
]

Operator = Literal["+", "-", "*", "/"]


def calculator(a: int, b: int, operator: Annotated[Operator, "operator"]) -> int:
    if operator == "+":
        return a + b
    elif operator == "-":
        return a - b
    elif operator == "*":
        return a * b
    elif operator == "/":
        return int(a / b)
    else:
        raise ValueError("Invalid operator")


assistant = ConversableAgent(
    name="Assistant",
    system_message=(
        "You are a helpful AI assistant. You can help with simple calculations. "
        "Return 'TERMINATE' when the task is done."
    ),
    llm_config={"config_list": config_list},
)

user_proxy = ConversableAgent(
    name="Tool Agent",
    llm_config=False,
    is_termination_msg=lambda msg: msg.get("content") is not None and "TERMINATE" in msg["content"],
    human_input_mode="NEVER",
)

assistant.register_for_llm(name="calculator", description="A simple calculator")(calculator)
user_proxy.register_for_execution(name="calculator")(calculator)

response = user_proxy.initiate_chat(
    assistant, message="What is (44231 + 13312 / (230 - 20)) * 4?"
)

トークンの使用状況を追跡する​

MLflow 3.2.0+AG2 のトークン使用状況の追跡をサポートします。通話ごとの使用状況は、 mlflow.chat.tokenUsage span 属性に記録されます。合計使用量がトレース情報に表示されます。

Python
import mlflow

last_trace_id = mlflow.get_last_active_trace_id()
trace = mlflow.get_trace(trace_id=last_trace_id)

total = trace.info.token_usage
print("Input:", total["input_tokens"], "Output:", total["output_tokens"], "Total:", total["total_tokens"])

for span in trace.data.spans:
    usage = span.get_attribute("mlflow.chat.tokenUsage")
    if usage:
        print(span.name, usage)

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

mlflow.ag2.autolog(disable=True)でAG2自動トレースを無効にするか、mlflow.autolog(disable=True)ですべての自動ログを無効にします。

に最終更新