在本文中，我们将展示如何通过向 Ollama Python SDK 的 chat 方法添加 Tracing 支持，从而为新的 LLM 提供商添加 MLflow Tracing。
MLflow Tracing 是 MLflow 中的一个可观测性工具，用于捕获 GenAI 应用和工作流的详细执行跟踪。除了捕获单个调用的输入、输出和元数据外，MLflow Tracing 还可以捕获中间步骤，例如工具调用、推理步骤、检索步骤或其他自定义步骤。
MLflow 为许多流行的 LLM 提供商和编排框架提供了 内置的 Tracing 支持。如果您正在使用这些提供商之一，只需一行代码即可启用 Tracing：mlflow.<provider>.autolog()。尽管 MLflow 的 autologging 功能覆盖了许多最广泛使用的 LLM 提供商和编排框架，但有时您可能需要为不受支持的提供商添加 Tracing，或定制超出 autologging 提供的 Tracing 功能。本文通过以下方式展示了 MLflow Tracing 的灵活性和可扩展性：
- 为不受支持的提供商（Ollama Python SDK）添加基本的 Tracing 支持
- 展示如何捕获简单的补全以及更复杂的工具调用工作流
- 说明如何在现有代码中以最小的改动添加 Tracing
我们将使用 Ollama Python SDK（一个用于 Ollama LLM 平台的开源 Python SDK）作为示例。我们将逐步讲解整个过程，展示如何使用 MLflow Tracing 捕获关键信息，同时保持与提供商 SDK 的干净集成。请注意，MLflow 确实对 Ollama 有 autologging 支持，但目前仅限于通过 OpenAI 客户端使用，而不是直接使用 Ollama Python SDK。
为新的提供商添加 MLflow Tracing：一般原则
MLflow 文档中有一个 出色的指南，介绍了如何贡献 MLflow Tracing。虽然在本例中我们不会向 MLflow 本身贡献代码，但我们将遵循相同的通用原则。
本文假设您对 MLflow Tracing 是什么以及它如何工作有基本了解。如果您刚开始学习，或者需要回顾，请参阅 Tracing 概念指南。
为新的提供商添加 Tracing 涉及几个关键考虑因素
-
理解提供商的核心功能：首先，我们需要了解为了获取所需的 Tracing 信息，需要跟踪哪些 API 方法。对于 LLM 推理提供商，这通常包括聊天补全、工具调用或嵌入生成等操作。在编排框架中，这可能涉及检索、推理、路由或各种自定义步骤。在我们的 Ollama 示例中，我们将重点关注聊天补全 API。此步骤会因提供商而异。
-
将操作映射到 Span： MLflow Tracing 使用不同的 *span 类型*来表示不同类型的操作。您可以在此处找到内置 span 类型的描述。不同的 span 类型在 MLflow UI 中显示方式不同，并且可以启用特定功能。在 span 内，我们还需要将提供商的输入和输出映射到 MLflow 期望的格式。MLflow 提供了用于记录聊天和工具输入输出的实用工具，这些信息会作为格式化消息显示在 MLflow UI 中。
为新的提供商添加 Tracing 时，我们的主要任务是将提供商的 API 方法映射到具有适当 span 类型的 MLflow Tracing span。
-
构建和保留关键数据：对于我们想要跟踪的每个操作，我们需要确定要保留的关键信息，并确保它被捕获并以有用的方式显示。例如，我们可能希望捕获控制操作行为的输入和配置数据、解释结果的输出和元数据、过早终止操作的错误等。查看类似提供商的 Tracing 和 Tracing 实现可以为如何构建和保留这些数据提供良好的起点。
向 Ollama Python SDK 添加 Tracing
现在我们对为新提供商添加 Tracing 的关键步骤有了高层理解，接下来让我们逐步完成过程，并向 Ollama Python SDK 添加 Tracing。
步骤 1：安装和测试 Ollama Python SDK
首先，我们需要安装 Ollama Python SDK 并确定在添加 Tracing 支持时需要关注哪些方法。您可以使用 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 已设置并正常工作。我们还知道在添加 Tracing 支持时需要重点关注的方法是：ollama.chat。
步骤 2：编写 Tracing Decorator
有几种方法可以向 Ollama 的 SDK 添加 Tracing——我们可以直接修改 SDK 代码，创建一个包装类，或使用 Python 的方法补丁功能。在本例中，我们将使用 decorator 来修补 SDK 的 chat 方法。这种方法使我们无需修改 SDK 代码或创建额外的包装类即可添加 Tracing，尽管它需要理解 Python 的 decorator 模式和 MLflow Tracing 的工作原理。
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 span 类型。这并非严格必要，因为我们目前只跟踪chat函数，但它展示了一种可以应用于其他方法的模式。这遵循了 Anthropic 提供商的参考实现，如 tracing 贡献指南中所推荐。 -
我们使用
functools.wraps定义一个 decoratortrace_ollama_chat，它修补了chat函数。这里有几个关键步骤：-
我们使用
mlflow.start_span开始一个新的 span。span 名称设置为 "ollama.chat"，span 类型设置为由_get_span_type返回的值。 -
我们使用
span.set_attribute将model_name设置为 span 的一个属性。这并非严格必要，因为模型名称将在输入中捕获，但它演示了如何在 span 上设置任意属性。 -
我们使用
span.set_inputs将消息作为输入记录到 span。我们通过访问kwargs字典从messages参数中获取这些消息。这些消息将被记录到 MLflow UI 中 span 的“inputs”部分。我们还将模型名称记录为输入，再次演示如何记录任意输入。
-
我们使用 MLflow 的
set_span_chat_messages实用函数来格式化输入消息，使其在 MLflow UI 的 Chat 面板中良好显示。此助手确保消息被正确格式化，并为每个消息角色显示适当的样式。 -
我们使用
func(*args, **kwargs)调用原始函数。这是 Ollama 的chat函数。 -
我们使用
span.set_outputs将函数的输出记录为 span 属性。这会将来自 Ollama API 的响应作为 span 的一个属性。这些输出将被记录到 MLflow UI 中 span 的“outputs”部分。
-
我们从响应中提取输出消息，并再次使用
set_span_chat_messages将其附加到聊天历史记录中，确保它出现在 MLflow UI 的 Chat 面板中。
-
最后，我们返回 API 调用的响应，不进行任何更改。现在，当我们使用
trace_ollama_chat修补 `chat` 函数时，该函数将被跟踪，但其行为将与正常情况相同。
-
需要注意的几点：
- 此实现使用了简单的 decorator 模式，在不修改底层 Ollama SDK 代码的情况下添加了 tracing。这使其成为一种轻量级且易于维护的方法。
- 使用
set_span_chat_messages确保输入和输出消息在 MLflow UI 的 Chat 面板中以用户友好的方式显示，从而易于跟踪对话流程。 - 还有其他几种方法可以实现这种 tracing 行为。我们可以编写一个 wrapper class，或者使用一个简单的 wrapper function，用
@mlflow.trace装饰chat函数。一些 orchestration frameworks 可能需要更复杂的方法，例如 callbacks 或 API hooks。有关更多详细信息，请参阅 MLflow Tracing 贡献指南。
步骤 3：修补 chat 方法并试用
现在我们有了 tracing decorator，我们可以修补 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 进行工具调用时，它本质上是在说“这个工具应该使用这些参数运行”，而不是自己运行工具。工具的实际执行是单独发生的，通常在应用程序代码中。
这是更新后的 tracing 代码版本，修补了 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 中检查跟踪，现在同时显示了可用工具和工具调用结果
编排：构建工具调用循环
到目前为止，Ollama 示例在每次进行聊天补全时只生成一个 span。但许多 GenAI 应用包含多个 LLM 调用、检索步骤、工具执行和其他自定义步骤。虽然我们在此不详细介绍如何向编排框架添加 Tracing，但我们将通过定义一个基于我们之前定义的工具的工具调用循环来阐述一些关键概念。
工具调用循环将遵循以下模式：
- 将用户提示作为输入
- 以工具调用作为响应
- 对于每个工具调用，执行工具并存储结果
- 将工具调用结果以
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
我们在此工具调用循环中处理 Tracing 的方法如下：
- 我们首先使用
mlflow.start_span为工具调用循环设置一个父 span。我们将 span 名称设置为 "ToolCallingLoop"，span 类型设置为 "CHAIN"，代表一个操作链。 - 我们将初始消息和可用工具记录为 span 的输入。这对于未来的调试可能很有帮助，因为它可以让我们验证工具是否可用且配置正确。
- 我们使用修补后的
chat函数进行第一次 LLM 调用。此调用已被我们的 decorator 跟踪，因此我们无需进行任何特殊操作来跟踪它。 - 我们迭代工具调用，执行每个工具并存储结果。每个工具执行都通过一个新的 span 进行跟踪，该 span 以工具函数名称命名。输入和输出作为 span 的属性进行记录。
- 我们将工具调用结果以
tool角色附加到消息历史记录中。这使得 LLM 可以在后续请求中看到工具调用的结果。它也使我们能够在 MLflow UI 中看到工具调用结果。 - 我们准备最终响应的消息，包括一个基于工具调用结果回答初始问题的提示。
- 我们使用修补后的
chat函数进行最终的 LLM 调用。同样，因为我们使用了修补后的函数，所以此调用已被跟踪。 - 我们为父 span 设置最终输出，包括来自 LLM 的最终响应和工具结果。
- 最后，我们使用
set_span_chat_messages将最终响应附加到 MLflow UI 的聊天历史记录中。请注意，为了保持简洁，我们只使用set_span_chat_messages将用户的初始查询和最终响应记录到父 span 中。我们可以点击嵌套的 span 查看工具调用结果和其他详细信息。
此过程创建了一个包含整个工具调用循环的完整跟踪，从初始请求到工具执行和最终响应。
我们可以按如下方式执行。但是请注意，在完全理解 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 Tracing 扩展到其内置提供商支持之外。我们从一个简单的例子开始——向 Ollama Python SDK 的 chat 方法添加 Tracing——并看到了如何通过轻量级的修补捕获每个聊天补全的详细信息。然后，我们在此基础上构建了一个更复杂的工具执行循环跟踪。
关键要点如下：
- MLflow Tracing 具有高度可定制性，可以适应没有 autologging 的提供商
- 添加基本的 Tracing 支持通常只需很少的代码修改。在本例中，我们修补了 Ollama Python SDK 的
chat方法，并编写了几行代码来添加 Tracing 支持。 - 用于简单 API 调用的相同原则可以扩展到包含多个步骤的复杂工作流。在本例中，我们跟踪了一个包含多个步骤和工具调用的工具调用循环。