Aller au contenu principal

Traçage de Haystack

Haystack est un framework d'orchestration d'IA open source pour la création d'applications LLM prêtes pour la production, de systèmes de recherche sémantique et de systèmes de questions-réponses.

MLflow Tracing fournit une capacité de traçage automatique pour Haystack. Vous pouvez activer le traçage pour Haystack en appelant la fonction mlflow.haystack.autolog, et les traces sont automatiquement consignées dans l'Experiment MLflow actif lors de l'invocation de pipelines et de composants.

Python
import mlflow

mlflow.haystack.autolog()

La trace MLflow capture automatiquement les informations suivantes concernant les exécutions de pipeline Haystack :

  • Pipelines et composants
  • Latences
  • Métadonnées des composants.
  • Utilisation et coût des jetons
  • Informations sur les accès au cache
  • Toute exception si levée
remarque

Sur les clusters de compute serverless, l'autologging n'est pas activé automatiquement. Vous devez appeler explicitement mlflow.haystack.autolog() pour activer le traçage automatique de cette intégration.

Prérequis

Pour utiliser MLflow Tracing avec Haystack, vous devez installer MLflow et les packages Haystack.

Pour les environnements de développement, installez le package MLflow complet avec les modules complémentaires Databricks et les packages Haystack :

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

Le package mlflow[databricks] complet inclut toutes les fonctionnalités pour le développement local et l’expérimentation sur Databricks.

remarque

MLflow 3 est fortement recommandé pour une meilleure expérience de traçage avec Haystack.

Avant d'exécuter les exemples, vous devrez configurer votre environnement :

Pour les utilisateurs en dehors des notebooks Databricks : Définissez vos variables d'environnement Databricks :

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

Pour les utilisateurs dans les Notebooks Databricks : ces identifiants sont définis automatiquement pour vous.

Clés API : Assurez-vous que les clés API de votre fournisseur LLM sont configurées. Pour les environnements de production, utilisez AI Gateway ou les secrets Databricks au lieu de valeurs codées en dur pour une gestion sécurisée des clés API.

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

Exemple d'utilisation

L’exemple suivant montre comment utiliser Haystack avec le suivi MLflow. Cet exemple crée un pipeline RAG (Génération augmentée de récupération) simple avec un récupérateur, un générateur de prompts et un LLM.

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])
attention

Pour les environnements de production, utilisez AI Gateway ou les secrets Databricks au lieu de valeurs codées en dur pour la gestion sécurisée des clés API.

Suivi de l'utilisation des jetons

MLflow suit automatiquement l'utilisation des jetons pour les pipelines Haystack lors de l'utilisation de MLflow version 3.4.0. ou ultérieure. Les informations d'utilisation des jetons incluent les jetons d'entrée, les jetons de sortie et les jetons totaux consommés pendant l'exécution du pipeline.

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}")

Les détails d'utilisation des jetons sont affichés dans l'interface utilisateur de suivi MLflow, vous permettant de surveiller et d'optimiser les performances et les coûts de votre pipeline.

Désactiver le suivi automatique

Le traçage automatique pour Haystack peut être désactivé globalement en appelant mlflow.haystack.autolog(disable=True) ou mlflow.autolog(disable=True).