Enforcing Minimum Cohort Size in a Databricks App with AI/BI Genie

[Fortinbras]

Hello, I'm currently building a Databricks App using AppKit and AI/BI Genie to let business users ask natural-language questions over registration and customer-demographic data. For privacy protection, demographic records are currently filtered so that only predefined demographic cohorts containing at least five distinct customers are available. However, Genie can join these records with registration data and apply additional filters or groupings, such as product, date, channel, or location. A cohort that was safe at the source level can therefore become a final result containing fewer than five distinct customers.

We need to enforce a rule equivalent to:

HAVING COUNT(DISTINCT customer_id) >= 5

at the final result level, after all joins, filters, and output grouping have been applied.

We are aiming for a robust privacy control rather than relying on prompt instructions alone, therefore we have considered a few different approaches, like client-side filtering or adding a server-side policy-aware proxy between the Genie plugin and the browser, but would appreciate guidance on the best Databricks-native pattern for this scenario:

• Is there a recommended way to centrally intercept/govern Genie query-result streams in an AppKit application?
  
• Are there established patterns for enforcing minimum cohort thresholds after arbitrary Genie-generated joins and filters?
• Would you recommend a server-side proxy, a custom agent/tool architecture, trusted SQL functions, or another approach?

[iyashk-DB · Accepted solution]

There's no built-in Genie feature that enforces a post-join HAVING threshold centrally, so you're right to think about this architecturally. The most reliable pattern in AppKit is to own the query execution layer yourself rather than letting Genie serve results directly.

Here's the approach: instead of embedding the Genie UI component directly, call the Genie REST API from your Databricks App's Python backend. When Genie responds to a user's question, the message contains a query attachment with the generated SQL text. Before showing anything to the user, you intercept that SQL, wrap it in a CTE, and re-execute it with your HAVING constraint enforced:

from databricks.sdk import WorkspaceClient

w = WorkspaceClient()

# Send the user's question to Genie
resp = w.genie.create_message_and_wait(
    space_id=SPACE_ID,
    conversation_id=conv.conversation_id,
    content=user_question,
)

# Pull the generated SQL from the first query attachment
generated_sql = None
for attachment in (resp.attachments or []):
    if attachment.query and attachment.query.query:
        generated_sql = attachment.query.query
        break

# Wrap with your privacy constraint before execution
if generated_sql:
    safe_sql = f"""
WITH _genie AS (
    {generated_sql}
)
SELECT *
FROM _genie
WHERE (
    SELECT COUNT(DISTINCT customer_id)
    FROM _genie
) >= 5
"""
    # Execute safe_sql against your warehouse

The exact wrapping depends on your schema. If Genie's output is already grouped (e.g. counts by product/region), you'd want to filter rows where the underlying customer count in each group meets the threshold. If it returns individual rows, a scalar subquery like above works for the whole result. You may need to adapt based on what cohort dimension Genie is grouping by.

As a second layer, use SQL Instructions in your Genie Agent configuration ("Always include HAVING COUNT(DISTINCT customer_id) >= 5 when aggregating or grouping demographic data") to nudge the LLM toward generating safe queries in the first place. This doesn't replace the code-level enforcement but reduces the cases where wrapping gets complicated.

The key difference between the two: SQL instructions are best-effort guidance to the model. The CTE wrapping at the app layer is your actual enforcement boundary.