Hi community,

I've implemented a custom agent using the Databricks Genie REST API (GetMessage endpoint) to mimic the core functionality of the Genie Agent library—specifically to extract SQL queries from the assistant's responses. This setup works as expected, and I'm able to retrieve both the generated SQL and relevant query results.

I’ve since integrated this agent into a multi-agent framework. However, when attempting to serve the model using MLflow, I encountered the following error during model loading:

[kjhd7] [2025-06-03 09:47:59 +0000]     python_model = cloudpickle.load(f)
[kjhd7] [2025-06-03 09:47:59 +0000]                    ^^^^^^^^^^^^^^^^^^^
[kjhd7] [2025-06-03 09:47:59 +0000] ModuleNotFoundError: No module named 'dbruntime'
[kjhd7] [2025-06-03 09:47:59 +0000] [12] [INFO] Worker exiting (pid: 12)
[kjhd7] [2025-06-03 09:48:03 +0000] An error occurred while loading the model: No module named 'dbruntime'

It seems like the model is trying to reference a Databricks-specific runtime module (dbruntime) that isn't available in the serving environment.

Has anyone else encountered this issue when serving models that reference internal Databricks components? Any recommendations for workarounds or best practices for decoupling model logic from Databricks-specific modules during serving?

Thanks in advance for any help!