Using Genie in Multi-Agent Systems — SQL Query Missing from Output

[Karthik_Karanm]

Hi Community,

I’ve implemented a multi-agent system using Genie in Databricks, following the guide on "Use Genie in multi-agent systems | Databricks Documentation". The system integrates well with vector indexes, and everything works fine during model logging and prediction phases.

I’m receiving all the expected results from the agents, but I’ve noticed that the SQL query is missing from the output. My expectation was to retrieve the generated SQL query from the Genie agent, but it doesn't appear in the response.

I have seen SQL queries in my poll results on MLFlow traces. How do I get this on the output.

Thanks in advance for your help!

#Databricks #GenieAgent #MultiAgent #SQLGeneration #VectorSearch #LLM #DatabricksGenAI #AIIntegration #GenAI #DatabricksCommunity

[Karthik_Karanm]

Hi Community,

I'm working with the Databricks Genie multi-agent framework and using the GenieAgent to generate SQL queries from natural language questions.

I've instantiated the agent like this:

from databricks.genie.agents import GenieAgent
from databricks.sdk import WorkspaceClient

genie_agent = GenieAgent(
    genie_space_id=GENIE_SPACE_ID,
    genie_agent_name="Genie",
    description=genie_agent_description,
    client=WorkspaceClient(
        host=os.getenv("DB_MODEL_SERVING_HOST_URL"),
        token=os.getenv("DATABRICKS_GENIE_PAT"),
    ),
)

The agent returns a result (i.e., the answer to the user's query), but it does not return the corresponding SQL query that was executed or generated. For my use case, I need to access the generated SQL as well, for transparency and debugging purposes.

Has anyone else encountered this?
Is there a specific way to extract the SQL query from the GenieAgent output?
Any configuration I might be missing?

Thanks in advance!

[Karthik_Karanm]

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!

[Karthik_Karanm]

Hello 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.

when I use this custom Genie agent without multi-agent framework,k I am able to serve the model successfully, but the problem is when I try to compile both structured and unstructured using multi-agent framework, it doesn't work.

Thanks,
Karthik.

[SP_6721]

Hi @Karthik_Karanm,

As per my understanding, the issue seems to be that the conversation state or authentication context isn’t being properly passed between agents in the multi-agent setup.

A few things you might try:

• Avoid caching session or conversation state, instead, create a fresh Genie client and start a new conversation for each request.
• Ensure any user tokens or auth context are passed through each agent to maintain proper access.
• Double-check that your custom agent’s input/output format aligns with what the multi-agent framework expects, even small mismatches can cause issues.

[mark_ott]

When using a custom Genie agent with the Databricks Genie REST API (GetMessage endpoint), extracting SQL queries and results works well as a standalone agent. However, integrating this custom agent into a multi-agent framework (especially when combining structured and unstructured sources) often leads to compilation, serving, or coordination issues that prevent the desired behavior.​

Common Causes & Challenges

• Databricks Runtime Dependencies: When serving a multi-agent model with MLflow, errors like ModuleNotFoundError: No module named 'dbruntime' are common. This often happens because custom agents or Genie code reference Databricks-specific modules not available outside the cluster or workspace runtime.​

• Authentication and State Passing: Multi-agent setups sometimes fail to pass conversation state or authentication context between agents correctly, resulting in missing SQL or improper task routing. Proper "on-behalf-of-user" authentication and context sharing are essential in multi-agent orchestration.​

• Agent Coordination and Permissions: Multi-agent supervisor patterns require carefully coordinated agent endpoints and explicit permissions for each subagent (Genie space, endpoint, function, or server). Missing permissions or incomplete setup can prevent subagents (including Genie) from contributing their results.​

Solutions & Best Practices

• Decouple Model Logic: Refactor agents or model code to avoid direct dependencies on Databricks runtime modules (like dbruntime). Try to ensure that any custom logic only uses libraries available in standard serving environments such as MLflow.​

• Ensure Proper Authentication: Use "on-behalf-of-user" authentication so that each agent can access Genie spaces and other resources in the user’s context. Review setup steps for multi-agent systems to make sure context and credentials are passed correctly between agents.​

• Agent Supervisor Setup: Use tools like Agent Bricks: Multi-Agent Supervisor, which orchestrate Genie and other agents while providing configuration for permissions and integration points. Grant explicit access for all agents and ensure correct endpoint configuration.​

• Structured & Unstructured Workflow Design: Leverage published notebooks and frameworks (like Mosaic AI Agent Framework or LangGraph) for orchestrating multi-agent workflows. These examples demonstrate composition for combining Genie agents with unstructured RAG or custom agents and walk through typical pitfalls.​

Next Steps

• Audit the multi-agent configuration to ensure all subagents have independent permissions and endpoints.

• If using MLflow, check the model’s pickle dependencies to confirm they are portable and not tied to Databricks-specific runtime.

• Review example notebooks provided in Databricks documentation for step-by-step multi-agent orchestration and troubleshooting.​