Pular para o conteúdo principal

Rastreando o Palheiro

Haystack é um framework de orquestração AI de código aberto para a construção de aplicações LLM prontas para produção, sistemas de busca semântica e sistemas de perguntas e respostas.

O MLflow Tracing oferece um recurso de rastreamento automático para o Haystack. Você pode habilitar o rastreamento para Haystack chamando a função mlflow.haystack.autolog , e os rastreamentos são registrados automaticamente no experimento MLflow ativo após a invocação do pipeline e dos componentes.

Python
import mlflow

mlflow.haystack.autolog()

O rastreamento MLflow captura automaticamente as seguintes informações sobre a execução pipeline Haystack:

  • tubulação e componentes
  • Latências
  • Metadados do componente
  • utilização e custo dos tokens
  • Informações sobre acertos de cache
  • Qualquer exceção, se lançada.
nota

Em clusters compute serverless , o registro automático de logs não é ativado automaticamente. Você deve chamar explicitamente mlflow.haystack.autolog() para habilitar o rastreamento automático para esta integração.

Pré-requisitos

Para usar MLflow Tracing com o Haystack, você precisa instalar MLflow e o pacote Haystack.

Para ambientes de desenvolvimento, instale o pacote completo MLflow com os extras Databricks e o pacote Haystack:

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

O pacote completo mlflow[databricks] inclui todos os recursos para desenvolvimento local e experimentação no Databricks.

nota

Recomenda-se fortemente o uso do MLflow 3 para obter a melhor experiência de rastreamento com o Haystack.

Antes de executar os exemplos, você precisará configurar seu ambiente:

Para usuários fora Databricks Notebook : Defina sua variável de ambiente Databricks :

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

Para usuários dentro Databricks Notebook : Essas credenciais são definidas automaticamente para você.

ChaveAPI : Certifique-se de que a chave API do seu provedor LLM esteja configurada. Para ambientes de produção, utilize segredosMosaic AI Gateway ou Databricks em vez de valores fixos para um gerenciamento seguro key API .

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

Exemplo de uso

O exemplo a seguir demonstra como usar o Haystack com o rastreamento do MLflow. Este exemplo cria um pipeline RAG (Retrieval-Augmented Generation) simples com um recuperador, um construtor de prompts e um 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])
atenção

Para ambientes de produção, utilize segredosMosaic AI Gateway ou Databricks em vez de valores fixos para um gerenciamento seguro key API .

Uso de tokens

MLflow rastreia automaticamente o uso de tokens para o pipeline Haystack ao usar a versão 3.4.0 MLflow ou mais tarde. As informações sobre o uso de tokens incluem tokens de entrada, tokens de saída e o total tokens consumidos durante a execução 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}")

Os detalhes de utilização dos tokens são exibidos na interface de rastreamento do MLflow, permitindo que você monitore e otimize o desempenho e os custos do seu pipeline.

Desativar rastreamento automático

O rastreamento automático para Haystack pode ser desativado globalmente chamando mlflow.haystack.autolog(disable=True) ou mlflow.autolog(disable=True).