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.