超越自动日志：为新的LLM提供商添加MLflow追踪功能
在这篇文章中，我们将展示如何通过为Ollama Python SDK的chat方法添加追踪支持，来为新的LLM提供商添加MLflow追踪功能。
MLflow Tracing 是 MLflow 中的一种可观测性工具，用于捕获 GenAI 应用程序和工作流程的详细执行轨迹。除了单个调用的输入、输出和元数据外，MLflow tracing 还可以捕获中间步骤，例如工具调用、推理步骤、检索步骤或其他自定义步骤。
MLflow 为许多流行的LLM提供商和编排框架提供内置追踪支持。如果您正在使用这些提供商之一，只需一行代码即可启用追踪：mlflow.。虽然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追踪的基本概念及其工作原理已有初步了解。如果您正在学习或需要复习，请查阅Tracing Concepts指南。
为新供应商添加追踪涉及几个关键考虑因素：
-
理解提供商的核心功能： 我们首先需要了解需要追踪哪些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)
Which will return:
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 provider 参考实现。 -
我们定义了一个装饰器，
trace_ollama_chat，使用functools.wraps，它修补了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的聊天面板中良好显示。这个辅助函数确保消息被正确格式化，并根据每个消息角色以适当的样式显示。 -
我们使用
func(*args, **kwargs)调用原始函数。这是Ollama的chat函数。 -
我们使用
span.set_outputs将函数的输出记录为跨度属性。这会接收来自Ollama API的响应并将其设置为跨度的属性。这些输出将被记录到MLflow UI中跨度的"outputs"部分。
-
我们从响应中提取输出消息，并再次使用
set_span_chat_messages将其附加到聊天历史记录中，确保它显示在MLflow UI的聊天面板中。
-
最后，我们直接返回API调用的响应，不做任何更改。现在，当我们用
trace_ollama_chat修补聊天函数时，该函数将被追踪，但其他行为保持正常。
-
需要注意的几点：
- 该实现采用简单的装饰器模式，在不修改底层Ollama SDK代码的情况下添加追踪功能。这使得该方法轻量且易于维护。
- 使用
set_span_chat_messages确保输入和输出消息以用户友好的方式显示在MLflow UI的聊天面板中，便于跟踪对话流程。 - 我们还可以通过其他几种方式实现这种追踪行为。我们可以编写一个包装类，或者使用一个简单的包装函数，用
@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聊天方法，记录可用工具并捕获工具调用：
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参数中提取了可用工具列表，使用tools = kwargs.get("tools", [])，将它们记录为输入，并使用set_span_chat_tools来捕获它们以便包含在聊天面板中。 - 我们在输出消息中为工具调用添加了特定处理，确保按照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中检查追踪记录，现在同时显示可用工具和工具调用结果：
编排：构建工具调用循环
到目前为止，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记录用户的初始查询和最终响应到父跨度。我们可以点击嵌套跨度来查看工具调用结果和其他详细信息。
这个过程创建了整个工具调用循环的全面追踪，从初始请求到工具执行再到最终响应。
我们可以按如下方式执行此操作。但是，请注意，在完全了解它将对你系统执行什么操作之前，你不应运行由LLMs生成或调用的任意代码。
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 Tracing 高度可定制，可适配不支持自动日志记录的提供程序
- 添加基础追踪支持通常只需少量代码改动。在此案例中，我们修补了Ollama Python SDK的
chat方法，并编写了几行代码来添加追踪支持。 - 用于简单API调用的相同原则可以扩展到具有多个步骤的复杂工作流。在这种情况下，我们追踪了一个包含多个步骤和工具调用的工具调用循环。