#使用 MLflow 跟踪 AI 智能体

MLflow Tracing 会记录 AI 智能体执行的步骤，包括模型请求和响应、链调用和工具调用、检索步骤、延迟、令牌使用量以及错误。在 Alauda AI 上，trace 会存储在所选的 MLflow experiment 中，并受到与 runs 和 models 相同的 workspace 身份验证和 Kubernetes RBAC 保护。

本指南介绍以下场景的推荐 tracing 路径：

• 使用 LangChain 构建的 Python 智能体。
• 通过 OpenAI 兼容 API 调用 Llama Stack，现称 OGX 的应用程序。
• 已经发出 OpenTelemetry traces 的 OGX 服务器或其他应用程序。

#前提条件

• 一个 MLflow tracking server 和一个你可以访问的 workspace。请参阅 MLflow Workspaces 和访问控制。
• MLFLOW_TRACKING_TOKEN 中的 Dex id token。请按照 将 MLflow Python SDK 与身份验证和 RBAC 一起使用 获取并续期 token。
• 用于 traces 的 experiment。创建或写入 traces 需要对 workspace 的 experiments 资源具有 create 和 update 权限。
• 智能体到 MLflow tracking server 及其模型 endpoint 的网络访问。

在启动智能体之前，先设置通用的 MLflow 连接变量：

export MLFLOW_TRACKING_URI=http://mlflow-tracking-server.kubeflow:5000
export MLFLOW_TRACKING_TOKEN='<dex-id-token>'
export MLFLOW_WORKSPACE=team-a

上面的 URI 是集群内 Service，由 OAuth proxy 提供前置代理。对于集群外的客户端，请使用 https://<platform>/clusters/<cluster>/mlflow。

#跟踪 LangChain 智能体

MLflow 的 LangChain 集成会自动为智能体、嵌套链、模型请求、检索和工具调用创建 span。它支持同步、异步、批处理和流式 LangChain 调用。

安装智能体依赖项：

pip install "mlflow[genai]" langchain langchain-openai

下面的示例使用工具调用智能体和一个与 OpenAI 兼容的模型 endpoint。为你环境中可用的模型服务设置 MODEL_BASE_URL、MODEL_API_KEY 和 MODEL_ID。

import os

import mlflow
from langchain.agents import create_agent
from langchain.tools import tool
from langchain_openai import ChatOpenAI


mlflow.set_experiment("agent-tracing")
mlflow.langchain.autolog()


@tool
def get_weather(city: str) -> str:
    """Return the current weather for a city."""
    # Replace this example result with a call to your weather service.
    return f"The weather in {city} is sunny and 24 C."


model = ChatOpenAI(
    model=os.environ["MODEL_ID"],
    base_url=os.environ["MODEL_BASE_URL"],
    api_key=os.environ["MODEL_API_KEY"],
)

agent = create_agent(
    model=model,
    tools=[get_weather],
    system_prompt="You are a concise travel assistant.",
)

result = agent.invoke(
    {"messages": [{"role": "user", "content": "Should I take a coat to Taipei today?"}]}
)
print(result["messages"][-1].content)

在应用启动期间、调用智能体之前执行一次 mlflow.langchain.autolog()。现有的 LangChain 应用代码无需其他更改。

#跟踪 OGX 应用程序

OGX 提供与 OpenAI 兼容的 Chat Completions 和 Responses API。对于调用 OGX 的 Python 应用程序，MLflow 的 OpenAI 集成是最简单的选择，并且可以完整跟踪 Responses API 调用。

安装客户端依赖项：

pip install "mlflow[genai]" openai

将 OpenAI 客户端指向 OGX API。@mlflow.trace 包装器会将 run_agent 执行的所有模型和工具活动归入一个父 span；mlflow.openai.autolog() 会为每个 OGX 请求创建嵌套 span。

import os

import mlflow
from openai import OpenAI


mlflow.set_experiment("ogx-agent-tracing")
mlflow.openai.autolog()

client = OpenAI(
    base_url=os.environ.get(
        "OGX_BASE_URL",
        "http://<ogx-service>.<namespace>.svc.cluster.local:8321/v1",
    ),
    api_key=os.environ.get("OGX_API_KEY", "not-used"),
)


@mlflow.trace(name="ogx-agent")
def run_agent(question: str) -> str:
    response = client.responses.create(
        model=os.environ["OGX_MODEL_ID"],
        input=question,
    )
    return response.output_text


print(run_agent("Give me a two-step plan for testing a Kubernetes operator."))

在多步骤智能体循环中发出的每次调用，都会嵌套在 ogx-agent trace 之下。当天请求使用函数工具时，MLflow 还会记录工具定义以及模型的工具调用响应。

#导出 OGX OpenTelemetry traces

当你希望跟踪 OGX 服务器本身、为非 Python 应用程序添加埋点，或通过 collector 路由 telemetry 时，请使用 OpenTelemetry。MLflow 通过 OTLP/HTTP 在 /v1/traces 接收 traces；它不接受 OTLP/gRPC 摄取。

首先，获取数值形式的 experiment ID：

import mlflow

experiment = mlflow.set_experiment("ogx-server-tracing")
print(experiment.experiment_id)

配置 OpenTelemetry Collector，以接收来自 OGX 的 traces 并将其转发到 Alauda MLflow：

receivers:
  otlp:
    protocols:
      http:
        endpoint: 0.0.0.0:4318

processors:
  batch: {}

exporters:
  otlphttp/mlflow:
    traces_endpoint: http://mlflow-tracking-server.kubeflow:5000/v1/traces
    headers:
      Authorization: "Bearer ${env:MLFLOW_TRACKING_TOKEN}"
      X-MLflow-Experiment-Id: "${env:MLFLOW_EXPERIMENT_ID}"
      X-MLFLOW-WORKSPACE: "${env:MLFLOW_WORKSPACE}"

service:
  pipelines:
    traces:
      receivers: [otlp]
      processors: [batch]
      exporters: [otlphttp/mlflow]

使用环境变量 MLFLOW_TRACKING_TOKEN、MLFLOW_EXPERIMENT_ID 和 MLFLOW_WORKSPACE 启动 collector。然后配置 OGX，将 OTLP/HTTP traces 发送到 collector：

export OTEL_EXPORTER_OTLP_ENDPOINT=http://otel-collector:4318
export OTEL_EXPORTER_OTLP_PROTOCOL=http/protobuf

opentelemetry-instrument \
  --traces_exporter otlp \
  --metrics_exporter none \
  --service_name ogx-server \
  -- \
  ogx run starter

OGX 使用 OpenTelemetry 自动埋点来采集服务器和受支持的 SDK 活动。MLflow 的 OTLP endpoint 只接收 traces；metrics 请单独路由到 metrics backend。

#查看 traces

1. 打开 Alauda AI → Tools → MLFlow。
2. 选择与应用程序或 collector 配置相同的 workspace 和 experiment。
3. 打开该 experiment 的 Traces 选项卡。
4. 选择一个 trace，查看其 span 层次结构、输入和输出、延迟、令牌使用量、工具调用以及错误。

#故障排查