追踪 Google Agent Development Kit (ADK)

MLflow Tracing 为 Google ADK 提供了自动追踪功能,Google ADK 是 Google 开发的一个灵活且模块化的 AI 代理框架。MLflow 通过 OpenTelemetry 集成支持 Google ADK 的追踪。
步骤 1:安装库
bash
pip install mlflow>=3.6.0 google-adk opentelemetry-exporter-otlp-proto-http
步骤 2:启动 MLflow Tracking Server
使用基于 SQL 的后端存储启动 MLflow Tracking Server
bash
mlflow server --backend-store-uri sqlite:///mlflow.db --port 5000
本示例使用 SQLite 作为后端存储。要使用其他类型的 SQL 数据库(如 PostgreSQL、MySQL 和 MSSQL),请按照 后端存储文档 中的说明更改存储 URI。文件型后端存储不支持 OpenTelemetry 摄取。
步骤 3:配置 OpenTelemetry
配置 OpenTelemetry Tracer 以将追踪导出到 MLflow Tracking Server 端点。
- 将端点设置为 MLflow Tracking Server 的
/v1/traces端点 (OTLP)。 - 将
x-mlflow-experiment-id标头设置为 MLflow 实验 ID。如果您没有实验 ID,请通过 Python SDK 或 MLflow UI 创建一个。
bash
export OTEL_EXPORTER_OTLP_ENDPOINT=https://:5000
export OTEL_EXPORTER_OTLP_HEADERS=x-mlflow-experiment-id=123
步骤 4:运行代理
像往常一样在 my_agent/agent.py 等 Python 脚本中定义和调用代理。Google ADK 将为您的代理生成追踪,并将其发送到 MLflow Tracking Server 端点。要为 Google ADK 启用追踪并将追踪发送到 MLflow,请在运行代理之前使用 OTLPSpanExporter 设置 OpenTelemetry Tracer Provider。
python
# my_agent/agent.py
from google.adk.agents import LlmAgent
from google.adk.tools import FunctionTool
from opentelemetry import trace
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
# Configure the tracer provider and add the exporter
tracer_provider = TracerProvider()
tracer_provider.add_span_processor(SimpleSpanProcessor(OTLPSpanExporter()))
trace.set_tracer_provider(tracer_provider)
def calculator(a: float, b: float) -> str:
"""Add two numbers and return the result.
Args:
a: First number
b: Second number
Returns:
The sum of a and b
"""
return str(a + b)
calculator_tool = FunctionTool(func=calculator)
root_agent = LlmAgent(
name="MathAgent",
model="gemini-2.0-flash-exp",
instruction=(
"You are a helpful assistant that can do math. "
"When asked a math problem, use the calculator tool to solve it."
),
tools=[calculator_tool],
)
使用 adk run 命令或 Web UI 运行代理。
bash
adk run my_agent
在 https://:5000 打开 MLflow UI,然后导航到实验以查看追踪。