Mosaic AI Agent Framework¹

This is the recommended method for deploying models (agents.deploy imported from databricks-agents)
It replaces the traditional method of serving ML models via serving endpoints (client.create_endpoint from mlflow.deployments)

Advantages	Disadvantages
Agent evaluation resources are automatically set up and native to Databricks	Requires strict compliance with a certain structure of requests and responses² (Input = `List[ChatAgentMessage]`, Context = `Optional[ChatContext]`, Custom Inputs = `Optional[dict[str, Any]]`, Output = `ChatAgentResponse`)
Enables multi-agent workflows with tool-calling	Incompatible with OpenAI `Response` schema, which might require one extra translation step if we allow switching from Claude back to OpenAI responses or if we move towards Azure³. Mostly, this would mean replacing the `custom_input`/`custom_output` keys with a `metadata` key instead, but file attachments would also be handled differently.
Can control custom outputs such as citations separately without them being included in message history	Need to write frontend class (either Web UI or Bot Framework) to comply with Databricks response specification and maintain through changes being made
Can return multiple messages and intermediate steps for multi-step workflows
Automatically set up review endpoints that can be shared to users outside of Databricks

ChatAgent

MLflow offers an experimental class called ChatAgent that replaces the original ChatModel
Imported from mlflow.pyfunc
Added in MLflow 2.20.2
ChatModel itself replaces the original PythonModel, which was highly customizable, but required a lot of manual configuration
By subclassing ChatAgent, you are making a guarantee that your subclass has a predict method implemented (and optionally a predict_stream method) that takes ChatAgentRequest's and outputs ChatAgentResponses's

predict

The structure of the predict is as follows:

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

Once implemented as part of your subclass, the whole class is logged in MLflow as follows:

with mlflow.start_run():
    logged_agent_info = mlflow.pyfunc.log_model(
        artifact_path="agent",
        python_model=os.path.join(os.getcwd(), "agent"),
        # Add serving endpoints, tools, and vector search indexes here
        resources=[],
    )

This automatically registers the model signature as using ChatAgentRequest for inputs and ChatAgentResponse for outputs
- You can still add an optional input_example param as long as it complies with the ChatAgentRequest schema
  - Mostly useful for Databricks serving since it will automatically fill in the example page on the model serving endpoint page
Testing the predict method after registration is available from mlflow:

input_example = {
                  "messages":
                  [
                    {"content": "Hello!", "role": "user"}
                  ]
                }

mlflow.models.predict(
   model_uri = logged_chain_info.model_uri,
   input_data = input_example,
)

Databricks warns against keeping state as part of the subclass as it may be used to handle multiple requests during serving
Instead, initialize state locally in the predict method, use it for the call, and forget about it after

ChatAgentRequest⁴

The structure of ChatAgentRequest is as follows:

class ChatAgentRequest:
 messages: List[ChatAgentMessage]      # The conversation messages
 context: ChatContext = None           # Tracking IDs for conversation
 custom_inputs: Dict[str, Any] = None  # Dictionary of arbitrary other inputs, must be JSON-serializable
 stream: bool = False                  # Whether to stream back responses

For most methods related to ChatAgent, you can also pass in a dictionary that complies with this spec

{
 "messages": [],
 "context": {
  "conversation_id": "abc12345",
  "user_id": 4510
 },
 "custom_inputs": {
  "hello": "world!"
 },
 "stream": false
}

ChatAgentMessage⁵

The structure of ChatAgentMessage is as follows:

class ChatAgentMessage:
 role: str                             # Entity sending message ("user", "system", "assistant", "tool")
 content: str                          # Contents of message, can be None if tool_calls is set
 name: str = None                      # Entity name
 id: str                               # Generated ID value, required as part of ChatAgentResponse/ChatAgentChunk
 tool_class: List[ToolCall] = None     # List of tool calls made
 tool_call_id: str = None              # ID of tool call
 attachments: Dict[str, str] = None    # Optional attachments

 def check_content_and_tool_calls(values): ...
 # classmethod that checks if either content or tool_calls has been set

 def check_tool_messages(values): ...
 # classmethod that checks if name and tool_call_id are set for tool messages

As with ChatAgentRequest, you can specify most of these values as part of a dictionary
You will lose out on the safety of checking content/tool calls are specified appropriately

ChatContext ⁶

The structure of ChatContext is as follows:

class ChatContext:
 conversation_id: str = None
 user_id: str = None

Can be specified as a dictionary as well
Entirely optional
Could be used to track message or conversation history for individual users

ChatAgentResponse⁷

The structure of ChatAgentResponse is as follows:

class ChatAgentResponse:
 messages: List[ChatAgentMessage]       # Chat responses returned by model
 finish_reason: str = None              # Reason for generation of messages stopping
 custom_outputs: Dict[str, Any] = None  # Arbitrary additional content, must be JSON-serializable
 usage: ChatUsageHistory = None         # Request's token usage

Can be specified as a dictionary as well
This is different from the output from Databricks when served as part of a model endpoint (see below for more information)

Mosaic AI Agent Evaluation⁸

Agent Evaluation doesn't currently support rendering traces for additional input fields (custom_input)
Use mlflow.evaluate as shown:

import mlflow
import pandas as pd

examples =  {
    "request": [
        "What is Spark?",
        "How do I convert a Spark DataFrame to Pandas?",
    ],
    "response": [
        "Spark is a data analytics framework.",
        "This is not possible as Spark is not a panda.",
    ],
    "retrieved_context": [ # Optional, needed for judging groundedness.
        [{"doc_uri": "doc1.txt", "content": "In 2013, Spark, a data analytics framework, was open sourced by UC Berkeley's AMPLab."}],
        [{"doc_uri": "doc2.txt", "content": "To convert a Spark DataFrame to Pandas, you can use toPandas()"}],
    ],
    "expected_response": [ # Optional, needed for judging correctness.
        "Spark is a data analytics framework.",
        "To convert a Spark DataFrame to Pandas, you can use the toPandas() method.",
    ]
}

result = mlflow.evaluate(
    data=pd.DataFrame(examples),    # Your evaluation set
    # model=logged_model.model_uri, # If you have an MLFlow model. `retrieved_context` and `response` will be obtained from calling the model.
    model_type="databricks-agent",  # Enable Mosaic AI Agent Evaluation
)

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

Evaluation Resources⁹

Databricks creates the following resources automatically upon deploying a model with agents.deploy:

A serving endpoint that handles requests to the model
AI Gateway¹⁰
This is a set of inference tables that log inputs and responses
Can be used to monitor data/model quality and debug serving
Has a predefined-schema and naming convention
Can log MLflow traces if the environment variable ENABLE_MLFLOW_TRACING is True
Service principals with credentials to access necessary resources specified in deployment
Review App with a static link that stakeholders can use to provide feedback¹¹

Using the traditional method of serving models, only the endpoint would be created.

Databricks Serving¹²¹³

First, a list of types used in this section:

int refers to a numeric value between 0 and the maximum integer value
double refers to a decimal value, which in this case would be between 0 and 1
long refers to a numeric value between 0 and the maximum long value
str refers to a string containing arbitrarily structured UTF-8 formatted text
uuid_str is a str object that contains a UUID code that is used on the Databricks side for tracking
bool: A boolean with either true or false
"json_str" is a JSON object dumped to str format
For example, "{\"role\": \"assistant\", \"content\": \"Hello world!\"}"

Requests

The format of a request to query a model endpoint on Databricks is a JSON object that follows this specification

{
  "messages": [                 // Comes from MLflow ChatAgentRequest
    {
      "role": "user",
      "content": "str"
    }
  ],
  "custom_inputs": {
   "key": "value"
  },
  "context": {
   "conversation_id": str,
   "user_id": str
  },
  "stream": bool,
  "max_tokens": long,
  "temperature": double
}

Fields:
- "messages": A list of ChatAgentMessage-formatted objects represented as JSON
  - "role": The role of the individual making the request. Should be set to "user" for queries submitted
  - "content": The body of the request containing the user's question
- "custom_inputs": An optional field formatted as JSON-serializable objects with keys to allow arbitrary inputs
- "context": An optional field with ChatContext-formatted JSON objects
  - "conversation_id: ID for tracking this conversation
  - "user_id": ID for the user making the query
- "stream": Whether to stream chunks of text back in the response
- "max_tokens": The maximum number of tokens expected back in the response
- "temperature": A measure of how creative the response should be in terms of next-word prediction

Responses

The format of response from querying a model endpoint on Databricks is a JSON object that follows the ChatAgentResponse specification:

{
  "messages": [
    {
      "role": "assistant",
      "content": "str",
      "id": "str"
    }
  ],
  "databricks_output": {                  // Unique to Databricks and not ChatAgent
    "databricks_request_id": "uuid_str"
  },
  "id": "uuid_str",
  "custom_outputs": {
    "key": "value"
  }
}

Fields:
- "messages": A list of JSON objects that have ChatAgentMessage elements. For our purposes, there will only be one element in the list with the following fields:
  - "role": The role of the message generator, either human or AI. Because this is being returned by Databricks, it will always be "assistant"
  - "content": The generated response
  - "id": A unique identifier for tracing the response that is currently unused
- "databricks_output": Output specific to Databricks that is not part of the ChatAgent model. This field is ignored by MLflow.
  - "databricks_request_id": Used for tracing requests through the inference tables that Databricks sets up.
- "id": Used for tracing responses through the inference tables that Databricks sets up
- "custom_outputs": Arbitrary outputs are allowed as long as they are associated with a key and JSON-serializable
  - For example, a citations field may appear as follows:

...
"custom_outputs": {
 "citations": [
   {
  "document_uri": "str",
  "chunk_text": "str",
  "metadata": {
    "content_id": int,
    "file_id": int,
    "metadata": "json_str",
    "chunk_metadata": "json_str"
   },
 ],
}
...

Fields
- "citations": A list of JSON objects with the following elements:
  - "document_uri": Direct link to document containing relevant text
  - "chunk_text": The full relevant text body
  - "metadata": Other information related to the document in question
    - "content_id": Row ID in the vector search index
    - "file_id": ID number for file from which text is sourced
    - "metadata": Any other information that does not fit in the other categories. Format may vary over time
    - "chunk_metadata": Information related to the relevant chunk of text and how it was created

Differences

The Databricks model serving endpoint response format is a hybrid of MLflow's ChatAgentResponse and Open AI's ChatCompletions
- The custom_outputs field is new to the platform and comes from the MLflow ChatAgentResponse structure
- The choices field with a single message is from Open AI's ChatCompletions
- By including both, the Databricks response format complies with neither
In addition, MLflow's ChatAgentResponse is supposed to allow multi-message responses back, but currently, only one message comes back in the choices field
The documentation page is currently incorrect (as of 5/20/2025) and does not show all aspects of the model response
- It leaves out the custom_response field

Recommendations

We should use agents.deploy to provision model endpoints, as it seems better than using client.create_endpoint
The main reason for doing so is the automated setup of inference tables to log user interactions
The cost is the need to rewrite our backend code to be able to handle these responses
- There is a serious concern that Databricks could evolve the API further, which would necessitate further rewrites
- The specification for our custom_output's (such as citations) would still be decided by us and not by Databricks
The other option would be to keep everything as plain text responses, create our own logging tables, and associated serving code
- This would still require us to standardize the custom output we have coming out of the model to ensure the backend code can handle it

Mosaic AI Agent Framework1

ChatAgent

predict

ChatAgentRequest4

ChatAgentMessage5

ChatContext 6

ChatAgentResponse7

Mosaic AI Agent Evaluation8

Evaluation Resources9

Databricks Serving1213