GenAI Agent with MLflow 3.0
Prerequisites: Run the following command to install MLflow 3.0 and OpenAI packages.
pip install --upgrade 'mlflow>=3.0.0rc0' --pre
pip install openai
This example demonstrates how to use MLflow to trace and evaluate OpenAI requests with prompts engineering. It showcases how to register prompts, generate traces, and assess response performance using evaluation datasets. The example also highlights the ability to track interactive traces and link them to the logged model for better observability.
Initialize OpenAI client​
First, we initialize an OpenAI client, use its chat completions API for answering questions, and utilize MLflow's prompt registry capability to manage and link prompts effectively.
As an initial step, set an active model with mlflow.set_active_model() API to enable linking any traces generated during app development to it.
from openai import OpenAI
import mlflow
mlflow.set_experiment("genai_example")
# set the active model for linking traces
mlflow.set_active_model(name="openai_model")
client = OpenAI()
question = "What is MLflow?"
# register a prompt so we can link it when logging the model
system_prompt = mlflow.register_prompt(
name="chatbot_prompt",
template="You are a chatbot that can answer questions about IT. Answer this question: {{question}}",
commit_message="Initial version of chatbot",
)
print(
client.completions.create(
prompt=system_prompt.format(question=question),
model="gpt-3.5-turbo-instruct",
temperature=0.1,
max_tokens=2000,
)
.choices[0]
.text
)
Switch to the Prompts tab to view the registered prompt and its details:
Test the model with tracing observability​
In this section, we manually test the model with example queries and leverage MLflow Tracing to analyze the outputs and debug potential issues. With autologging enabled, MLflow automatically creates a LoggedModel and links the generated traces to it, ensuring seamless observability.
# Enable autologging so that interactive traces from the client are automatically linked to a LoggedModel
mlflow.openai.autolog()
questions = [
"What is MLflow Tracking and how does it work?",
"What is Unity Catalog?",
"What are user-defined functions (UDFs)?",
]
outputs = []
for question in questions:
outputs.append(
client.completions.create(
prompt=system_prompt.format(question=question),
model="gpt-3.5-turbo-instruct",
temperature=0.1,
max_tokens=2000,
)
.choices[0]
.text
)
logged_model = mlflow.last_logged_model()
traces = mlflow.search_traces(model_id=logged_model.model_id)
print(traces)
# trace_id trace ... assessments request_id
# 0 9f1548c66b4647e6ae67336dd1b83fe5 Trace(trace_id=9f1548c66b4647e6ae67336dd1b83fe5) ... [] 9f1548c66b4647e6ae67336dd1b83fe5
# 1 c574c9825c6443b09b4c8cf0862f6c9a Trace(trace_id=c574c9825c6443b09b4c8cf0862f6c9a) ... [] c574c9825c6443b09b4c8cf0862f6c9a
# 2 5e29af629f144ebf852a55fa11daf7d7 Trace(trace_id=5e29af629f144ebf852a55fa11daf7d7) ... [] 5e29af629f144ebf852a55fa11daf7d7
#
# [3 rows x 12 columns]
Check out the Models tab in the experiment to view the new logged model with name openai_model.
On the Logged Model page, you can view detailed information, including the model_id.
Navigating to the Traces tab of the model's page, you can view the traces that were just generated.
Evaluate the agent's performance​
Finally, we use mlflow.evaluate() to assess the performance of the logged model on an evaluation dataset.
This step involves calculating additional metrics, such as latency and answer correctness, to gain deeper insights into the model's behavior and accuracy.
# Prepare the eval dataset in a pandas DataFrame
import pandas as pd
from mlflow.entities import LoggedModelInput
eval_df = pd.DataFrame(
{
"messages": questions,
"expected_response": [
"""MLflow Tracking is a key component of the MLflow platform designed to record and manage machine learning experiments. It enables data scientists and engineers to log parameters, code versions, metrics, and artifacts in a systematic way, facilitating experiment tracking and reproducibility.\n\nHow It Works:\n\nAt the heart of MLflow Tracking is the concept of a run, which is an execution of a machine learning code. Each run can log the following:\n\nParameters: Input variables or hyperparameters used in the model (e.g., learning rate, number of trees). Metrics: Quantitative measures to evaluate the model's performance (e.g., accuracy, loss). Artifacts: Output files like models, datasets, or images generated during the run. Source Code: The version of the code or Git commit hash used. These logs are stored in a tracking server, which can be set up locally or on a remote server. The tracking server uses a backend storage (like a database or file system) to keep a record of all runs and their associated data.\n\n Users interact with MLflow Tracking through its APIs available in multiple languages (Python, R, Java, etc.). By invoking these APIs in the code, you can start and end runs, and log data as the experiment progresses. Additionally, MLflow offers autologging capabilities for popular machine learning libraries, automatically capturing relevant parameters and metrics without manual code changes.\n\nThe logged data can be visualized using the MLflow UI, a web-based interface that displays all experiments and runs. This UI allows you to compare runs side-by-side, filter results, and analyze performance metrics over time. It aids in identifying the best models and understanding the impact of different parameters.\n\nBy providing a structured way to record experiments, MLflow Tracking enhances collaboration among team members, ensures transparency, and makes it easier to reproduce results. It integrates seamlessly with other MLflow components like Projects and Model Registry, offering a comprehensive solution for managing the machine learning lifecycle.""",
"""Unity Catalog is a feature in Databricks that allows you to create a centralized inventory of your data assets, such as tables, views, and functions, and share them across different teams and projects. It enables easy discovery, collaboration, and reuse of data assets within your organization.\n\nWith Unity Catalog, you can:\n\n1. Create a single source of truth for your data assets: Unity Catalog acts as a central repository of all your data assets, making it easier to find and access the data you need.\n2. Improve collaboration: By providing a shared inventory of data assets, Unity Catalog enables data scientists, engineers, and other stakeholders to collaborate more effectively.\n3. Foster reuse of data assets: Unity Catalog encourages the reuse of existing data assets, reducing the need to create new assets from scratch and improving overall efficiency.\n4. Enhance data governance: Unity Catalog provides a clear view of data assets, enabling better data governance and compliance.\n\nUnity Catalog is particularly useful in large organizations where data is scattered across different teams, projects, and environments. It helps create a unified view of data assets, making it easier to work with data across different teams and projects.""",
"""User-defined functions (UDFs) in the context of Databricks and Apache Spark are custom functions that you can create to perform specific tasks on your data. These functions are written in a programming language such as Python, Java, Scala, or SQL, and can be used to extend the built-in functionality of Spark.\n\nUDFs can be used to perform complex data transformations, data cleaning, or to apply custom business logic to your data. Once defined, UDFs can be invoked in SQL queries or in DataFrame transformations, allowing you to reuse your custom logic across multiple queries and applications.\n\nTo use UDFs in Databricks, you first need to define them in a supported programming language, and then register them with the SparkSession. Once registered, UDFs can be used in SQL queries or DataFrame transformations like any other built-in function.\n\nHere\'s an example of how to define and register a UDF in Python:\n\n```python\nfrom pyspark.sql.functions import udf\nfrom pyspark.sql.types import IntegerType\n\n# Define the UDF function\ndef multiply_by_two(value):\n return value * 2\n\n# Register the UDF with the SparkSession\nmultiply_udf = udf(multiply_by_two, IntegerType())\n\n# Use the UDF in a DataFrame transformation\ndata = spark.range(10)\nresult = data.withColumn("multiplied", multiply_udf(data.id))\nresult.show()\n```\n\nIn this example, we define a UDF called `multiply_by_two` that multiplies a given value by two. We then register this UDF with the SparkSession using the `udf` function, and use it in a DataFrame transformation to multiply the `id` column of a DataFrame by two.""",
],
"predictions": outputs,
}
)
# Start a run to represent the evaluation job
with mlflow.start_run() as evaluation_run:
eval_dataset = mlflow.data.from_pandas(
df=eval_df,
name="eval_dataset",
targets="expected_response",
predictions="predictions",
)
mlflow.log_input(dataset=eval_dataset)
# Run the evaluation based on extra metrics
# Current active model will be automatically used
result = mlflow.evaluate(
data=eval_dataset,
model_type="question-answering",
extra_metrics=[
mlflow.metrics.latency(),
mlflow.metrics.genai.answer_correctness("openai:/gpt-4o"),
],
# This is needed since answer_correctness looks for 'inputs' field
evaluator_config={"col_mapping": {"inputs": "messages"}},
)
print(result.tables["eval_results_table"])
# messages ... answer_correctness/v1/justification
# 0 What is MLflow Tracking and how does it work? ... The output is mostly correct and aligns well w...
# 1 What is Unity Catalog? ... The output is completely incorrect as it descr...
# 2 What are user-defined functions (UDFs)? ... The output accurately describes user-defined f...
#
# [3 rows x 7 columns]
Navigating to the active model, you can see the metrics and their details displayed in the MLflow UI. This includes metrics like latency and answer correctness, providing insights into the model's performance.
