Log and register AI agents
Important
This feature is in Public Preview.
Log AI agents using Mosaic AI Agent Framework. Logging an agent is the basis of the development process. Logging captures a “point in time” of the agent’s code and configuration so you can evaluate the quality of the configuration.
Requirements
Create an AI agent before logging it.
Code-based logging
Databricks recommends using MLflow’s Models from Code functionality when logging agents.
In this approach, the agent’s code is captured as a Python file, and the Python environment is captured as a list of packages. When the agent is deployed, the Python environment is restored, and the agent’s code is executed to load the agent into memory so it can be invoked when the endpoint is called.
You can couple this approach with the use of pre-deployment validation APIs like mlflow.models.predict() to ensure that the agent runs reliably when deployed for serving.
The code that logs the agent or agent must be in a separate notebook from the agent code. This notebook is called a driver notebook. For an example notebook, see Example notebooks.
Infer Model Signature during logging
During logging, you must define an MLflow Model Signature, which specifies the agent’s input and output schema. The signature validates inputs and outputs to ensure that the agent interacts correctly with downstream tools like AI Playground and the review app. It also guides other applications on how to use the agent effectively.
Databricks recommends using MLflow’s Model Signature inferencing capabilities to automatically generate the agent’s signature based on an input example you provide. This approach is more convenient than manually defining the signature.
The LangChain and PyFunc examples below use Model Signature inferencing.
If you would rather explicitly define a Model Signature yourself at logging time, see MLflow docs - How to log models with signatures.
Code-based logging with LangChain
The following instructions and code sample show you how to log an agent with LangChain.
Create a notebook or Python file with your code. For this example, the notebook or file is named
agent.py
. The notebook or file must contain a LangChain agent, referred to here aslc_agent
.Include mlflow.models.set_model(lc_agent) in the notebook or file.
Create a new notebook to serve as the driver notebook (called
driver.py
in this example).In the driver notebook, use the following code to run
agent.py
and log the results to an MLflow model:mlflow.langchain.log_model(lc_model="/path/to/agent.py", resources=list_of_databricks_resources)
The
resources
parameter declares Databricks-managed resources needed to serve the agent, such as a vector search index or serving endpoint that serves a foundation model. For more information, see Specify resources for automatic authentication passthrough.Deploy the model. See Deploy an agent for generative AI application.
When the serving environment is loaded,
agent.py
is executed.When a serving request comes in,
lc_agent.invoke(...)
is called.
import mlflow
code_path = "/Workspace/Users/first.last/agent.py"
config_path = "/Workspace/Users/first.last/config.yml"
# Input example used by MLflow to infer Model Signature
input_example = {
"messages": [
{
"role": "user",
"content": "What is Retrieval-augmented Generation?",
}
]
}
# example using langchain
with mlflow.start_run():
logged_agent_info = mlflow.langchain.log_model(
lc_model=code_path,
model_config=config_path, # If you specify this parameter, this configuration is used by agent code. The development_config is overwritten.
artifact_path="agent", # This string is used as the path inside the MLflow model where artifacts are stored
input_example=input_example, # Must be a valid input to the agent
example_no_conversion=True, # Required
)
print(f"MLflow Run: {logged_agent_info.run_id}")
print(f"Model URI: {logged_agent_info.model_uri}")
# To verify that the model has been logged correctly, load the agent and call `invoke`:
model = mlflow.langchain.load_model(logged_agent_info.model_uri)
model.invoke(example)
Code-based logging with PyFunc
The following instructions and code sample show you how to log an agent with PyFunc.
Create a notebook or Python file with your code. For this example, the notebook or file is named
agent.py
. The notebook or file must contain a PyFunc class, namedPyFuncClass
.Include
mlflow.models.set_model(PyFuncClass)
in the notebook or file.Create a new notebook to serve as the driver notebook (called
driver.py
in this example).In the driver notebook, use the following code to run
agent.py
and log the results to an MLflow model:mlflow.pyfunc.log_model(python_model="/path/to/agent.py", resources=list_of_databricks_resources)
The
resources
parameter declares Databricks-managed resources needed to serve the agent, such as a vector search index or serving endpoint that serves a foundation model. For more information, see Specify resources for automatic authentication passthrough.Deploy the model. See Deploy an agent for generative AI application.
When the serving environment is loaded,
agent.py
is executed.When a serving request comes in,
PyFuncClass.predict(...)
is called.
import mlflow
from mlflow.models.resources import (
DatabricksServingEndpoint,
DatabricksVectorSearchIndex,
)
code_path = "/Workspace/Users/first.last/agent.py"
config_path = "/Workspace/Users/first.last/config.yml"
# Input example used by MLflow to infer Model Signature
input_example = {
"messages": [
{
"role": "user",
"content": "What is Retrieval-augmented Generation?",
}
]
}
with mlflow.start_run():
logged_agent_info = mlflow.pyfunc.log_model(
python_model=agent_notebook_path,
artifact_path="agent",
input_example=input_example,
resources=resources_path,
example_no_conversion=True,
resources=[
DatabricksServingEndpoint(endpoint_name="databricks-mixtral-8x7b-instruct"),
DatabricksVectorSearchIndex(index_name="prod.agents.databricks_docs_index"),
]
)
print(f"MLflow Run: {logged_agent_info.run_id}")
print(f"Model URI: {logged_agent_info.model_uri}")
# To verify that the model has been logged correctly, load the agent and call `invoke`:
model = mlflow.pyfunc.load_model(logged_agent_info.model_uri)
model.invoke(example)
Specify resources for automatic authentication passthrough
AI agents often need to authenticate to other resources to complete tasks. For example, an agent may need to access a Vector Search index to query unstructured data.
As described in Authentication for dependent resources, Model Serving supports authenticating to both Databricks-managed and external resources when you deploy the agent.
For the most common Databricks resource types, Databricks supports and recommends declaring resource dependencies for the agent upfront during logging. This enables automatic authentication passthrough when you deploy the agent - Databricks automatically provisions, rotates, and manages short-lived credentials to securely access these resource dependencies from within the agent endpoint.
To enable automatic authentication passthrough, specify dependent resources using the resources
parameter of the log_model() API, as shown in the following code.
import mlflow
from mlflow.models.resources import (
DatabricksVectorSearchIndex,
DatabricksServingEndpoint,
DatabricksSQLWarehouse,
DatabricksFunction,
DatabricksGenieSpace,
DatabricksTable,
)
with mlflow.start_run():
logged_agent_info = mlflow.pyfunc.log_model(
python_model=agent_notebook_path,
artifact_path="agent",
input_example=input_example,
example_no_conversion=True,
# Specify resources for automatic authentication passthrough
resources=[
DatabricksVectorSearchIndex(index_name="prod.agents.databricks_docs_index"),
DatabricksServingEndpoint(endpoint_name="databricks-mixtral-8x7b-instruct"),
DatabricksServingEndpoint(endpoint_name="databricks-bge-large-en"),
DatabricksSQLWarehouse(warehouse_id="your_warehouse_id"),
DatabricksFunction(function_name="ml.tools.python_exec"),
DatabricksGenieSpace(genie_space_id="your_genie_space_id"),
DatabricksTable(table_name="your_table_name"),
]
)
Databricks recommends you manually specify resources
for all agent flavors.
Note
If you do not specify resources when logging LangChain agents using mlflow.langchain.log_model(...)
, MLflow performs best-effort automatic inference of resources. However, this may not capture all dependencies, resulting in authorization errors when serving or querying the agent.
The following table lists the Databricks resources that support automatic authentication passthrough and the minimum mlflow
version required to log the resource.
Resource type | Minimum mlflow version required to log the resource |
---|---|
Vector Search index | Requires mlflow 2.13.1 or above |
Model serving endpoint | Requires mlflow 2.13.1 or above |
SQL warehouse | Requires mlflow 2.16.1 or above |
Unity Catalog function | Requires mlflow 2.16.1 or above |
Genie space | Requires mlflow 2.17.1 or above |
Unity Catalog table | Requires mlflow 2.18.0 or above |
Register the agent to Unity Catalog
Before you deploy the agent, you must register the agent to Unity Catalog. Registering the agent packages it as a model in Unity Catalog. As a result, you can use Unity Catalog permissions for authorization for resources in the agent.
import mlflow
mlflow.set_registry_uri("databricks-uc")
catalog_name = "test_catalog"
schema_name = "schema"
model_name = "agent_name"
model_name = catalog_name + "." + schema_name + "." + model_name
uc_model_info = mlflow.register_model(model_uri=logged_agent_info.model_uri, name=model_name)