Sledovatelnost agenta pomocí trasování MLflow

Tento článek popisuje, jak přidat pozorovatelnost k aplikacím generující AI pomocí trasování MLflow v Databricks.

Co je trasování MLflow?

Sledování MLflow poskytuje komplexní pozorovatelnost pro aplikace generativní umělé inteligence od vývoje po nasazení. Trasování je plně integrované se sadou nástrojů generativní AI od Databricks a zaznamenává podrobné přehledy v celém životním cyklu vývoje a nasazení.

Zde jsou klíčové případy použití trasování v aplikacích generativní umělé inteligence:

• zjednodušené ladění: Trasování poskytuje přehled o jednotlivých krocích aplikace generativní umělé inteligence, což tak usnadňuje diagnostiku a řešení problémů.

• offline vyhodnocení: Trasování generuje cenné údaje pro vyhodnocení agenta, což vám umožní měřit a zlepšit kvalitu agentů v průběhu času.

• monitorování produkčního prostředí: Trasování poskytuje přehled o chování agenta a podrobných krocích provádění, které vám umožní monitorovat a optimalizovat výkon agenta v produkčním prostředí.

• protokoly auditu: Sledování MLflow generuje komplexní protokoly auditu akcí a rozhodnutí agenta. To je důležité pro zajištění dodržování předpisů a podpory ladění v případě, že dojde k neočekávaným problémům.

Požadavky

Trasování MLflow je k dispozici od verze MLflow 2.13.0 a vyšší. Databricks doporučuje nainstalovat nejnovější verzi MLflow pro přístup k nejnovějším funkcím a vylepšením.

%pip install mlflow>=2.13.0 -qqqU
%restart_python

Automatické trasování

automatického zalogování MLflow umožňuje rychle instrumentovat agenta přidáním jednoho řádku do kódu mlflow.<library>.autolog().

MLflow podporuje automatické protokolování pro nejoblíbenější knihovny pro vytváření agentů. Další informace o každé autorské knihovně najdete v dokumentaci k automatickému logování MLflow:

Zakázání automatického protokolování

Trasování v rámci automatického logování je standardně povoleno v prostředí Databricks Runtime 15.4 ML a vyšších verzích pro následující knihovny:

• LangChain
• Langgraph
• OpenAI
• LlamaIndex

Pokud chcete zakázat trasování automatického protokolování těchto knihoven, spusťte v poznámkovém sešitu následující příkaz:

`mlflow.<library>.autolog(log_traces=False)`

Ruční přidání trasování

Ačkoli automatické protokolování poskytuje pohodlný způsob, jak instrumentovat agenty, můžete chtít provádět instrumentaci svého agenta podrobněji nebo přidat další trasování, která automatické protokolování nezachytí. V těchto případech použijte rozhraní API pro sledování MLflow k ručnímu přidání záznamů.

Rozhraní API pro trasování MLflow jsou rozhraní API s nízkým kódem pro přidávání tras, aniž byste se museli starat o správu struktury sledování. MLflow automaticky určuje odpovídající nadřazené a podřízené vztahy pomocí Pythonového stacku.

Kombinování automatického přihlašování a ručního trasování

Rozhraní API pro ruční trasování je možné použít s automatickým zaznamenáváním. MLflow kombinuje rozsahy vytvořené automatickým protokolováním a ručním trasováním a vytvoří kompletní trasování spuštění agenta. Příklad kombinování automatického logování a ručního trasování najdete viz Instrumentace nástroje volajícího agenta pomocí MLflow Tracing.

Trasování funkcí pomocí dekorátoru @mlflow.trace

Nejjednodušší způsob, jak ručně instrumentovat kód, je ozdobit funkci @mlflow.trace dekorátorem. dekorátor sledování MLflow vytvoří "span" v rozsahu dekorované funkce, která představuje jednotku vykonání ve sledování a zobrazí se jako jeden řádek ve vizualizaci sledování. Rozsah zachycuje vstup a výstup funkce, latenci a všechny výjimky vyvolané funkcí.

Například následující kód vytvoří rozsah pojmenovaný my_function, který zachycuje vstupní argumenty x a y a výstup.

import mlflow

@mlflow.trace
def add(x: int, y: int) -> int:
  return x + y

Můžete také přizpůsobit název rozsahu, typ rozsahu a přidat do rozsahu vlastní atributy:

from mlflow.entities import SpanType

@mlflow.trace(
  # By default, the function name is used as the span name. You can override it with the `name` parameter.
  name="my_add_function",
  # Specify the span type using the `span_type` parameter.
  span_type=SpanType.TOOL,
  # Add custom attributes to the span using the `attributes` parameter. By default, MLflow only captures input and output.
  attributes={"key": "value"}
)
def add(x: int, y: int) -> int:
  return x + y

Trasování libovolných bloků kódu pomocí správce kontextu

Pokud chcete vytvořit rozsah pro libovolný blok kódu, nejen pro funkci, použijte mlflow.start_span() jako správce kontextu, který zabalí blok kódu. Interval začíná při vstupu do kontextu a končí při výstupu z kontextu. Vstupy a výstupy span by měly být poskytovány ručně pomocí metod setter objektu span, který poskytuje správce kontextu. Další informace najdete v dokumentaci MLflow (o obslužné rutině kontextu).

with mlflow.start_span(name="my_span") as span:
  span.set_inputs({"x": x, "y": y})
  result = x + y
  span.set_outputs(result)
  span.set_attribute("key", "value")

Knihovny trasování na nižší úrovni

MLflow také poskytuje nízkoúrovňové API pro explicitní kontrolu struktury sledování stromu. Podívejte se na dokumentaci k MLflow - ruční instrumentace.

Příklad trasování : Kombinování automatického a ručního logování

Následující příklad kombinuje automatické protokolování OpenAI a ruční trasování k úplnému instrumentaci agenta pro volání nástrojů.

import json
from openai import OpenAI
import mlflow
from mlflow.entities import SpanType

client = OpenAI()

# Enable OpenAI autologging to capture LLM API calls
# (*Not necessary if you are using the Databricks Runtime 15.4 ML and above, where OpenAI autologging is enabled by default)
mlflow.openai.autolog()

# Define the tool function. Decorate it with `@mlflow.trace` to create a span for its execution.
@mlflow.trace(span_type=SpanType.TOOL)
def get_weather(city: str) -> str:
  if city == "Tokyo":
    return "sunny"
  elif city == "Paris":
    return "rainy"
  return "unknown"


tools = [
  {
    "type": "function",
    "function": {
      "name": "get_weather",
      "parameters": {
        "type": "object",
        "properties": {"city": {"type": "string"}},
      },
    },
  }
]

_tool_functions = {"get_weather": get_weather}

# Define a simple tool-calling agent
@mlflow.trace(span_type=SpanType.AGENT)
def run_tool_agent(question: str):
  messages = [{"role": "user", "content": question}]

  # Invoke the model with the given question and available tools
  response = client.chat.completions.create(
    model="gpt-4o-mini",
    messages=messages,
    tools=tools,
  )
  ai_msg = response.choices[0].message
  messages.append(ai_msg)

  # If the model requests tool calls, invoke the function(s) with the specified arguments
  if tool_calls := ai_msg.tool_calls:
    for tool_call in tool_calls:
      function_name = tool_call.function.name
      if tool_func := _tool_functions.get(function_name):
        args = json.loads(tool_call.function.arguments)
        tool_result = tool_func(**args)
      else:
        raise RuntimeError("An invalid tool is returned from the assistant!")

      messages.append(
        {
          "role": "tool",
          "tool_call_id": tool_call.id,
          "content": tool_result,
        }
      )

    # Send the tool results to the model and get a new response
    response = client.chat.completions.create(
      model="gpt-4o-mini", messages=messages
    )

  return response.choices[0].message.content

# Run the tool calling agent
question = "What's the weather like in Paris today?"
answer = run_tool_agent(question)

Přidávání poznámek k trasování pomocí značek

Značky trasy MLflow jsou páry klíč-hodnota, které umožňují přidat vlastní metadata k trasám, jako je ID konverzace, ID uživatele, hash Git commit, atd. Značky se zobrazují v uživatelském rozhraní MLflow pro filtrování a vyhledávání tras.

Značky je možné nastavit na probíhající nebo dokončené trasování pomocí rozhraní API MLflow nebo uživatelského rozhraní MLflow. Následující příklad ukazuje přidání značky do probíhajícího trasování pomocí rozhraní mlflow.update_current_trace() API.

@mlflow.trace
def my_func(x):
    mlflow.update_current_trace(tags={"fruit": "apple"})
    return x + 1

Další informace o označování trasování a jejich použití k filtrování a vyhledávání trasování najdete v dokumentaci k MLflow – Nastavení značek trasování.

Kontrola trasování

Pokud chcete zkontrolovat trasování po spuštění agenta, použijte jednu z následujících možností:

• Vizualizace v řádku: V poznámkových blocích Databricks se tras vykreslí přímo ve výstupu buňky.
• experimentu MLflow: V Databricks přejděte na Experimenty> Vyberte experiment >Trasování a prohlédněte si všechny trasování experimentu.
• spuštění MLflow: Když agent běží během aktivního spuštění MLflow, stopy se zobrazí na stránce běhu v uživatelském rozhraní MLflow.
• Uživatelské rozhraní pro vyhodnocení agenta: Ve vyhodnocení agenta systému Mosaic AI můžete zkontrolovat trasování jednotlivých spuštění agentů kliknutím na Zobrazit podrobnosti trasování ve výsledcích vyhodnocení.
• rozhraní API pro vyhledávání tras: Pokud chcete trasování načíst prostřednictvím kódu programu, použijte rozhraní API pro vyhledávání tras .

Vyhodnocení agentů pomocí stop

Data trasování slouží jako cenný prostředek pro vyhodnocení agentů. Zachycením podrobných informací o provádění modelů je MLflow Tracing nástrojem pro offline vyhodnocení. Data trasování můžete použít k vyhodnocení výkonu vašeho agenta vůči zlaté datové sadě, identifikaci problémů a zlepšení výkonu agenta.

%pip install -U mlflow databricks-agents
%restart_python

import mlflow

# Get the recent 50 successful traces from the experiment
traces = mlflow.search_traces(
    max_results=50,
    filter_string="status = 'OK'",
)

traces.drop_duplicates("request", inplace=True) # Drop duplicate requests.
traces["trace"] = traces["trace"].apply(lambda x: x.to_json()) # Convert the trace to JSON format.

# Evaluate the agent with the trace data
mlflow.evaluate(data=traces, model_type="databricks-agent")

Další informace o vyhodnocení agenta najdete v tématu Spuštění vyhodnocení a zobrazení výsledků.

Monitorování nasazených agentů s odvozovacími tabulkami

Jakmile je agent nasazen do Mosaic AI Model Serving, můžete k monitorování agenta použít inference tabulky. Tabulky odvozování obsahují podrobné protokoly požadavků, odpovědí, trasování agentů a zpětné vazby agentů z aplikace pro kontrolu. Tyto informace vám umožní ladit problémy, monitorovat výkon a vytvořit zlatou datovou sadu pro offline vyhodnocení.

Pokud chcete povolit tabulky odvozování pro nasazení agenta, přečtěte si téma Povolení odvozovacích tabulek pro agenty AI.

Dotazování online trasování

K dotazování tabulky odvozování a analýze výsledků použijte poznámkový blok.

Pokud chcete vizualizovat trasování, spusťte display(<the request logs table>) a vyberte řádky, které chcete zkontrolovat:

# Query the inference table
df = spark.sql("SELECT * FROM <catalog.schema.my-inference-table-name>")
display(df)

Monitorování agentů

Podívejte se na Jak sledovat vaši aplikaci Gen AI.

Trasování latence nadměrné zátěže

Trasování se zapisuje asynchronně, aby se minimalizoval dopad na výkon. Trasování ale stále přidává latenci rychlosti odezvy koncového bodu, zejména pokud je velikost trasování pro každý požadavek na odvozování velká. Databricks doporučuje otestovat koncový bod, abyste pochopili, jak trasování ovlivňuje latenci, ještě před nasazením do produkčního prostředí.

Následující tabulka obsahuje přibližné odhady dopadu na latenci podle velikosti sledování:

Řešení problémů

Informace o řešení potíží a běžné dotazy najdete v dokumentaci k MLflow: Průvodce trasováním a MLflow dokumentace: Nejčastější dotazy