Agent Evaluation (MLflow 2) : métriques personnalisées, directives et étiquettes d'experts du domaine

Ouvrir dans Databricks

Remarque Ce notebook décrit l'Agent Evaluation de MLflow 2. Databricks recommande d'utiliser MLflow 3 pour l'évaluation et le monitoring des applications d'IA. Pour plus d'informations sur MLflow 3, consultez l'évaluation et le monitoring sur MLflow 3 et la migration vers MLflow 3.

Ce Notebook démontre comment évaluer une application d'IA en utilisant les juges LLM propriétaires d'Agent Evaluation, des métriques personnalisées et des étiquettes d'experts du domaine. Il démontre :

• Comment charger les Logs de production (traces) dans un dataset d'évaluation.
• Comment exécuter une évaluation et effectuer une analyse des causes profondes.
• Comment créer des métriques personnalisées pour détecter automatiquement les problèmes de qualité.
• Comment envoyer les logs de production aux experts en la matière pour l'étiquetage et l'évolution du dataset d'évaluation.

Pour préparer votre agent à la préproduction, consultez le guide de démarrage rapide de l'agent.

Pour des informations générales sur l'évaluation des agents dans MLflow 2, consultez la documentation sur l'Agent Evaluation.

Exigences

• Voir les exigences d'Agent Evaluation.
• Cluster Serverless ou classique exécutant Databricks Runtime 15.4 LTS ou version ultérieure, ou Databricks Runtime for Machine Learning 15.4 LTS ou version ultérieure.
• Accès CREATE TABLE dans un schéma Unity Catalog

Python

%pip install -U -qqqq 'mlflow>=2.20.3' 'langchain==0.3.20' 'langgraph==0.3.4' 'databricks-langchain>=0.3.0' pydantic 'databricks-agents>=0.17.2' uv databricks-sdk
dbutils.library.restartPython()

Sélectionnez un schéma Unity Catalog.

Assurez-vous de disposer des autorisations CREATE TABLE dans ce schéma. Par default, ces valeurs sont définies sur le catalogue et le schéma par default de votre workspace.

Python

# Get the workspace default UC catalog / schema
uc_default_location = spark.sql("select current_catalog() as current_catalog, current_schema() as current_schema").collect()[0]
current_catalog = uc_default_location["current_catalog"]
current_schema = uc_default_location["current_schema"]


# Modify the UC catalog / schema here or at the top of the notebook in the widget editor
dbutils.widgets.text("uc_catalog", current_catalog)
dbutils.widgets.text("uc_schema", current_schema)
UC_CATALOG = dbutils.widgets.get("uc_catalog")
UC_SCHEMA = dbutils.widgets.get("uc_schema")
UC_PREFIX = f"{UC_CATALOG}.{UC_SCHEMA}"

Agent simple d'appel d'outils

La cellule suivante définit un agent simple d’appel d’outil, créé avec LangGraph, qui dispose de 2 outils :

1. multiply, qui prend 2 nombres et les multiplie
2. query_docs, qui prend un ensemble de mots-clés et renvoie des documents pertinents sur Databricks à l'aide de la recherche par mot-clé.

Aux fins de ce Notebook de démo, il n'est pas important de savoir comment le code de l'Agent fonctionne – cette démo se concentre sur l'évaluation de la qualité de l'Agent.

Remarque : Agent Evaluation fonctionne avec toute application d’IA, quelle que soit sa conception, tant que l’application peut accepter une entrée Dict[str, Any] et retourne une sortie Dict[str, Any].

Pour plus d'exemples d'outils à ajouter à votre agent, consultez la documentation sur les outils d'agent.

Python

from typing import Any, Generator, Optional, Sequence, Union
from langchain_core.tools import tool
import mlflow
from databricks_langchain import ChatDatabricks
from langchain_core.language_models import LanguageModelLike
from langchain_core.runnables import RunnableConfig, RunnableLambda
from langchain_core.tools import BaseTool
from langgraph.graph import END, StateGraph
from langgraph.graph.graph import CompiledGraph
from langgraph.graph.state import CompiledStateGraph
from langgraph.prebuilt.tool_node import ToolNode
from mlflow.langchain.chat_agent_langgraph import ChatAgentState, ChatAgentToolNode
from mlflow.pyfunc import ChatAgent
from mlflow.types.agent import (
    ChatAgentChunk,
    ChatAgentMessage,
    ChatAgentResponse,
    ChatContext,
)
import pandas as pd

mlflow.langchain.autolog()
LLM_ENDPOINT_NAME = "databricks-meta-llama-3-3-70b-instruct"


# Example docs in our vector store.
DOCS = [
    mlflow.entities.Document(
        metadata={"doc_uri": "uri1.txt"},
        page_content="Databricks has managed MLFlow, which has Tracing for observing any GenAI application",
    )
]

SYSTEM_PROMPT = "You are an assistant that answers user's questions by calling tools.  Always try to answer the user's question!"

@tool
def multiply(a: int, b: int) -> int:
    """Multiply two numbers."""
    return a * b

@tool
@mlflow.trace(span_type="RETRIEVER")
def query_docs(keywords: list[str]) -> list[mlflow.entities.Document]:
    """
    Use this tool to search for Databricks product documentation.

    Args:
      keywords: a set of individual keywords to find relevant docs for. Each item of the array must be a single word.
    Returns:
    A list of documents that match the keywords.
    """
    if len(keywords) == 0:
        return []
    result = []
    for doc in DOCS:
        score = sum(
            (keyword.lower() in doc.page_content.lower())
            for keyword in keywords
        )
        result.append({
                "page_content": doc.page_content,
                "metadata": {
                    "doc_uri": doc.metadata["doc_uri"],
                    "score": score,
                },
        })
    ranked_docs = sorted(result, key=lambda x: x["metadata"]["score"], reverse=True)
    cutoff_docs = []
    context_budget_left = 8_000
    for doc in ranked_docs:
        content = doc["page_content"]
        doc_len = len(content)
        if context_budget_left < doc_len:
            cutoff_docs.append(
                {**doc, "page_content": content[:context_budget_left]}
            )
            break
        else:
            cutoff_docs.append(doc)
        context_budget_left -= doc_len
    return cutoff_docs

def create_tool_calling_agent(
    model: LanguageModelLike,
    tools: Union[ToolNode, Sequence[BaseTool]],
    system_prompt: Optional[str] = None,
) -> CompiledGraph:
    model = model.bind_tools(tools)

    # Define the function that determines which node to go to
    def should_continue(state: ChatAgentState):
        messages = state["messages"]
        last_message = messages[-1]
        # If there are function calls, continue. else, end
        if last_message.get("tool_calls"):
            return "continue"
        else:
            return "end"

    if system_prompt:
        preprocessor = RunnableLambda(
            lambda state: [{"role": "system", "content": system_prompt}]
            + state["messages"]
        )
    else:
        preprocessor = RunnableLambda(lambda state: state["messages"])
    model_runnable = preprocessor | model

    def call_model(
        state: ChatAgentState,
        config: RunnableConfig,
    ):
        response = model_runnable.invoke(state, config)

        return {"messages": [response]}

    workflow = StateGraph(ChatAgentState)

    workflow.add_node("agent", RunnableLambda(call_model))
    workflow.add_node("tools", ChatAgentToolNode(tools))

    workflow.set_entry_point("agent")
    workflow.add_conditional_edges(
        "agent",
        should_continue,
        {
            "continue": "tools",
            "end": END,
        },
    )
    workflow.add_edge("tools", "agent")

    return workflow.compile()

class LangGraphChatAgent(ChatAgent):
    def __init__(self, agent: CompiledStateGraph):
        self.agent = agent

    def predict(
        self,
        messages: list[ChatAgentMessage],
        context: Optional[ChatContext] = None,
        custom_inputs: Optional[dict[str, Any]] = None,
    ) -> ChatAgentResponse:
        request = {"messages": self._convert_messages_to_dict(messages)}

        messages = []
        for event in self.agent.stream(request, stream_mode="updates"):
            for node_data in event.values():
                messages.extend(
                    ChatAgentMessage(**msg) for msg in node_data.get("messages", [])
                )
        return ChatAgentResponse(messages=messages)

    def predict_stream(
        self,
        messages: list[ChatAgentMessage],
        context: Optional[ChatContext] = None,
        custom_inputs: Optional[dict[str, Any]] = None,
    ) -> Generator[ChatAgentChunk, None, None]:
        request = {"messages": self._convert_messages_to_dict(messages)}
        for event in self.agent.stream(request, stream_mode="updates"):
            for node_data in event.values():
                yield from (
                    ChatAgentChunk(**{"delta": msg}) for msg in node_data["messages"]
                )

tools = [multiply, query_docs]
llm = ChatDatabricks(endpoint=LLM_ENDPOINT_NAME)
agent = create_tool_calling_agent(llm, tools, SYSTEM_PROMPT)
AGENT = LangGraphChatAgent(agent)

Sélectionner les Logs de (pré)production

Ce Notebook de démonstration génère des Logs de production d'exemple pour démontrer les nouvelles fonctionnalités d'Agent Evaluation. Normalement, ces Logs proviendraient d'un agent (de pré-)production. La cellule suivante appelle directement l'agent et enregistre les traces dans MLflow.

REMARQUE : Le traçage MLflow visualise chaque trace (avec pagination) dans la sortie de la cellule lorsque l'agent est appelé ou que les traces sont récupérées à l'aide de mlflow.search_traces.

Après avoir terminé le Notebook, si vous avez déjà un agent déployé sur Databricks, localisez le request_ids à examiner dans la table d'inférence <model_name>_payload_request_logs. La table d'inférence se trouve dans le même catalogue et schéma Unity Catalog où le modèle a été enregistré. L'exemple de code pour cela se trouve près du bas de ce notebook.

Python

import mlflow

# Fake production logs.  Normally, these would come from a (pre-)production agent, but for this demo, they are generated here.
examples = [
    "How much is 423 * 124",
    "If I go to the store 13 times and go 3 more times, how many visits did I do?",
    "Does Databricks have GenAI observability?",
    "Does Databricks support spark 3.5?",
    "How do I get a discount on Databricks?"
]

# The following code calls the agent and logs the traces in an MLflow run. These traces become the evaluation dataset.
with mlflow.start_run(run_name="example-production-logs") as run:
    for example in examples:
        AGENT.predict({"messages": [{"role": "user", "content": example}]})

requests = mlflow.search_traces(run_id=run.info.run_id)

Charger les traces dans un dataset d'évaluation

Important : avant d'exécuter cette cellule, assurez-vous que les valeurs des widgets uc_catalog et uc_schema sont définies sur un schéma Unity Catalog pour lequel vous disposez des autorisations CREATE TABLE. La réexécution de cette cellule recréera le dataset d'évaluation.

Python

from databricks.agents import datasets
from databricks.sdk.errors.platform import NotFound

# Make sure you have updated the uc_catalog & uc_schema widgets to a valid catalog/schema where you have CREATE TABLE permissions.
UC_TABLE_NAME = f'{UC_PREFIX}.agent_evaluation_set'

# Remove the evaluation dataset if it already exists
try:
  datasets.delete_dataset(UC_TABLE_NAME)
except NotFound:
  pass

# Create the evaluation dataset
dataset = datasets.create_dataset(UC_TABLE_NAME)

# Add the traces from the production logs gathered in the previous cell.
dataset.insert(requests)

# Show the resulting evaluation set
display(spark.table(UC_TABLE_NAME))

Exécuter une évaluation

Juges intégrés d'Agent Evaluation

• Juges qui s'exécutent sans étiquettes de vérité terrain ni récupération dans les traces :

  • guidelines: Permet aux développeurs d'écrire des listes de contrôle ou des grilles d'évaluation en langage clair, améliorant la transparence et la confiance avec les parties prenantes de l'entreprise grâce à des grilles d'évaluation structurées et faciles à comprendre.
  • safety: Vérifie que la réponse est sûre.
  • relevance_to_query: vérifie que la réponse est pertinente.

• Pour les traces avec des documents récupérés (étendues de type RETRIEVER) :

  • groundedness: Détecte les hallucinations.
  • chunk_relevance: Pertinence au niveau du segment par rapport à la query.

• Une fois les étiquettes de vérité terrain collectées à l’aide de l’application Review, deux juges supplémentaires deviennent disponibles :

  • correctness: Ignoré jusqu'à ce que des étiquettes comme expected_facts soient collectées.
  • context_sufficiency: Ignoré jusqu'à ce que des étiquettes comme expected_facts soient collectées.

Consultez la liste complète des juges intégrés et découvrez comment exécuter un sous-ensemble de juges ou personnaliser les juges.

Métriques personnalisées

• Vérifiez la qualité des appels d'outils :

  • tool_calls_are_logical: Affirme que les outils sélectionnés dans la trace étaient logiques compte tenu de la demande de l'utilisateur.
  • grounded_in_tool_outputs: Affirme que les réponses du LLM sont basées sur les sorties des outils et ne sont pas hallucinantes.

• Mesurer le coût et la latence de l'agent :

  • latency: Extrait la latence de la trace MLflow.
  • cost: Extrait le nombre total de jetons utilisés et multiplie par le taux de jetons du LLM.

Ce notebook crée des métriques personnalisées qui utilisent des juges Databricks appelables. Les métriques personnalisées peuvent être n'importe quelle fonction Python. Pour plus d'exemples de métriques personnalisées, consultez la référence du juge LLM.

Définir les métriques personnalisées

Python

from databricks.agents.evals import judges
from mlflow.evaluation import Assessment
from databricks.agents.evals import metric
from mlflow.entities import SpanType


@metric
def tool_calls_are_logical(request, tool_calls):
    # If no tool calls, don't run this metric
    if len(tool_calls) == 0:
        return None

    # This assumes that the tools available to the FIRST LLM call are the same as what is presented to all other LLM calls.  Adjust if this doesn't hold true for a given use case.
    available_tools = tool_calls[0].available_tools

    # Get ALL called tools across ALL LLM calls - this will happen if the LLM does multiple iterations to call tools (e.g., calls a set of tools & then decides to call more tools based on that output)
    requested_tools = []
    for item in tool_calls:
        requested_tools.append(
            {"tool_name": item.tool_name, "tool_call_args": item.tool_call_args}
        )

    is_logical = judges.guideline_adherence(
        request=f"User's request: {request}\nAvailable tools: {available_tools}",
        response=str(requested_tools),
        guidelines=[
            "The response is a set of selected tool calls. The selected tools must be logical, given the user's request."
        ],
    )
    # See https://docs.databricks.com/aws/en/generative-ai/agent-evaluation/llm-judge-reference#examples-6 or
    # https://learn.microsoft.com/en-us/azure/databricks/generative-ai/agent-evaluation/llm-judge-reference#examples-6
    return Assessment(
        name="tool_calls_are_logical",
        value=is_logical.value,
        rationale=is_logical.rationale,
    )


@metric
def grounded_in_tool_outputs(request, response, tool_calls):
    # If no tool calls, don't run this metric
    if len(tool_calls) == 0:
        return None

    # Customize the built-in groundedness judge for the tool calling outputs
    tool_outputs = [{'result': t.tool_call_result["content"], 'args': t.tool_call_args, 'name': t.tool_name} for t in tool_calls]
    contexts = []

    # Format the tool calls as "Called tool tool_name(param1=value, param2=value) that returned ```return value```"".
    for tool in tool_outputs:
        args_str = ', '.join(f"{k}={v}" for k, v in tool['args'].items())
        contexts.append(f"Called tool `{tool['name']}({args_str})` that returned ```{tool['result']}```")


    context_to_evaluate = "\n".join(contexts)

    # Extract the user's request & LLM's response
    user_request = next(item for item in request['messages'] if item['role'] == 'user')['content']
    assistant_response = response['messages'][-1]["content"]

    # Create a guidelines judge to evaluate if the assistant's response is grounded in the context of the tool calls.
    out = judges.guideline_adherence(
        request=f"<user_request>{user_request}<user_request><context_to_evaluate>{context_to_evaluate}<context_to_evaluate>",
        response=f"<assistant_response>{assistant_response}<assistant_response>",
        guidelines=["The <assistant_response>'s to the <user_request> must be grounded in the <context_to_evaluate> which represent tools that were called when trying to answer the <user_request>."]
    )

    return Assessment(
        name="grounded_in_tool_outputs", value=out.value, rationale=out.rationale
    )


@metric
def is_answer_relevant(request, response):
    # Extract the user's request & LLM's response
    user_request = next(item for item in request['messages'] if item['role'] == 'user')['content']
    assistant_response = response['messages'][-1]["content"]

    # Use the guideline's judge to assess the relevance of the LLM's response.  This approach (rather than the built-in answer_relevance judge) accounts for the fact that the LLM may (correctly) refuse to answer a question that violates the defined policies.
    out = judges.guideline_adherence(
        request=request,
        response=assistant_response,
        guidelines=["Determine if the response provides an answer to the user's request.  A refusal to answer is considered relevant.  However, if the response is NOT a refusal BUT also doesn't provide relevant information, then the answer is not relevant."]
    )
    return Assessment(
        name="is_answer_relevant", value=out.value, rationale=out.rationale
    )

@metric
def latency(trace):
    return trace.info.execution_time_ms / 1000

@metric
def cost(trace):
    INPUT_TOKEN_COST = 2 # per 1M tokens
    OUTPUT_TOKEN_COST = 15 # per 1M tokens
    input_tokens = trace.search_spans(span_type=SpanType.CHAT_MODEL)[0].outputs['llm_output']['prompt_tokens']
    output_tokens = trace.search_spans(span_type=SpanType.CHAT_MODEL)[0].outputs['llm_output']['completion_tokens']
    cost = ((input_tokens/1000000) * INPUT_TOKEN_COST) + ((output_tokens/1000000) * OUTPUT_TOKEN_COST)
    return round(cost, 3)

Exécuter l’évaluation

Python

# Define global guidelines.  Guidelines are plain language
guidelines = {'pricing': ["The agent should always refuse to answer questions about product pricing; it should never provide anything more than 'I can't talk about pricing'."]}

with mlflow.start_run(run_name="eval-prod-logs"):
    eval_results = mlflow.evaluate(
        # Each row["inputs"] from the dataset is passed to the model. Any dict[str, Any] is supported as inputs.
        model=lambda inputs: AGENT.predict(inputs),
        data=spark.table(UC_TABLE_NAME),
        model_type="databricks-agent",
        # Enable custom metrics
        extra_metrics=[grounded_in_tool_outputs, tool_calls_are_logical, is_answer_relevant, latency, cost],

        # Configure which built-in judges are used and customize the guidelines used
        evaluator_config={
            "databricks-agent": {"global_guidelines": guidelines, "metrics": [
                "chunk_relevance", # Check if the retrieved documents are relevant to the user's query
                "guideline_adherence", # Run the global guidelines defined in `guidelines`
                # Disable the built-in groundedness & relevance judge in favor of the custom-defined version of these metrics
                # "groundedness",
                # "relevance_to_query",
                "safety", # Check if the LLM's response has any toxicity
                # context_sufficiency & correctness require labeled ground truth, which is collected later in this notebook, so they are disabled for now.
                # "context_sufficiency",
                # "correctness",
            ],},

        },
    )
    # Review the evaluation results in the MLflow UI (see console output), or access them in place:
    display(eval_results.tables["eval_results"])

Problèmes détectés

Les résultats de l’évaluation révèlent quelques problèmes :

• L'agent a appelé l'outil multiply lorsque la query nécessitait une sommation.
• La question concernant Spark n'est pas représentée dans le dataset, et le juge chunk_relevance a détecté ce problème.
• Le LLM répond à des questions tarifaires, ce qui enfreint la consigne.

L'agent a correctement utilisé l'outil multiplication et l'outil query_docs pour les 2 autres requêtes.

Corrigez les problèmes et réévaluez

Maintenant qu'il existe un ensemble d'évaluation avec des juges à essayer, corrigez les problèmes en :

• Amélioration de l'invite système pour indiquer à l'agent que l'absence d'appels d'outils est acceptable.
• Ajout d’un document à la base de connaissances concernant la dernière version de Spark.
• Ajout d'un nouvel outil.

Python

SYSTEM_PROMPT_v2="""You are an assistant that answers user's questions by calling tools.  Only call a tool if it directly helps with the request.  If the user asks about product pricing or discounts, state 'I can't talk about pricing'."""

DOCS = [
    mlflow.entities.Document(
        metadata={"doc_uri": "uri1.txt"},
        page_content="Databricks has managed MLFlow, which has Tracing for observing any GenAI application",
    ),
    # This is a new document about spark.
    mlflow.entities.Document(
        metadata={"doc_uri": "uri2.txt"},
        page_content="The latest spark version in databricks in 3.5.0",
    )
]

@tool
def add(a: int, b: int) -> int:
    """Adds two numbers."""
    return a + b

tools_v2 = [multiply, query_docs, add]

agent_v2 = create_tool_calling_agent(llm, tools_v2, SYSTEM_PROMPT_v2)
AGENT_v2 = LangGraphChatAgent(agent_v2)

with mlflow.start_run(run_name="updated-model") as run:
    eval_results = mlflow.evaluate(
        # Each row["inputs"] from the dataset is passed to the model. Any dict[str, Any] is supported as inputs.
        model=lambda inputs: AGENT_v2.predict(inputs),
        data=spark.table(UC_TABLE_NAME),
        model_type="databricks-agent",
        # Enable custom metrics
        extra_metrics=[grounded_in_tool_outputs, tool_calls_are_logical, is_answer_relevant, latency, cost],

        # Configure which built-in judges are used and customize the guidelines used
        evaluator_config={
            "databricks-agent": {"global_guidelines": guidelines, "metrics": [
                "chunk_relevance", # Check if the retrieved documents are relevant to the user's query
                "guideline_adherence", # Run the global guidelines defined in `guidelines`
                # "groundedness", # Disable the built-in groundedness in favor of the custom-defined version
                # "relevance_to_query", # Check if the LLM's response is relevant to the user's query
                "safety", # Check if the LLM's response has any toxicity
                # context_sufficiency & correctness require labeled ground truth, which is collected later in this notebook, so they are disabled for now.
                # "context_sufficiency",
                # "correctness",
            ],},

        },
    )
    display(eval_results.tables["eval_results"])

Collecter les attentes (libellés de vérité terrain)

Après avoir amélioré l’agent, assurez-vous que certaines réponses contiennent toujours les faits exacts.

Utilisez l'application d’examen pour envoyer des évaluations à une session d’étiquetage pour que les PME fournissent :

• expected_facts pour activer les juges correctness et context_sufficiency.
• guidelines afin que les PME puissent ajouter des critères en langage clair pour chaque question en fonction de leur contexte métier. Cela étend les directives déjà définies au niveau global.
• Les experts ont-ils apprécié la réponse ? Ainsi, les parties prenantes peuvent être sûres que le nouveau modèle est meilleur. Ceci utilise un schéma d'étiquette personnalisé.

**Remarque** : cette session d'étiquetage utilise des traces précalculées de l'exécution d'évaluation précédente, au lieu d'un *agent en direct*. Voir la fin du Notebook pour savoir comment déployer un agent sur Databricks.

Python

from databricks.agents import review_app

# OPTIONAL: Update the assigned_users widget with a comma separated list of users to assign the review app to.
# If not provided, only the user running this notebook will be granted access to the review app.
ASSIGNED_USERS = []

my_review_app = review_app.get_review_app()

my_review_app.create_label_schema(
  name="good_response",
  # Type can be "expectation" or "feedback".
  type="feedback",
  title="Is this a good response?",
  input=review_app.label_schemas.InputCategorical(options=["Yes", "No"]),
  instruction="Optional: provide a rationale below.",
  enable_comment=True,
  overwrite=True
)

my_session = my_review_app.create_labeling_session(
    name="collect_facts",
    assigned_users=ASSIGNED_USERS, # If not provided, only the user running this notebook will be granted access
    # Built-in labeling schemas: EXPECTED_FACTS, GUIDELINES, EXPECTED_RESPONSE
    label_schemas=[review_app.label_schemas.GUIDELINES,review_app.label_schemas.EXPECTED_FACTS,  "good_response"],
)

traces_from_the_updated_model = mlflow.search_traces(run_id=run.info.run_id)
my_session.add_traces(traces_from_the_updated_model)

# Share with the SME.
print("Review App URL:", my_review_app.url)
print("Labeling session URL: ", my_session.url)

Réévaluation avec les données collectées expected_facts

Une fois que les PME ont terminé l'étiquetage, synchronisez les étiquettes dans le dataset d'évaluation et réévaluez. Le juge correctness s'exécute pour toute ligne d'évaluation avec expected_facts.

Python

# Check the progress of the labeling session by selecting traces associated with the labeling session run.
def is_response_good(assessments):
    for assessment in assessments:
        if assessment.name == "good_response":
            return assessment.feedback.value == "Yes"
    return None

# View how many labels the SME provided.
traces = mlflow.search_traces(run_id=my_session.mlflow_run_id)
response_values = traces["assessments"].apply(is_response_good).value_counts(dropna=False)
print(
    f"Got {response_values.get(True, 0)} good responses, "
    f"{response_values.get(False, 0)} bad responses, and "
    f"{response_values.get(None, 0)} not yet labeled.")

# Move the SME's labels to the evaluation dataset created earlier.
my_session.sync_expectations(to_dataset=UC_TABLE_NAME)

with mlflow.start_run(run_name="with-human-labels") as run:
    eval_results = mlflow.evaluate(
        # Each row["inputs"] from the dataset is passed to the model. Any dict[str, Any] is supported as inputs.
        model=lambda inputs: AGENT_v2.predict(inputs),
        data=spark.table(UC_TABLE_NAME),
        model_type="databricks-agent",
        # Enable custom metrics
        extra_metrics=[grounded_in_tool_outputs, tool_calls_are_logical, is_answer_relevant, latency, cost],

        # Configure which built-in judges are used and customize the guidelines used
        evaluator_config={
            "databricks-agent": {"global_guidelines": guidelines, "metrics": [
                "chunk_relevance", # Check if the retrieved documents are relevant to the user's query
                "guideline_adherence", # Run the global guidelines defined in `guidelines`
                # "groundedness", # Disable the built-in groundedness in favor of the custom-defined version
                # "relevance_to_query", # Check if the LLM's response is relevant to the user's query
                "safety", # Check if the LLM's response has any toxicity
                # context_sufficiency & correctness can now be enabled since labeled ground truth has been collected.
                "context_sufficiency",
                "correctness",
            ],},

        },
    )
    display(eval_results.tables["eval_results"])

Facultatif : Déploiement de l'Agent dans Databricks

Enregistrer l'agent en tant que modèle MLflow

Stockez le dernier agent dans un fichier agent.py autonome et enregistrez-le en tant que code. Voir MLflow - Modèles à partir du code.

Python

%%writefile agent.py
from typing import Any, Generator, Optional, Sequence, Union
from langchain_core.tools import tool
import mlflow
from databricks_langchain import ChatDatabricks
from langchain_core.language_models import LanguageModelLike
from langchain_core.runnables import RunnableConfig, RunnableLambda
from langchain_core.tools import BaseTool
from langgraph.graph import END, StateGraph
from langgraph.graph.graph import CompiledGraph
from langgraph.graph.state import CompiledStateGraph
from langgraph.prebuilt.tool_node import ToolNode
from mlflow.langchain.chat_agent_langgraph import ChatAgentState, ChatAgentToolNode
from mlflow.pyfunc import ChatAgent
from mlflow.types.agent import (
    ChatAgentChunk,
    ChatAgentMessage,
    ChatAgentResponse,
    ChatContext,
)

mlflow.langchain.autolog()
LLM_ENDPOINT_NAME = "databricks-meta-llama-3-3-70b-instruct"
# Example docs in our vector store.
DOCS = [
    mlflow.entities.Document(
        metadata={"doc_uri": "uri1.txt"},
        page_content="Databricks has managed MLFlow, which has Tracing for observing any GenAI application",
    ),
    # This is a new document about spark.
    mlflow.entities.Document(
        metadata={"doc_uri": "uri2.txt"},
        page_content="The latest spark version in databricks in 3.5.0",
    )
]
SYSTEM_PROMPT="""You are an assistant that answers user's questions by calling tools.  Only call a tool if it directly helps with the request.  If the user asks about product pricing or discounts, state 'I can't talk about pricing'."""

@tool
def add(a: int, b: int) -> int:
    """Adds two numbers."""
    return a + b

@tool
def multiply(a: int, b: int) -> int:
    """Multiply two numbers."""
    return a * b

@tool
@mlflow.trace(span_type="RETRIEVER")
def query_docs(keywords: list[str]) -> list[mlflow.entities.Document]:
    """
    Use this tool to search for Databricks product documentation.

    Args:
      keywords: a set of individual keywords to find relevant docs for. Each item of the array must be a single word.
    Returns:
    A list of documents that match the keywords.
    """
    if len(keywords) == 0:
        return []
    result = []
    for doc in DOCS:
        score = sum(
            (keyword.lower() in doc.page_content.lower())
            for keyword in keywords
        )
        result.append({
                "page_content": doc.page_content,
                "metadata": {
                    "doc_uri": doc.metadata["doc_uri"],
                    "score": score,
                },
        })
    ranked_docs = sorted(result, key=lambda x: x["metadata"]["score"], reverse=True)
    cutoff_docs = []
    context_budget_left = 8_000
    for doc in ranked_docs:
        content = doc["page_content"]
        doc_len = len(content)
        if context_budget_left < doc_len:
            cutoff_docs.append(
                {**doc, "page_content": content[:context_budget_left]}
            )
            break
        else:
            cutoff_docs.append(doc)
        context_budget_left -= doc_len
    return cutoff_docs

def create_tool_calling_agent(
    model: LanguageModelLike,
    tools: Union[ToolNode, Sequence[BaseTool]],
    system_prompt: Optional[str] = None,
) -> CompiledGraph:
    model = model.bind_tools(tools)

    # Define the function that determines which node to go to
    def should_continue(state: ChatAgentState):
        messages = state["messages"]
        last_message = messages[-1]
        # If there are function calls, continue. else, end
        if last_message.get("tool_calls"):
            return "continue"
        else:
            return "end"

    if system_prompt:
        preprocessor = RunnableLambda(
            lambda state: [{"role": "system", "content": system_prompt}]
            + state["messages"]
        )
    else:
        preprocessor = RunnableLambda(lambda state: state["messages"])
    model_runnable = preprocessor | model

    def call_model(
        state: ChatAgentState,
        config: RunnableConfig,
    ):
        response = model_runnable.invoke(state, config)

        return {"messages": [response]}

    workflow = StateGraph(ChatAgentState)

    workflow.add_node("agent", RunnableLambda(call_model))
    workflow.add_node("tools", ChatAgentToolNode(tools))

    workflow.set_entry_point("agent")
    workflow.add_conditional_edges(
        "agent",
        should_continue,
        {
            "continue": "tools",
            "end": END,
        },
    )
    workflow.add_edge("tools", "agent")

    return workflow.compile()

class LangGraphChatAgent(ChatAgent):
    def __init__(self, agent: CompiledStateGraph):
        self.agent = agent

    def predict(
        self,
        messages: list[ChatAgentMessage],
        context: Optional[ChatContext] = None,
        custom_inputs: Optional[dict[str, Any]] = None,
    ) -> ChatAgentResponse:
        request = {"messages": self._convert_messages_to_dict(messages)}

        messages = []
        for event in self.agent.stream(request, stream_mode="updates"):
            for node_data in event.values():
                messages.extend(
                    ChatAgentMessage(**msg) for msg in node_data.get("messages", [])
                )
        return ChatAgentResponse(messages=messages)

    def predict_stream(
        self,
        messages: list[ChatAgentMessage],
        context: Optional[ChatContext] = None,
        custom_inputs: Optional[dict[str, Any]] = None,
    ) -> Generator[ChatAgentChunk, None, None]:
        request = {"messages": self._convert_messages_to_dict(messages)}
        for event in self.agent.stream(request, stream_mode="updates"):
            for node_data in event.values():
                yield from (
                    ChatAgentChunk(**{"delta": msg}) for msg in node_data["messages"]
                )

tools = [multiply, query_docs, add]
llm = ChatDatabricks(endpoint=LLM_ENDPOINT_NAME)
agent = create_tool_calling_agent(llm, tools, SYSTEM_PROMPT)
# print(agent.invoke({"messages": [{"role": "user", "content": "What is 423 * 124"}]}))
AGENT = LangGraphChatAgent(agent)
mlflow.models.set_model(AGENT)

Python

import mlflow
from mlflow.models.resources import DatabricksServingEndpoint

resources = [DatabricksServingEndpoint(endpoint_name=LLM_ENDPOINT_NAME)]

with mlflow.start_run():
    logged_agent_info = mlflow.pyfunc.log_model(
        artifact_path="agent",
        python_model="agent.py",
        pip_requirements=[
            "mlflow",
            "langchain",
            "langgraph==0.3.4",
            "databricks-langchain",
            "pydantic",
        ],
        resources=resources,
    )

Enregistrer le modèle dans Unity Catalog et le déployer

Python

from databricks import agents

mlflow.set_registry_uri("databricks-uc")
UC_MODEL_NAME = f"{UC_PREFIX}.agent_model"

uc_registered_model_info = mlflow.register_model(
    model_uri=logged_agent_info.model_uri, name=UC_MODEL_NAME
)
deployment = agents.deploy(UC_MODEL_NAME, uc_registered_model_info.version, tags = {"endpointSource": "agent-eval-demo"}, deploy_feedback_model=False)

Étiqueter un agent actif

Créez une autre session d'étiquetage qui communique avec l'agent nouvellement déployé. Au lieu d'ajouter des traces, ajoutez le dataset d'évaluation à la session. L'appel de add_agent() active également le mode de chat en direct de l'application d'évaluation, ce qui permet aux utilisateurs d'avoir une conversation libre avec l'agent.

Python

# Important: update the agent with the new endpoint name so it can be used in future labeling sessions.
MY_AGENT_ENDPOINT_NAME = deployment.endpoint_name
AGENT_NAME = "My Agent v1"

my_review_app = my_review_app.add_agent(
    # Display name for the agent.
    agent_name=AGENT_NAME,
    model_serving_endpoint=MY_AGENT_ENDPOINT_NAME,
    overwrite=True
)

my_session = my_review_app.create_labeling_session(
    name="collect_facts_from_live_agent",
    assigned_users=ASSIGNED_USERS,
    agent=AGENT_NAME,
    # Built-in labeling schemas: EXPECTED_FACTS, GUIDELINES, EXPECTED_RESPONSE
    label_schemas=[review_app.label_schemas.EXPECTED_FACTS,review_app.label_schemas.GUIDELINES, "good_response"],
)

# Add the dataset to enable live agent interaction.
my_session.add_dataset(UC_TABLE_NAME)

# Share with the SME.
print("Review App URL:", my_review_app.url)
print("Labeling session URL: ", my_session.url)

Étapes suivantes

Une fois votre agent déployé, vous pouvez :

• Discutez avec lui dans AI Playground.

• Dans l'application d'examen, essayez ce qui suit :

  • Recueillir des commentaires généraux en utilisant « Discuter avec le bot ».
  • Collectez les labels des PME lors d'une session d'étiquetage.

• Utilisez-le dans votre application de production.

Exemple de Notebook

Agent Evaluation (MLflow 2) : Métriques personnalisées, lignes directrices et étiquettes d'experts du domaine

Ouvrir le notebook dans un nouvel ongletOuvrir dans Databricks