custom-pyfunc-schema-agent(Python)

Loading...

Pyfunc custom schema agent notebook

This notebook shows you how to write a Python AI agent compatible with Mosaic AI Agent Framework that accepts custom inputs and returns custom outputs.

To ensure compatibility, the agent must conform to Mosaic AI Agent Framework schema requirements, see (AWS | Azure).

Model-as-code notebook

Mosaic AI Agent Framework uses MLflows Models-as-code development workflow, which requires two notebooks:

  • An agent notebook that defines the agent's logic (this notebook)
  • A driver notebook that logs, registers, and deploys the agent
    • You can find the driver notebook for this agent, custom-pyfunc-schema-driver, on Databricks documentation (AWS | Azure)

For more information on Model-as-code, see MLflow's Models as code guide.

Requirements

This notebook requires a Unity Catalog enabled workspace.

2

Define the agent's predict and predict_stream methods

Databricks recommends defining custom Python code agents as subclasses of MLflow's mlflow.pyfunc.ChatModel. Subclassing ChatModel allows you to easily adhere to the chat completion schema recommended by Mosaic AI Agent Framework, ensuring your agent can be used with downstream Databricks features like AI Playground.

The following cell defines an agent as a subclass of ChatModel. It implements the predict_stream method to support streaming responses and the predict method to process non-streaming requests.

4

5

Test the agent

Uncomment the following cell to test the agent.

7

Next steps

Rerun the cells above to iterate and test the agent.

See the driver notebook, custom-pyfunc-schema-driver (AWS | Azure), to learn how to log, register, and deploy this agent.