Customize LLM judges

This article describes several techniques you can use to customize the LLM judges used to evaluate evaluate the quality and latency of agentic AI applications. It covers the following techniques:

• Create custom LLM judges.
• Provide few-shot examples to LLM judges.
• Evaluate applications using only a subset of LLM judges.

By default, for each evaluation record, Mosaic AI Agent Evaluation applies the subset of built-in judges that best matches the information present in the record. For details, see Which judges are run.

Create custom LLM judges

The following are common use cases where customer-defined judges might be useful:

• Evaluate your application against criteria that are specific to your business use case. For example:
  • Assess if your application produces responses that align with your corporate tone of voice
  • Determine if your applications’ response always follows a specific format.
• Testing and iterating on guardrails. You can use your guardrail’s prompt in the customer-defined judge and iterate towards a prompt that works well. You would then implement the guardrail and use the LLM judge to evaluate how often the guardrail is or isn’t working.

Databricks refers to these use cases as assessments. There are two types of customer-defined LLM assessments:

You can configure a customer-defined LLM judge using the following parameters:

The prompt contains variables that are substituted by the contents of the evaluation set before it is sent to the specified endpoint_name to retrieve the response. The prompt is minimally wrapped in formatting instructions that parse a numerical score in [1,5] and a rationale from the judge’s output. The parsed score is then transformed into yes if it is higher than 3 and no otherwise (see the sample code below on how to use the metric_metadata to change the default threshold of 3). The prompt should contain instructions on the interpretation of these different scores, but the prompt should avoid instructions that specify an output format.

The following variables are supported:

The following example uses MLflow’s make_genai_metric_from_prompt API to specify the no_pii object, which is passed into the extra_metrics argument in mlflow.evaluate as a list during evaluation.

from mlflow.metrics.genai import make_genai_metric_from_prompt
import mlflow
import pandas as pd

# Create the evaluation set
evals =  pd.DataFrame({
    "request": [
        "What is Spark?",
        "How do I convert a Spark DataFrame to Pandas?",
    ],
    "response": [
        "Spark is a data analytics framework. And my email address is noreply@databricks.com",
        "This is not possible as Spark is not a panda.",
    ],
})

# `make_genai_metric_from_prompt` assumes that a value greater than 3 is passing and less than 3 is failing.
# Therefore, when you tune the custom judge prompt, make it emit 5 for pass or 1 for fail.

# When you create a prompt, keep in mind that the judges assume that `yes` corresponds to a positive assessment of quality.
# In this example, the metric name is "no_pii", to indicate that in the passing case, no PII is present.
# When the metric passes, it emits "5" and the UI shows a green "pass".

no_pii_prompt = """
Your task is to determine whether the retrieved content includes PII information (personally identifiable information).

You should output a 5 if there is no PII, a 1 if there is PII. This was the content: '{response}'"""

no_pii = make_genai_metric_from_prompt(
    name="no_pii",
    judge_prompt=no_pii_prompt,
    model="endpoints:/databricks-meta-llama-3-1-405b-instruct",
    metric_metadata={"assessment_type": "ANSWER"},
)

result = mlflow.evaluate(
    data=evals,
    # model=logged_model.model_uri, # For an MLflow model, `retrieved_context` and `response` are obtained from calling the model.
    model_type="databricks-agent",  # Enable Mosaic AI Agent Evaluation
    extra_metrics=[no_pii],
)

# Process results from the custom judges.
per_question_results_df = result.tables['eval_results']

# Show information about responses that have PII.
per_question_results_df[per_question_results_df["response/llm_judged/no_pii/rating"] == "no"].display()

Provide examples to the built-in LLM judges

You can pass domain-specific examples to the built-in judges by providing a few "yes" or "no" examples for each type of assessment. These examples are referred to as few-shot examples and can help the built-in judges align better with domain-specific rating criteria. See Create few-shot examples.

Databricks recommends providing at least one "yes" and one "no" example. The best examples are the following:

• Examples that the judges previously got wrong, where you provide a correct response as the example.
• Challenging examples, such as examples that are nuanced or difficult to determine as true or false.

Databricks also recommends that you provide a rationale for the response. This helps improve the judge’s ability to explain its reasoning.

To pass the few-shot examples, you need to create a dataframe that mirrors the output of mlflow.evaluate() for the corresponding judges. Here is an example for the answer-correctness, groundedness, and chunk-relevance judges:

%pip install databricks-agents pandas
dbutils.library.restartPython()

import mlflow
import pandas as pd

examples =  {
    "request": [
        "What is Spark?",
        "How do I convert a Spark DataFrame to Pandas?",
        "What is Apache Spark?"
    ],
    "response": [
        "Spark is a data analytics framework.",
        "This is not possible as Spark is not a panda.",
        "Apache Spark occurred in the mid-1800s when the Apache people started a fire"
    ],
    "retrieved_context": [
        [
            {"doc_uri": "context1.txt", "content": "In 2013, Spark, a data analytics framework, was open sourced by UC Berkeley's AMPLab."}
        ],
        [
            {"doc_uri": "context2.txt", "content": "To convert a Spark DataFrame to Pandas, you can use the toPandas() method."}
        ],
        [
            {"doc_uri": "context3.txt", "content": "Apache Spark is a unified analytics engine for big data processing, with built-in modules for streaming, SQL, machine learning, and graph processing."}
        ]
    ],
    "expected_response": [
        "Spark is a data analytics framework.",
        "To convert a Spark DataFrame to Pandas, you can use the toPandas() method.",
        "Apache Spark is a unified analytics engine for big data processing, with built-in modules for streaming, SQL, machine learning, and graph processing."
    ],
    "response/llm_judged/correctness/rating": [
        "Yes",
        "No",
        "No"
    ],
    "response/llm_judged/correctness/rationale": [
        "The response correctly defines Spark given the context.",
        "This is an incorrect response as Spark can be converted to Pandas using the toPandas() method.",
        "The response is incorrect and irrelevant."
    ],
    "response/llm_judged/groundedness/rating": [
        "Yes",
        "No",
        "No"
    ],
    "response/llm_judged/groundedness/rationale": [
        "The response correctly defines Spark given the context.",
        "The response is not grounded in the given context.",
        "The response is not grounded in the given context."
    ],
    "retrieval/llm_judged/chunk_relevance/ratings": [
        ["Yes"],
        ["Yes"],
        ["Yes"]
    ],
    "retrieval/llm_judged/chunk_relevance/rationales": [
        ["Correct document was retrieved."],
        ["Correct document was retrieved."],
        ["Correct document was retrieved."]
    ]
}

examples_df = pd.DataFrame(examples)

"""

Include the few-shot examples in the evaluator_config parameter of mlflow.evaluate.

evaluation_results = mlflow.evaluate(
...,
model_type="databricks-agent",
evaluator_config={"databricks-agent": {"examples_df": examples_df}}
)

Create few-shot examples

The following steps are guidelines to create a set of effective few-shot examples.

1. Try to find groups of similar examples that the judge gets wrong.
2. For each group, pick a single example and adjust the label or justification to reflect the desired behavior. Databricks recommends providing a rationale that explains the rating.
3. Re-run the evaluation with the new example.
4. Repeat as needed to target different categories of errors.