超越 Autolog：为新 LLM 提供商添加 MLflow 追踪
在本文中，我们将展示如何通过为 Ollama Python SDK 的 chat 方法添加追踪支持，为新的 LLM 提供商添加 MLflow 追踪。
MLflow 追踪是 MLflow 中的一个可观测性工具，用于捕获 GenAI 应用程序和工作流的详细执行追踪。除了捕获单个调用的输入、输出和元数据外，MLflow 追踪还可以捕获中间步骤，例如工具调用、推理步骤、检索步骤或其他自定义步骤。
MLflow 为许多流行的 LLM 提供商和编排框架提供了内置的追踪支持。如果您使用的是这些提供商之一，只需一行代码即可启用追踪：mlflow.<provider>.autolog()。虽然 MLflow 的自动日志记录功能涵盖了许多最广泛使用的 LLM 提供商和编排框架，但有时您可能需要为不受支持的提供商添加追踪，或自定义追踪以超出自动日志记录的范围。本文演示了 MLflow 追踪的灵活性和可扩展性，方法是
- 为不受支持的提供商（Ollama Python SDK）添加基本追踪支持
- 展示如何捕获简单的补全和更复杂的工具调用工作流
- 说明如何在对现有代码进行最少更改的情况下添加追踪
我们将以 Ollama Python SDK 为例，它是 Ollama LLM 平台的开源 Python SDK。我们将逐步完成此过程，展示如何在保持与提供商 SDK 干净集成的同时，使用 MLflow 追踪捕获关键信息。请注意，MLflow 确实对 Ollama 具有自动日志记录支持，但目前仅限于通过 OpenAI 客户端使用，而不是直接与 Ollama Python SDK 一起使用。
为新提供商添加 MLflow 追踪：一般原则
MLflow 文档中有一份出色的指南，介绍了如何为 MLflow 追踪做贡献。尽管在此示例中我们不会为 MLflow 本身做贡献，但我们将遵循相同的总体原则。
本文假设您对 MLflow 追踪是什么及其工作原理有基本的了解。如果您是初学者或需要复习，请查阅追踪概念指南。
向新提供商添加追踪涉及几个关键考虑因素
-
了解提供商的关键功能：我们首先需要了解需要追踪哪些 API 方法才能获得我们想要的追踪信息。对于 LLM 推理提供商，这通常涉及聊天补全、工具调用或嵌入生成等操作。在编排框架中，这可能涉及检索、推理、路由或各种自定义步骤等操作。在我们的 Ollama 示例中，我们将重点关注聊天补全 API。此步骤将根据提供商的不同而有很大差异。
-
将操作映射到跨度：MLflow 追踪使用不同的跨度类型来表示不同类型的操作。您可以在此处找到内置跨度类型的描述。不同类型的跨度在 MLflow UI 中显示方式不同，并且可以启用特定功能。在跨度内部，我们还希望将提供商的输入和输出映射到 MLflow 所期望的格式。MLflow 提供用于记录聊天和工具输入和输出的实用程序，这些内容随后会作为格式化消息显示在 MLflow UI 中。
向新提供商添加追踪时，我们的主要任务是将提供商的 API 方法映射到具有适当跨度类型的 MLflow 追踪跨度。
-
结构化和保留关键数据：对于我们想要追踪的每项操作，我们需要确定想要保留的关键信息，并确保以有用的方式捕获和显示这些信息。例如，我们可能希望捕获控制操作行为的输入和配置数据、解释结果的输出和元数据、过早终止操作的错误等。查看类似提供商的追踪和追踪实现可以为如何结构化和保留这些数据提供一个良好的起点。
向 Ollama Python SDK 添加追踪
现在我们对向新提供商添加追踪的关键步骤有了大致了解，让我们逐步完成此过程，并向 Ollama Python SDK 添加追踪。
步骤 1：安装和测试 Ollama Python SDK
首先，我们需要安装 Ollama Python SDK 并确定在添加追踪支持时需要关注哪些方法。您可以使用 pip install ollama-python 安装 Ollama Python SDK。
如果您使用过 OpenAI Python SDK，Ollama Python SDK 会让您感到非常熟悉。以下是我们如何使用它来执行聊天补全调用
from ollama import chat
from rich import print
response = chat(model="llama3.2",
messages = [
{"role": "user", "content": "Briefly describe the components of an MLflow model"}
]
)
print(response)
这将返回
ChatResponse(
model='llama3.2',
created_at='2025-01-30T15:57:39.097119Z',
done=True,
done_reason='stop',
total_duration=7687553708,
load_duration=823704250,
prompt_eval_count=35,
prompt_eval_duration=3414000000,
eval_count=215,
eval_duration=3447000000,
message=Message(
role='assistant',
content="In MLflow, a model consists of several key components:\n\n1. **Model Registry**: A centralized
storage for models, containing metadata such as the model's name, version, and description.\n2. **Model Version**:
A specific iteration of a model, represented by a unique version number. This can be thought of as a snapshot of
the model at a particular point in time.\n3. **Model Artifacts**: The actual model code, parameters, and data used
to train the model. These artifacts are stored in the Model Registry and can be easily deployed or reused.\n4.
**Experiment**: A collection of runs that use the same hyperparameters and model version to train and evaluate a
model. Experiments help track progress, provide reproducibility, and facilitate collaboration.\n5. **Run**: An
individual instance of training or testing a model using a specific experiment. Runs capture the output of each
run, including metrics such as accuracy, loss, and more.\n\nThese components work together to enable efficient
model management, version control, and reproducibility in machine learning workflows.",
images=None,
tool_calls=None
)
)
我们已经验证了 Ollama Python SDK 已设置并正在工作。我们也知道在添加追踪支持时需要关注的方法：ollama.chat。
步骤 2：编写追踪装饰器
我们可以通过几种方式向 Ollama 的 SDK 添加追踪——我们可以直接修改 SDK 代码、创建包装类，或者使用 Python 的方法修补功能。在此示例中，我们将使用装饰器来修补 SDK 的 chat 方法。这种方法允许我们在不修改 SDK 代码或创建额外包装类的情况下添加追踪，但它确实需要了解 Python 的装饰器模式和 MLflow 追踪的工作原理。
import mlflow
from mlflow.entities import SpanType
from mlflow.tracing.utils import set_span_chat_messages
from functools import wraps
from ollama import chat as ollama_chat
def _get_span_type(task_name: str) -> str:
span_type_mapping = {
"chat": SpanType.CHAT_MODEL,
}
return span_type_mapping.get(task_name, SpanType.UNKNOWN)
def trace_ollama_chat(func):
@wraps(func)
def wrapper(*args, **kwargs):
with mlflow.start_span(
name="ollama.chat",
span_type=_get_span_type("chat"),
) as span:
# Set model name as a span attribute
model_name = kwargs.get("model", "")
span.set_attribute("model_name", model_name)
# Log the inputs
input_messages = kwargs.get("messages", [])
span.set_inputs({
"messages": input_messages,
"model": model_name,
})
# Set input messages
set_span_chat_messages(span, input_messages)
# Make the API call
response = func(*args, **kwargs)
# Log the outputs
if hasattr(response, 'to_dict'):
output = response.to_dict()
else:
output = response
span.set_outputs(output)
output_message = response.message
# Append the output message
set_span_chat_messages(span, [{"role": output_message.role, "content": output_message.content}], append=True)
return response
return wrapper
让我们分解代码并看看它是如何工作的。
-
我们首先定义一个帮助函数
_get_span_type，它将 Ollama 方法映射到 MLflow 跨度类型。这并非严格必需，因为我们目前只追踪chat函数，但它展示了一种可以应用于其他方法的模式。这遵循 Anthropic 提供商的参考实现，如追踪贡献指南中所建议的那样。 -
我们使用
functools.wraps定义一个装饰器trace_ollama_chat，用于修补chat函数。这里有几个关键步骤-
我们使用
mlflow.start_span启动一个新跨度。跨度名称设置为 "ollama.chat"，跨度类型设置为_get_span_type返回的值。 -
我们使用
span.set_attribute将model_name设置为跨度的属性。这并非严格必需，因为模型名称将捕获在输入中，但它说明了如何向跨度设置任意属性。 -
我们使用
span.set_inputs将消息作为输入记录到跨度。我们通过访问kwargs字典从messages参数中获取这些信息。这些消息将记录到 MLflow UI 中跨度的“inputs”部分。我们还将模型名称记录为输入，再次说明如何记录任意输入。
-
我们使用 MLflow 的
set_span_chat_messages实用函数来格式化输入消息，使其能在 MLflow UI 的“Chat”面板中以美观的方式显示。此帮助程序确保消息得到正确格式化，并以适合每个消息角色的样式显示。 -
我们使用
func(*args, **kwargs)调用原始函数。这是 Ollama 的chat函数。 -
我们将函数输出记录为跨度属性，使用
span.set_outputs。这会获取来自 Ollama API 的响应并将其设置为跨度的属性。这些输出将记录到 MLflow UI 中跨度的“outputs”部分。
-
我们从响应中提取输出消息，并再次使用
set_span_chat_messages将其追加到聊天历史记录中，确保它出现在 MLflow UI 的“Chat”面板中。
-
最后，我们返回 API 调用的响应，不做任何更改。现在，当我们使用
trace_ollama_chat修补 chat 函数时，该函数将被追踪，但其行为将与以前一样。
-
一些需要注意的点
- 此实现使用简单的装饰器模式，在不修改底层 Ollama SDK 代码的情况下添加了追踪。这使其成为一种轻量级且可维护的方法。
- 使用
set_span_chat_messages可确保输入和输出消息在 MLflow UI 的“Chat”面板中以用户友好的方式显示，从而易于遵循对话流程。 - 我们还可以通过其他几种方式实现此追踪行为。我们可以编写一个包装类，或者使用一个简单的包装函数，用
@mlflow.trace装饰chat函数。一些编排框架可能需要更复杂的方法，例如回调或 API 钩子。有关更多详细信息，请参阅MLflow 追踪贡献指南。
步骤 3：修补 chat 方法并尝试使用它
现在我们有了一个追踪装饰器，我们可以修补 Ollama 的 chat 方法并尝试使用它。
original_chat = ollama_chat
chat = trace_ollama_chat(ollama_chat)
此代码有效地修补了当前范围内的 ollama.chat 函数。我们首先将原始函数存储在 original_chat 中以备后用，然后将 chat 重新分配给装饰后的版本。这意味着我们代码中任何后续对 chat() 的调用都将使用被追踪的版本，同时仍保留原始功能。
现在，当我们调用 chat() 时，该方法将被追踪，结果将记录到 MLflow UI 中
mlflow.set_experiment("ollama-tracing")
response = chat(model="llama3.2",
messages = [
{"role": "user", "content": "Briefly describe the components of an MLflow model"}
]
)
追踪工具和工具调用
Ollama Python SDK 支持工具调用。我们希望记录两件主要事情
- 可供 LLM 使用的工具
- 实际的工具调用，包括特定的工具及其传递给它的参数。
请注意，“工具调用”指的是 LLM 指定应使用哪个工具以及应将哪些参数传递给它的信息——而不是该工具的实际执行。当 LLM 进行工具调用时，它本质上是在说“应使用这些参数运行此工具”，而不是运行工具本身。工具的实际执行是分开的，通常在应用程序代码中进行。
这是更新后的追踪代码版本，修补了 Ollama chat 方法，记录了可用工具并捕获了工具调用
from mlflow.entities import SpanType
from mlflow.tracing.utils import set_span_chat_messages, set_span_chat_tools
from functools import wraps
from ollama import chat as ollama_chat
import json
from uuid import uuid4
def _get_span_type(task_name: str) -> str:
span_type_mapping = {
"chat": SpanType.CHAT_MODEL,
}
return span_type_mapping.get(task_name, SpanType.UNKNOWN)
def trace_ollama_chat(func):
@wraps(func)
def wrapper(*args, **kwargs):
with mlflow.start_span(
name="ollama.chat",
span_type=_get_span_type("chat"),
) as span:
# Set model name as a span attribute
model_name = kwargs.get("model", "")
span.set_attribute("model_name", model_name)
# Log the inputs
input_messages = kwargs.get("messages", [])
tools = kwargs.get("tools", [])
span.set_inputs({
"messages": input_messages,
"model": model_name,
"tools": tools,
})
# Set input messages and tools
set_span_chat_messages(span, input_messages)
if tools:
set_span_chat_tools(span, tools)
# Make the API call
response = func(*args, **kwargs)
# Log the outputs
if hasattr(response, "to_dict"):
output = response.to_dict()
else:
output = response
span.set_outputs(output)
output_message = response.message
# Prepare the output message for span
output_span_message = {
"role": output_message.role,
"content": output_message.content,
}
# Handle tool calls if present
if output_message.tool_calls:
tool_calls = []
for tool_call in output_message.tool_calls:
tool_calls.append({
"id": str(uuid4()),
"type": "function",
"function": {
"name": tool_call.function.name,
"arguments": json.dumps(tool_call.function.arguments),
}
})
output_span_message["tool_calls"] = tool_calls
# Append the output message
set_span_chat_messages(span, [output_span_message], append=True)
return response
return wrapper
这里的关键变化是
- 我们使用
tools = kwargs.get("tools", [])从tools参数中提取了可用工具列表，将它们记录为输入，并使用set_span_chat_tools将它们捕获以包含在“Chat”面板中。 - 我们为工具调用在输出消息中添加了特定的处理，确保根据ToolCall规范对其进行格式化。
现在让我们用一个简单的提示计算工具来测试一下。工具是根据OpenAI 规范来定义的，用于工具调用。
chat = trace_ollama_chat(ollama_chat)
tools = [
{
"type": "function",
"function": {
"name": "calculate_tip",
"description": "Calculate the tip amount based on the bill amount and tip percentage",
"parameters": {
"type": "object",
"properties": {
"bill_amount": {
"type": "number",
"description": "The total bill amount"
},
"tip_percentage": {
"type": "number",
"description": "The percentage of the bill to be given as a tip, given as a whole number."
}
},
"required": ["bill_amount", "tip_percentage"]
}
}
}
]
response = chat(
model="llama3.2",
messages=[
{"role": "user", "content": "What is the tip for a $187.32 bill with a 22% tip?"}
],
tools=tools,
)
我们可以检查 MLflow UI 中的追踪，现在“Chat”面板中会显示可用工具和工具调用结果
编排：构建工具调用循环
到目前为止，Ollama 示例仅在进行聊天补全时生成单个跨度。但许多 GenAI 应用程序包括多个 LLM 调用、检索步骤、工具执行和其他自定义步骤。虽然我们不会在此处详细介绍如何向编排框架添加追踪，但我们将通过定义一个基于前面定义的工具的工具调用循环来说明一些关键概念。
工具调用循环将遵循以下模式
- 接收用户提示作为输入
- 以工具调用形式响应
- 对于每个工具调用，执行该工具并存储结果
- 使用
tool角色将工具调用结果追加到消息历史记录中 - 再次调用 LLM 并传入工具调用结果，提示它对用户的提示提供最终答案
这是一个只进行了一次工具调用的实现。
class ToolExecutor:
def __init__(self):
self.tools = [
{
"type": "function",
"function": {
"name": "calculate_tip",
"description": "Calculate the tip amount based on the bill amount and tip percentage",
"parameters": {
"type": "object",
"properties": {
"bill_amount": {
"type": "number",
"description": "The total bill amount"
},
"tip_percentage": {
"type": "number",
"description": "The percentage of the bill to be given as a tip, represented as a whole number."
}
},
"required": ["bill_amount", "tip_percentage"]
}
}
}
]
# Map tool names to their Python implementations
self.tool_implementations = {
"calculate_tip": self._calculate_tip
}
def _calculate_tip(self, bill_amount: float, tip_percentage: float) -> float:
"""Calculate the tip amount based on the bill amount and tip percentage."""
bill_amount = float(bill_amount)
tip_percentage = float(tip_percentage)
return round(bill_amount * (tip_percentage / 100), 2)
def execute_tool_calling_loop(self, messages):
"""Execute a complete tool calling loop with tracing."""
with mlflow.start_span(
name="ToolCallingLoop",
span_type="CHAIN",
) as parent_span:
# Set initial inputs
parent_span.set_inputs({
"initial_messages": messages,
"available_tools": self.tools
})
# Set input messages
set_span_chat_messages(parent_span, messages)
# First LLM call (already traced by our chat method patch)
response = chat(
messages=messages,
model="llama3.2",
tools=self.tools,
)
messages.append(response.message)
tool_calls = response.message.tool_calls
tool_results = []
# Execute tool calls
for tool_call in tool_calls:
with mlflow.start_span(
name=f"ToolExecution_{tool_call.function.name}",
span_type="TOOL",
) as tool_span:
# Parse tool inputs
tool_inputs = tool_call.function.arguments
tool_span.set_inputs(tool_inputs)
# Execute tool
func = self.tool_implementations.get(tool_call.function.name)
if func is None:
raise ValueError(f"No implementation for tool: {tool_call.function.name}")
result = func(**tool_inputs)
tool_span.set_outputs({"result": result})
tool_results.append({
"tool_call_id": str(uuid4()),
"output": str(result)
})
messages.append({
"role": "tool",
"tool_call_id": str(uuid4()),
"content": str(result)
})
# Prepare messages for final response
messages.append({
"role": "user",
"content": "Answer the initial question based on the tool call results. Do not refer to the tool call results in your response. Just give a direct answer."
})
# Final LLM call (already traced by our chat method patch)
final_response = chat(
messages=messages,
model="llama3.2"
)
# Set the final output for the parent span
parent_span.set_outputs({
"final_response": final_response.message.content,
"tool_results": tool_results
})
print(final_response)
# set output messages
set_span_chat_messages(parent_span, [final_response.message.model_dump()], append=True)
return final_response
这是我们在该工具调用循环中处理追踪的方式
- 我们首先使用
mlflow.start_span为工具调用循环设置一个父跨度。我们将跨度名称设置为 "ToolCallingLoop"，跨度类型设置为 "CHAIN"，表示操作链。 - 我们将初始消息和可用工具记录为跨度的输入。这有助于未来的调试，因为它允许我们验证工具是否已正确提供和配置。
- 我们使用我们修补过的
chat函数进行第一次 LLM 调用。此调用已由我们的装饰器追踪，因此我们无需进行任何特殊操作即可追踪它。 - 我们遍历工具调用，执行每个工具并将结果存储起来。每次工具执行都会使用一个新跨度进行追踪，该跨度以工具函数名称命名。输入和输出作为属性记录在跨度上。
- 我们将工具调用结果追加到消息历史记录中，使用
tool角色。这允许 LLM 在后续请求中看到工具调用的结果。它还允许我们在 MLflow UI 中看到工具调用结果。 - 我们准备最终响应的消息，包括一个提示，要求根据工具调用结果回答最初的问题。
- 我们使用我们修补过的
chat函数进行最终的 LLM 调用。同样，因为我们使用的是修补过的函数，所以此调用已被追踪。 - 我们将最终输出设置为父跨度的输出，包括来自 LLM 的最终响应和工具结果。
- 最后，我们使用
set_span_chat_messages将最终响应追加到 MLflow UI 中的聊天历史记录中。请注意，为保持简洁明了，我们只使用set_span_chat_messages将用户的初始查询和最终响应记录到父跨度中。我们可以点击嵌套的跨度以查看工具调用结果和其他详细信息。
此过程创建了从初始请求到工具执行和最终响应的整个工具调用循环的综合追踪。
我们可以如下执行此操作。但是，请注意，在完全了解 LLM 生成或调用的任意代码将对您的系统执行什么操作之前，您不应运行它。
executor = ToolExecutor()
response = executor.execute_tool_calling_loop(
messages=[
{"role": "user", "content": "What is the tip for a $235.32 bill with a 22% tip?"}
]
)
结果如下所示的追踪
结论
本文展示了如何将 MLflow 追踪扩展到其内置提供商支持之外。我们从一个简单的示例开始——向 Ollama Python SDK 的 chat 方法添加追踪——并看到了如何通过轻量级修补，我们可以捕获关于每次聊天补全的详细信息。然后，我们在此基础上构建，以追踪更复杂的工具执行循环。
关键要点是
- MLflow 追踪高度可定制，可以适应尚无自动日志记录支持的提供商
- 添加基本追踪支持通常只需对代码进行最少的更改。在此案例中，我们修补了 Ollama Python SDK 的
chat方法，并编写了几行代码以添加追踪支持。 - 用于简单 API 调用的相同原则可以扩展到具有多个步骤的复杂工作流。在此案例中，我们追踪了一个包含多个步骤和工具调用的工具调用循环。