Complete your Databricks User Groups profile!

Fill out a few details about yourself so the community can get to know you.
Genie Community

Databricks Genie with ML: serving your custom models through Genie spaces with ai_query()

Summary: m5u95p shares a lean architectural pattern for serving custom ML models within the Genie Space, aimed at simplifying predictive queries without needing additional agents or apps. The method involves wrapping a model-serving endpoint in a Unity Catalog SQL function that interfaces with ai_query(). This approach is praised by mgjpmp and mpv6th for its simplicity and effectiveness. m8cx48 appreciates the shared method but suggests exploring recent advances such as the Relational Transformer for potentially reducing the need for custom solutions in handling tabular data. The discussion also includes references to a full write-up and a GitHub repository for deeper insights.
AI Summary

Genie is great at answering questions over your data — but what about predictive questions like "Is this week's spending anomalous?" Those need a custom ML model, and historically that meant building a separate agent or app on top of Genie.

I wanted to share a lean pattern that keeps everything inside the Genie Space. 👇

A lean architecture: Genie + ai_query()

Wrap your custom model-serving endpoint in a Unity Catalog SQL function that calls ai_query(). Add that function to your Genie Space, and Genie figures out when to invoke it from the function name and parameter comments — no router, no extra orchestration agent.

How?

  1. Train models (you can parallelize with applyInPandas).

  2. Serve the custom MLflow model on a serverless Model Serving endpoint.

  3. Wrap the endpoint in a UC SQL function around ai_query(), with rich comments on each parameter.

  4. Add the function to your Genie Space — now users can ask predictive questions in plain language.

The metadata you put on the function is what makes this work: clear names and parameter descriptions are how Genie decides to reach for the model.

Note : Need to route across multiple tools (several Genie spaces, endpoints, RAG)? That's where Agent Bricks Supervisor comes in — but for a single predictive capability, this approach is much simpler.

Try it

Would love to hear if others are exposing custom models through Genie this way — what use cases are you tackling? 🧱

4 comments

This is a really clean pattern. I like that it keeps everything inside the Genie Space and avoids overengineering for a single predictive use case.

agreed!

Thanks for sharing Gabriele! Very valuable and great write-up

Thanks for sharing!

Did you check the recent advances in Relational Transformer? Here is a link to one of the paper : https://arxiv.org/abs/2510.06377

and the associated github:

https://github.com/snap-stanford/relational-transformer/blob/main/README.md

I hope soon we will see more of these models developed and deployed so that we won’t need to build “custom” approach for some straight forward tasks on tabular data.