追踪 Mistral

MLflow 追踪功能为您与 Mistral AI 模型的交互提供了可观测性。通过调用 mlflow.mistral.autolog() 函数启用 Mistral 自动追踪后，使用 Mistral SDK 将在交互式开发过程中自动记录生成的追踪信息。

请注意，目前仅支持对文本生成 API 的同步调用，异步 API 和流式方法不会被追踪。

示例用法

import os

from mistralai import Mistral

import mlflow

# Turn on auto tracing for Mistral AI by calling mlflow.mistral.autolog()
mlflow.mistral.autolog()

# Configure your API key.
client = Mistral(api_key=os.environ["MISTRAL_API_KEY"])

# Use the chat complete method to create new chat.
chat_response = client.chat.complete(
    model="mistral-small-latest",
    messages=[
        {
            "role": "user",
            "content": "Who is the best French painter? Answer in one short sentence.",
        },
    ],
)
print(chat_response.choices[0].message)

Token 使用量

MLflow >= 3.2.0 版本支持对 Mistral 的 Token 使用量进行追踪。每次 LLM 调用的 Token 使用量将被记录在 mlflow.chat.tokenUsage 属性中。整个追踪过程中的总 Token 使用量将记录在追踪信息对象的 token_usage 字段中。

import json
import mlflow

mlflow.mistral.autolog()

# Configure your API key.
client = Mistral(api_key=os.environ["MISTRAL_API_KEY"])

# Use the chat complete method to create new chat.
chat_response = client.chat.complete(
    model="mistral-small-latest",
    messages=[
        {
            "role": "user",
            "content": "Who is the best French painter? Answer in one short sentence.",
        },
    ],
)

# 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']}")

== Total token usage: ==
  Input tokens: 16
  Output tokens: 25
  Total tokens: 41

== Detailed usage for each LLM call: ==
Chat.complete:
  Input tokens: 16
  Output tokens: 25
  Total tokens: 41

禁用自动跟踪

可以通过调用 mlflow.mistral.autolog(disable=True) 或 mlflow.autolog(disable=True) 来全局禁用对 Mistral 的自动追踪。