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

Smolagentsの追跡

Smolagents は、ミニマリズムとコンポーザビリティを重視した軽量のエージェント フレームワークです。

MLflow Tracing は Smolagents と統合して、軽量エージェント ワークフローの合理化されたトレースをキャプチャします。mlflow.smolagents.autologで有効にします。

注記

MLflow 自動トレースでは、同期呼び出しのみがサポートされます。非同期 API メソッドとストリーミング メソッドはトレースされません。

前提 条件

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

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

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

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

注記

最適なトレース エクスペリエンスを得るには、MLflow 3 をお勧めします。

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

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

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

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

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

Bash
export OPENAI_API_KEY="your-openai-api-key"
# Add other provider keys as needed

使用例

Python
import mlflow

mlflow.smolagents.autolog()

from smolagents import CodeAgent, LiteLLMModel
import mlflow

# Turn on auto tracing for Smolagents by calling mlflow.smolagents.autolog()
mlflow.smolagents.autolog()

model = LiteLLMModel(model_id="openai/gpt-4o-mini", api_key=API_KEY)
agent = CodeAgent(tools=[], model=model, add_base_tools=True)

result = agent.run(
"Could you give me the 118th number in the Fibonacci sequence?",
)

実行 Smolagents ワークフロー 通常どおり。 トレースはエクスペリメントUIに表示されます。

トークン追跡の使用状況

MLflow は、 mlflow.chat.tokenUsage 属性への各エージェント呼び出しのトークン使用状況をログに記録します。トレース全体のトークン使用量の合計は、トレース情報オブジェクトの token_usage フィールドで確認できます。

Python
import json
import mlflow

mlflow.smolagents.autolog()

model = LiteLLMModel(model_id="openai/gpt-4o-mini", api_key=API_KEY)
agent = CodeAgent(tools=[], model=model, add_base_tools=True)

result = agent.run(
"Could you give me the 118th number in the Fibonacci sequence?",
)

# Get the trace object just created
last_trace_id = mlflow.get_last_active_trace_id()
trace = mlflow.get_trace(trace_id=last_trace_id)

# Print the token usage
total_usage = trace.info.token_usage
print("== Total token usage: ==")
print(f" Input tokens: {total_usage['input_tokens']}")
print(f" Output tokens: {total_usage['output_tokens']}")
print(f" Total tokens: {total_usage['total_tokens']}")

# Print the token usage for each LLM call
print("\n== Detailed usage for each LLM call: ==")
for span in trace.data.spans:
if usage := span.get_attribute("mlflow.chat.tokenUsage"):
print(f"{span.name}:")
print(f" Input tokens: {usage['input_tokens']}")
print(f" Output tokens: {usage['output_tokens']}")
print(f" Total tokens: {usage['total_tokens']}")

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

mlflow.smolagents.autolog(disable=True)で無効にするか、mlflow.autolog(disable=True)でグローバルに無効にします。