Rastreando AG2
MLflow Tracing integra-se ao AG2 (antigo AutoGen 0.2) para capturar traços unificados de conversas multiagentes e fluxo de trabalho. A integração instrumenta automaticamente os loops do agente e a execução da ferramenta — basta ligar paramlflow.ag2.autolog
.
import mlflow
mlflow.ag2.autolog()
A integração fornece visibilidade abrangente sobre:
- Qual agente é chamado em turnos diferentes
- Mensagens passadas entre agentes
- Chamadas de LLM e ferramentas feitas por cada agente, organizadas por agente e turno
- Latências
- Qualquer exceção, se levantada
Pré-requisitos
Para usar o MLflow Tracing com o AG2, o senhor precisa instalar o MLflow e o pacote AG2 (AutoGen) relevante.
- Development
- Production
Para ambientes de desenvolvimento, instale o pacote completo do MLflow com os extras do Databricks e o AutoGen:
pip install --upgrade "mlflow[databricks]>=3.1" pyautogen
O pacote completo do mlflow[databricks]
inclui todos os recursos para desenvolvimento local e experimentação no Databricks.
Para implantações de produção, instale mlflow-tracing
e AutoGen:
pip install --upgrade mlflow-tracing pyautogen
O pacote mlflow-tracing
é otimizado para uso na produção.
O MLflow 3 é recomendado para obter a melhor experiência de rastreamento com o AG2.
Antes de executar os exemplos, você precisará configurar seu ambiente:
Para usuários fora do Databricks Notebook : Defina seu Databricks variável de ambiente:
export DATABRICKS_HOST="https://your-workspace.cloud.databricks.com"
export DATABRICKS_TOKEN="your-personal-access-token"
Para usuários do Databricks Notebook : Essas credenciais são definidas automaticamente para o senhor.
Exemplo básico
import os
from typing import Annotated, Literal
from autogen import ConversableAgent
import mlflow
# Enable auto-tracing for AG2
mlflow.ag2.autolog()
# Track to Databricks (optional if already configured)
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/ag2-tracing-demo")
# Define a simple multi-agent workflow using AG2 (AutoGen 0.2)
config_list = [
{
"model": "gpt-4o-mini",
# Requires OPENAI_API_KEY in env for this example
"api_key": os.environ.get("OPENAI_API_KEY"),
}
]
Operator = Literal["+", "-", "*", "/"]
def calculator(a: int, b: int, operator: Annotated[Operator, "operator"]) -> int:
if operator == "+":
return a + b
elif operator == "-":
return a - b
elif operator == "*":
return a * b
elif operator == "/":
return int(a / b)
else:
raise ValueError("Invalid operator")
assistant = ConversableAgent(
name="Assistant",
system_message=(
"You are a helpful AI assistant. You can help with simple calculations. "
"Return 'TERMINATE' when the task is done."
),
llm_config={"config_list": config_list},
)
user_proxy = ConversableAgent(
name="Tool Agent",
llm_config=False,
is_termination_msg=lambda msg: msg.get("content") is not None and "TERMINATE" in msg["content"],
human_input_mode="NEVER",
)
assistant.register_for_llm(name="calculator", description="A simple calculator")(calculator)
user_proxy.register_for_execution(name="calculator")(calculator)
response = user_proxy.initiate_chat(
assistant, message="What is (44231 + 13312 / (230 - 20)) * 4?"
)
Rastrear o uso de tokens
MLflow 3.2.0+ suporta o acompanhamento do uso de tokens para a AG2. O uso por chamada é registrado no atributo mlflow.chat.tokenUsage
span; o uso total aparece nas informações de rastreamento.
import mlflow
last_trace_id = mlflow.get_last_active_trace_id()
trace = mlflow.get_trace(trace_id=last_trace_id)
total = trace.info.token_usage
print("Input:", total["input_tokens"], "Output:", total["output_tokens"], "Total:", total["total_tokens"])
for span in trace.data.spans:
usage = span.get_attribute("mlflow.chat.tokenUsage")
if usage:
print(span.name, usage)
Desativar o rastreamento automático
Desative o rastreamento automático do AG2 com mlflow.ag2.autolog(disable=True)
ou desative todo o registro mlflow.autolog(disable=True)
automático com.