跟踪 Anthropic

MLflow 跟踪为 Anthropic LLM 提供自动跟踪功能。 通过调用 mlflow.anthropic.autolog 函数为 Anthropic 启用自动跟踪，MLflow 将捕获嵌套跟踪，并在调用 Anthropic Python SDK 时将其记录在活跃的 MLflow 实验中。

import mlflow

mlflow.anthropic.autolog()

MLflow 跟踪会自动捕获有关 Anthropic 调用的以下信息：

• 提示和完成响应
• 潜伏期
• 模型名称
• 其他元数据（例如 temperature， max_tokens如果指定）。
• 在响应中返回时进行函数调用
• 引发的任何异常

先决条件

若要将 MLflow 跟踪与人类学配合使用，需要安装 MLflow 和 Anthropic SDK。

开发

对于开发环境，请安装包含 Databricks 附加程序和 anthropic 的完整 MLflow 软件包：

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

完整 mlflow[databricks] 包包括用于 Databricks 的本地开发和试验的所有功能。

生产

对于生产部署，请安装 mlflow-tracing 和 anthropic：

pip install --upgrade mlflow-tracing anthropic

包 mlflow-tracing 已针对生产用途进行优化。

在运行以下示例之前，需要配置环境：

对于不使用 Databricks 笔记本的用户：设置 Databricks 环境变量：

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

对于 Databricks 笔记本中的用户：这些凭据会自动为您设置。

API 密钥：确保配置人类 API 密钥。 对于生产用途，我们建议使用 Mosaic AI 网关或 Databricks 密钥，而不是环境变量：

export ANTHROPIC_API_KEY="your-anthropic-api-key"

受支持的 API

MLflow 支持以下人类 API 的自动跟踪：

(*1)MLflow 2.21.0 中添加了异步支持。

若要请求对其他 API 的支持，请在 GitHub 上打开 功能请求 。

基本示例

import anthropic
import mlflow
import os

# Ensure your ANTHROPIC_API_KEY is set in your environment
# os.environ["ANTHROPIC_API_KEY"] = "your-anthropic-api-key" # Uncomment and set if not globally configured

# Enable auto-tracing for Anthropic
mlflow.anthropic.autolog()

# Set up MLflow tracking to Databricks
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/anthropic-tracing-demo")

# Configure your API key.
client = anthropic.Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])

# Use the create method to create new message.
message = client.messages.create(
    model="claude-3-5-sonnet-20241022",
    max_tokens=1024,
    messages=[
        {"role": "user", "content": "Hello, Claude"},
    ],
)

异步

import anthropic
import mlflow
import os

# Ensure your ANTHROPIC_API_KEY is set in your environment
# os.environ["ANTHROPIC_API_KEY"] = "your-anthropic-api-key" # Uncomment and set if not globally configured

# Enable trace logging
mlflow.anthropic.autolog()

# Set up MLflow tracking to Databricks if not already configured
# mlflow.set_tracking_uri("databricks")
# mlflow.set_experiment("/Shared/anthropic-async-demo")

client = anthropic.AsyncAnthropic()

response = await client.messages.create(
    model="claude-3-5-sonnet-20241022",
    max_tokens=1024,
    messages=[
        {"role": "user", "content": "Hello, Claude"},
    ],
)

高级示例：工具呼叫代理

MLflow 跟踪自动从 Anthropic 模型中捕获工具调用响应。 响应中的函数指令将在跟踪 UI 中突出显示。 此外，可以使用@mlflow.trace修饰器为工具函数添加批注，以便为工具执行创建跨度。

以下示例使用 Anthropic 工具调用和针对 Anthropic 的 MLflow 跟踪实现简单的函数调用代理。 该示例进一步使用异步人类 SDK，以便代理可以在不阻止的情况下处理并发调用。

import json
import anthropic
import mlflow
import asyncio
from mlflow.entities import SpanType
import os

# Ensure your ANTHROPIC_API_KEY is set in your environment
# os.environ["ANTHROPIC_API_KEY"] = "your-anthropic-api-key" # Uncomment and set if not globally configured

# Set up MLflow tracking to Databricks if not already configured
# mlflow.set_tracking_uri("databricks")
# mlflow.set_experiment("/Shared/anthropic-tool-agent-demo")

# Assuming autolog is enabled globally or called earlier
# mlflow.anthropic.autolog()

client = anthropic.AsyncAnthropic()
model_name = "claude-3-5-sonnet-20241022"

# Define the tool function. Decorate it with `@mlflow.trace` to create a span for its execution.
@mlflow.trace(span_type=SpanType.TOOL)
async def get_weather(city: str) -> str:
    if city == "Tokyo":
        return "sunny"
    elif city == "Paris":
        return "rainy"
    return "unknown"

tools = [
    {
        "name": "get_weather",
        "description": "Returns the weather condition of a given city.",
        "input_schema": {
            "type": "object",
            "properties": {"city": {"type": "string"}},
            "required": ["city"],
        },
    }
]

_tool_functions = {"get_weather": get_weather}

# Define a simple tool calling agent
@mlflow.trace(span_type=SpanType.AGENT)
async def run_tool_agent(question: str):
    messages = [{"role": "user", "content": question}]

    # Invoke the model with the given question and available tools
    ai_msg = await client.messages.create(
        model=model_name,
        messages=messages,
        tools=tools,
        max_tokens=2048,
    )
    messages.append({"role": "assistant", "content": ai_msg.content})

    # If the model requests tool call(s), invoke the function with the specified arguments
    tool_calls = [c for c in ai_msg.content if c.type == "tool_use"]
    for tool_call in tool_calls:
        if tool_func := _tool_functions.get(tool_call.name):
            tool_result = await tool_func(**tool_call.input)
        else:
            raise RuntimeError("An invalid tool is returned from the assistant!")

        messages.append(
            {
                "role": "user",
                "content": [
                    {
                        "type": "tool_result",
                        "tool_use_id": tool_call.id,
                        "content": tool_result,
                    }
                ],
            }
        )

    # Send the tool results to the model and get a new response
    response = await client.messages.create(
        model=model_name,
        messages=messages,
        max_tokens=2048,
    )

    return response.content[-1].text

# Run the tool calling agent
cities = ["Tokyo", "Paris", "Sydney"]
questions = [f"What's the weather like in {city} today?" for city in cities]
answers = await asyncio.gather(*(run_tool_agent(q) for q in questions))

for city, answer in zip(cities, answers):
    print(f"{city}: {answer}")

禁用自动跟踪

可以通过调用 mlflow.anthropic.autolog(disable=True) 或 mlflow.autolog(disable=True) 全局禁用 Anthropic 的自动跟踪。