メインコンテンツまでスキップ

干し草の山をたどる

Haystackは、本番運用対応のLLMアプリケーション、セマンティック検索システム、質問応答システムを構築するためのオープンソースAIオーケストレーション フレームワークです。

MLflow Tracing は、 Haystack の自動トレース機能を提供します。mlflow.haystack.autolog関数を呼び出すことで Haystack のトレースを有効にできます。パイプラインとコンポーネントの呼び出し時にトレースがアクティブなMLflowエクスペリメントに自動的に記録されます。

Python
import mlflow

mlflow.haystack.autolog()

MLflow トレースは、Haystack パイプラインの実行に関する次の情報を自動的にキャプチャします。

  • パイプラインとコンポーネント
  • レイテンシー
  • コンポーネントメタデータ
  • トークンの使用とコスト
  • キャッシュヒット情報
  • 発生した例外
注記

サーバレス コンピュート クラスターでは、自動ログは自動的に有効になりません。 この統合の自動トレースを有効にするには、 mlflow.haystack.autolog()明示的に呼び出す必要があります。

前提条件​

Haystack で MLflow Tracing を使用するには、MLflow と Haystack パッケージをインストールする必要があります。

開発環境では、Databricks エクストラと Haystack パッケージを含む完全な MLflow パッケージをインストールします。

Bash
pip install --upgrade "mlflow[databricks]>=3.1" haystack-ai

完全な mlflow[databricks] パッケージには、Databricks でのローカル開発と実験のためのすべての機能が含まれています。

注記

Haystack で最高のトレース体験を得るには、MLflow 3 を強くお勧めします。

例を実行する前に、環境を設定する必要があります。

Databricks ノートブックの外部のユーザーの場合 : Databricks 環境変数を設定します。

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

Databricks ノートブック内のユーザーの場合 : これらの資格情報は自動的に設定されます。

API キー : LLM プロバイダーの API キーが設定されていることを確認します。本番運用環境の場合は、安全なAPIキー管理のために、ハードコードされた値の代わりにMosaic AI GatewayまたはDatabricksシークレットを使用します。

Bash
export OPENAI_API_KEY="your-openai-api-key"
# Add other provider keys as needed

使用例​

次の例は、Haystack を MLflow トレースとともに使用する方法を示しています。この例では、リトリーバー、プロンプト ビルダー、LLM を使用して、単純な RAG (Retrieval-Augmented Generation) パイプラインを作成します。

Python
import mlflow
import os

from haystack import Pipeline
from haystack.components.builders import PromptBuilder
from haystack.components.generators import OpenAIGenerator
from haystack.components.retrievers.in_memory import InMemoryBM25Retriever
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack import Document

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

# Enable auto tracing for Haystack
mlflow.haystack.autolog()

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

# Create a simple document store with sample documents
document_store = InMemoryDocumentStore()
document_store.write_documents([
    Document(content="Paris is the capital of France."),
    Document(content="Berlin is the capital of Germany."),
    Document(content="Rome is the capital of Italy."),
])

# Build a simple RAG pipeline
template = """
Given the following documents, answer the question.

Documents:
{% for doc in documents %}
  
  {{ doc.content }}
{% endfor %}

Question: {{ question }}
Answer:
"""

pipe = Pipeline()
pipe.add_component("retriever", InMemoryBM25Retriever(document_store=document_store))
pipe.add_component("prompt_builder", PromptBuilder(template=template))
pipe.add_component("llm", OpenAIGenerator(model="gpt-4o-mini"))

pipe.connect("retriever", "prompt_builder.documents")
pipe.connect("prompt_builder", "llm")

# Run the pipeline - trace will be automatically logged
result = pipe.run({
    "retriever": {"query": "What is the capital of France?"},
    "prompt_builder": {"question": "What is the capital of France?"}
})

print(result["llm"]["replies"][0])
警告

本番運用環境の場合は、安全なAPIキー管理のために、ハードコードされた値の代わりにMosaic AI GatewayまたはDatabricksシークレットを使用します。

トークン使用状況の追跡​

MLflow バージョン 3.4.0 を使用すると、MLflow は Haystack パイプラインのトークンの使用状況を自動的に追跡します。またはそれ以降。VPN 使用情報には、パイプラインの実行中に消費された入力ウイルス、出力ウイルス、および合計消費ウイルスが含まれます。

Python
import mlflow

mlflow.haystack.autolog()

from haystack import Pipeline
from haystack.components.generators import OpenAIGenerator

# Create and run a pipeline
pipe = Pipeline()
pipe.add_component("llm", OpenAIGenerator(model="gpt-4o-mini"))

# Run the pipeline and retrieve trace information
with mlflow.start_span(name="haystack_pipeline_run") as span:
    result = pipe.run({"llm": {"prompt": "What is the capital of France?"}})
    print(result["llm"]["replies"][0])

# Token usage is automatically logged and visible in the MLflow UI
trace_info = mlflow.get_last_active_trace()
print(f"Trace ID: {trace_info.request_id}")

トークンの使用状況の詳細は MLflow トレース UI に表示されるため、パイプラインのパフォーマンスとコストを監視して最適化できます。

自動トレースを無効にする​

Haystack の自動トレースは、 mlflow.haystack.autolog(disable=True)またはmlflow.autolog(disable=True)を呼び出すことによってグローバルに無効にすることができます。

に最終更新