Allow verified queries to be marked for onboarding

journeys/iteration.py

@@ -36,11 +36,6 @@
     yaml_to_semantic_model,
 )
 from semantic_model_generator.protos import semantic_model_pb2
-from semantic_model_generator.snowflake_utils.env_vars import (
-    SNOWFLAKE_ACCOUNT_LOCATOR,
-    SNOWFLAKE_HOST,
-    SNOWFLAKE_USER,
-)
 from semantic_model_generator.validate_model import validate
 
 
@@ -234,6 +229,11 @@ def edit_verified_query(
 
             elif st.session_state.get("successful_sql", False):
                 # Moved outside the `if run:` block to ensure it's always evaluated
+                mark_as_onboarding = st.checkbox(
+                    "Mark as onboarding question",
+                    key=f"edit_onboarding_idx_{message_index}",
+                    help="Mark this question as an onboarding verified query.",
+                )
                 save = st.button(
                     "Save as verified query",
                     use_container_width=True,
@@ -243,12 +243,18 @@ def edit_verified_query(
                     sql_no_analyst_comment = user_updated_sql.replace(
                         " /* Generated by Cortex Analyst */", ""
                     )
-                    add_verified_query(question, sql_no_analyst_comment)
+                    add_verified_query(
+                        question,
+                        sql_no_analyst_comment,
+                        is_onboarding_question=mark_as_onboarding,
+                    )
                     st.session_state["editing"] = False
                     st.session_state["confirmed_edits"] = True
 
 
-def add_verified_query(question: str, sql: str) -> None:
+def add_verified_query(
+    question: str, sql: str, is_onboarding_question: bool = False
+) -> None:
     """Save verified question and SQL into an in-memory list with additional details."""
     # Verified queries follow the Snowflake definitions.
     verified_query = semantic_model_pb2.VerifiedQuery(
@@ -257,6 +263,7 @@ def add_verified_query(question: str, sql: str) -> None:
         sql=sql,
         verified_by=st.session_state["user_name"],
         verified_at=int(time.time()),
+        use_as_onboarding_question=is_onboarding_question,
     )
     st.session_state.semantic_model.verified_queries.append(verified_query)
     st.success(
@@ -307,6 +314,11 @@ def display_content(
                 df = pd.read_sql(sql, conn)
                 st.dataframe(df, hide_index=True)
 
+                mark_as_onboarding = st.checkbox(
+                    "Mark as onboarding question",
+                    key=f"onboarding_idx_{message_index}",
+                    help="Mark this question as an onboarding verified query.",
+                )
                 left, right = st.columns(2)
                 if right.button(
                     "Save as verified query",
@@ -317,7 +329,9 @@ def display_content(
                     cleaned_sql = sql_no_cte.replace(
                         " /* Generated by Cortex Analyst */", ""
                     )
-                    add_verified_query(question, cleaned_sql)
+                    add_verified_query(
+                        question, cleaned_sql, is_onboarding_question=mark_as_onboarding
+                    )
 
                 if left.button(
                     "Edit",

semantic_model_generator/protos/semantic_model.proto

@@ -307,6 +307,9 @@ message VerifiedQuery {
   int64 verified_at = 5 [(optional) = true];
   // Name of the person who verified this query.
   string verified_by = 6 [(optional) = true];
+  // Whether to always include in this question in the suggested questions
+  // module
+  bool use_as_onboarding_question = 7 [(optional) = true];
 }
 
 // VerifiedQueryRepository is a simply a collection of verified queries.