跳到主要内容

追踪 LlamaIndex🦙

LlamaIndex Tracing via autolog

LlamaIndex 是一个用于构建代理式生成 AI 应用程序的开源框架，它允许大语言模型处理任何格式的数据。

MLflow Tracing 为 LlamaIndex 提供了自动追踪功能。您可以通过调用 mlflow.llama_index.autolog() 函数来为 LlamaIndex 启用追踪，在调用 LlamaIndex 引擎和工作流时，嵌套的追踪信息会自动记录到活动的 MLflow 实验中。

import mlflow

mlflow.llama_index.autolog()
提示

MLflow 与 LlamaIndex 的集成不仅仅是追踪。MLflow 为 LlamaIndex 提供了全面的追踪体验，包括模型追踪、索引管理和评估。请查看 MLflow LlamaIndex Flavor 了解更多信息！

示例用法​

首先，让我们下载一个测试数据来创建一个玩具索引。

!mkdir -p data
!curl -L https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt -o ./data/paul_graham_essay.txt

将它们加载到一个简单的内存向量索引中。

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader

documents = SimpleDirectoryReader("data").load_data()
index = VectorStoreIndex.from_documents(documents)

现在您可以启用 LlamaIndex 自动追踪并开始查询索引。

import mlflow

# Enabling tracing for LlamaIndex
mlflow.llama_index.autolog()

# Optional: Set a tracking URI and an experiment
mlflow.set_tracking_uri("https://:5000")
mlflow.set_experiment("LlamaIndex")

# Query the index
query_engine = index.as_query_engine()
response = query_engine.query("What was the first program the author wrote?")

Token 用量​

MLflow >= 3.2.0 版本支持 LlamaIndex 的 Token 用量追踪。每次 LLM 调用的 Token 用量将被记录在 mlflow.chat.tokenUsage 属性中。整个追踪过程中的总 Token 用量将在追踪信息对象的 token_usage 字段中提供。

import json
import mlflow
from llama_index.llms.openai import OpenAI

mlflow.llama_index.autolog()

# Use the chat complete method to create new chat.
llm = OpenAI(model="gpt-3.5-turbo")
Settings.llm = llm
response = llm.chat(
    [ChatMessage(role="user", content="What is the capital of France?")]
)

# Get the trace object just created
last_trace_id = mlflow.get_last_active_trace_id()
trace = mlflow.get_trace(trace_id=last_trace_id)

# Print the token usage
total_usage = trace.info.token_usage
print("== Total token usage: ==")
print(f"  Input tokens: {total_usage['input_tokens']}")
print(f"  Output tokens: {total_usage['output_tokens']}")
print(f"  Total tokens: {total_usage['total_tokens']}")

# Print the token usage for each LLM call
print("\n== Detailed usage for each LLM call: ==")
for span in trace.data.spans:
    if usage := span.get_attribute("mlflow.chat.tokenUsage"):
        print(f"{span.name}:")
        print(f"  Input tokens: {usage['input_tokens']}")
        print(f"  Output tokens: {usage['output_tokens']}")
        print(f"  Total tokens: {usage['total_tokens']}")
== Total token usage: ==
  Input tokens: 14
  Output tokens: 7
  Total tokens: 21

== Detailed usage for each LLM call: ==
OpenAI.chat:
  Input tokens: 14
  Output tokens: 7
  Total tokens: 21

LlamaIndex 工作流​

Workflow 是 LlamaIndex 的下一代 GenAI 编排框架。它被设计为一个灵活且可解释的框架，用于构建任意的 LLM 应用程序，例如代理、RAG 流程、数据提取管道等。MLflow 支持追踪、评估和记录 Workflow 对象，这使得它们更具可观察性和可维护性。

通过调用相同的 mlflow.llama_index.autolog()，LlamaIndex 工作流的自动追踪功能即可开箱即用。

要了解更多关于 MLflow 与 LlamaIndex Workflow 集成的信息，请继续阅读以下教程。

禁用自动跟踪​

可以通过调用 mlflow.llama_index.autolog(disable=True) 或 mlflow.autolog(disable=True) 来全局禁用 LlamaIndex 的自动追踪。