databao 0.3.3.dev2__py3-none-any.whl → 0.3.3.dev3__py3-none-any.whl

{databao-0.3.3.dev2.dist-info → databao-0.3.3.dev3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: databao
-Version: 0.3.3.dev2
+Version: 0.3.3.dev3
 License-File: LICENSE
 Requires-Python: >=3.11
 Requires-Dist: adbc-driver-manager>=1.10.0

{databao-0.3.3.dev2.dist-info → databao-0.3.3.dev3.dist-info}/RECORD

@@ -1,8 +1,8 @@
 databao_cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-databao_cli/__main__.py,sha256=BFA_15eQIiQOvT26APsazBTmnynYk7_OlipcVHiQcpU,8428
+databao_cli/__main__.py,sha256=-ZySAVV7bABN_PwleKCZUx06gUEcnJvBN2kjZQYV2yE,9027
 databao_cli/executor_utils.py,sha256=WFUuUoxfHGsxp2wfjpgCeGs4v1EGKl8ZKltzKp4VFmI,1733
 databao_cli/commands/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-databao_cli/commands/app.py,sha256=u9odVAWyEJHFcaxWb-quCHIb8RRGEtkInU_2hW45yfs,645
+databao_cli/commands/app.py,sha256=mCUN8qJlF7tj23fRNHAAzlZp50GJgwMVUxDJwaR2vcA,813
 databao_cli/commands/ask.py,sha256=NGvjJwdYbBhqyqSFGknaNmgd-O2X-HrI3acH7pJ5psA,7143
 databao_cli/commands/build.py,sha256=alsyExwwI23QyG49h5d6QL6nf4hjGHtluciepLKY2LQ,506
 databao_cli/commands/context_engine_cli.py,sha256=U88IlVS_PQpcQiNp86lCxx1uAQLyGlfcgNXNqUsg1pY,1241
@@ -24,18 +24,18 @@ databao_cli/mcp/tools/databao_ask/tool.py,sha256=-1541H8KDY18jS_qYuB1AFu_n0pRHjC
 databao_cli/project/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 databao_cli/project/layout.py,sha256=WcIUaNBD-njKiGmFiTDKwcOk1NgRV5CQ6NjdNVTnu30,1586
 databao_cli/ui/__init__.py,sha256=5U7huKF5FMi7RcgidxcbFfCx2hKKAJj6JxD4CgvS7n0,39
-databao_cli/ui/app.py,sha256=Y4DQRKoijASHwrHgs_OvoA2_aoFuHaqztKtAtlgwvTE,17858
-databao_cli/ui/cli.py,sha256=JNuE9n8oAaYkTEMRHs588Vi5YtEVYVPqBEvWwbirOgk,1022
+databao_cli/ui/app.py,sha256=qmoklnHCFPSRA2HXWGL_375f2qQwEV-Ob8twQMnx-r8,19379
+databao_cli/ui/cli.py,sha256=nRzaGBHoJTzTSEQ_8yB_5G0EoGC-do-crA4DzV_RnXQ,1281
 databao_cli/ui/project_utils.py,sha256=4LG8OimVxuRa2Ke0rtn7AxE-7OuAK3qnvVNxLk11z3c,1614
 databao_cli/ui/streaming.py,sha256=aBaw8UC01lTL0M2AfvT-VRRVFamAbhP4BLrfDLC4f50,1376
-databao_cli/ui/suggestions.py,sha256=EuDUsG0gA2T1q-lqRas5lYOhXX4ZjZAlzkKc7EvbkEM,7217
+databao_cli/ui/suggestions.py,sha256=vXkmr8dmRL0xhIemCa6oifGgpTClJsVTf-PQ9hJQaqw,8486
 databao_cli/ui/assets/bao.png,sha256=aJOABLmxPmn7pfCAeiUbXgho_SilUcdpCtKDYm0_LLs,54635
 databao_cli/ui/components/__init__.py,sha256=bcxC59NGEPVfDfb4VwSoj1mHuWVJnWd-aNEmMvZMmlI,31
-databao_cli/ui/components/chat.py,sha256=jQ-Q9620UuygLtHzc8wF4hjOOCRjfGJGSm7ecMusr1k,13516
+databao_cli/ui/components/chat.py,sha256=ttxBhLoRH4ljm0xwQilns1c33yQTffVJNF0BTfhO6jA,19478
 databao_cli/ui/components/datasource_form.py,sha256=qaGfBakCnkHtitcMAj4rnhx7z6aYxhXI3msFGMKk65g,6306
-databao_cli/ui/components/datasource_manager.py,sha256=vP0Z9gISfKzdYKestYZHuGSAkMi-rMAuMiCPRpCT8cM,9495
+databao_cli/ui/components/datasource_manager.py,sha256=v_djWJZ7xetvkYuFLBRsAl7o3uY7qqPFgpIMAx7ARbc,9635
 databao_cli/ui/components/icons.py,sha256=IdFBEwTXK_gBaDIoTCdGaFAV2kGhP1pcM94-ISmp9rA,1755
-databao_cli/ui/components/results.py,sha256=5vkCenMQUDj4Nt9r9cGHhDppfIXZiRTvHiQKazmXqPs,12488
+databao_cli/ui/components/results.py,sha256=yXsiboayqH54GytO2dJB6-5piYFDugOM_JtqEXXy3wE,13246
 databao_cli/ui/components/sidebar.py,sha256=AuL02hDvmCO64yzstncUoW64byZdfy4vBZdRIjVVe4A,6027
 databao_cli/ui/components/status.py,sha256=W9RG250yef39ZTLfGW7vVfRYJ3DYPb0JjKkWgK-TPfo,3953
 databao_cli/ui/models/__init__.py,sha256=fwnl7W0XfhKPlCgHnQ33SEF4uldigo_Ndm5h6SDJDLA,279
@@ -44,20 +44,20 @@ databao_cli/ui/models/settings.py,sha256=xrmtHpmhsL8H1PpxWtkA6qJRegT-95u6phPm50Y
 databao_cli/ui/pages/__init__.py,sha256=_pRh8Qxs02DzS4r9z3-LPjUQTWte0NTa0Zu4la6dHRc,60
 databao_cli/ui/pages/agent_settings.py,sha256=huBZAmM9McA_e3W6iyJPfF23y5nFqYohLqwH7h_hlC4,12639
 databao_cli/ui/pages/chat.py,sha256=yLJt3Jm_62ar2jUFB1OUAQqG5AkaOa9zciMiQaStASk,7779
-databao_cli/ui/pages/context_settings.py,sha256=aT_Hx5XlJj5eEfvLDqnSLtNWt5q30QuijJj1AwdRx2g,5297
+databao_cli/ui/pages/context_settings.py,sha256=wI0_In_19tSLqEnXpe-vslNzl7eYLyhii_FmIt4jdlY,5085
 databao_cli/ui/pages/general_settings.py,sha256=wXtNpnvFK-BD_AYExbym0aFluDOg2J_aMWshFKfqZTg,3874
-databao_cli/ui/pages/welcome.py,sha256=VmDrxnTar-0y3ggVdC2lZYNGA2VlsYqi4tv4e7_3IWc,12512
-databao_cli/ui/services/__init__.py,sha256=K_5kJPty2sBEUpGkql-PQQouWZsl-Pywj54CLogLmXc,1334
+databao_cli/ui/pages/welcome.py,sha256=bXtdJ35LG48KNEdSJHaeJpkZ-uHUPeAA1aL7a4ofeDs,12924
+databao_cli/ui/services/__init__.py,sha256=lmPHHuOVdOA-c_POOMkD1044SRBb6-Jl1Gv-2g4Z63U,1368
 databao_cli/ui/services/build_service.py,sha256=St78meRjvT5BJzVhHusWTSL76NESuQGryxRD0CtnGwQ,8033
 databao_cli/ui/services/chat_persistence.py,sha256=kb7KTT2WCNFPcSz99m653cAMV8mN2hVbD1pMyAVaX40,9346
 databao_cli/ui/services/chat_title.py,sha256=cN5isa3NOp2_yHvLtTuO1ZuOII7D064KfiGxGQuxeT8,5630
 databao_cli/ui/services/dce_operations.py,sha256=mHGmLit270BC8br9wBxSGwCewdUjts8MOa40QS_rxSo,7199
 databao_cli/ui/services/llm_models.py,sha256=EQQ1ATSaDbs7bHVSwWpLHcYFbO1VA0nTl-t3HZJz14Q,2779
-databao_cli/ui/services/query_executor.py,sha256=9-l171v944_EhagiCPaxI-SOy1t-pYHmaeTu9-mV0Yo,7606
+databao_cli/ui/services/query_executor.py,sha256=TEcHFUp7yuo7GdTtGJXJmuF8IAkctjc2B3BHbSFLz7c,10879
 databao_cli/ui/services/settings_persistence.py,sha256=LFgO21guu0CfcKSRWWwUkc7nREwSe8iLnu3GSE8JhCg,1878
 databao_cli/ui/services/storage.py,sha256=wCSAjL8WiAghX2JaZVCDxb52NCMwL010ziqo6xcwqd8,2305
-databao-0.3.3.dev2.dist-info/METADATA,sha256=iDyAQHee0Iu8dHn5BEo3f3p9pjXLestuTELsHr1IGz0,3374
-databao-0.3.3.dev2.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
-databao-0.3.3.dev2.dist-info/entry_points.txt,sha256=fKKysivJDgZbz5wKQXsbPALm2WTpY1RPHqZmZ_hu1f4,53
-databao-0.3.3.dev2.dist-info/licenses/LICENSE,sha256=_u6l_NE7cSV2FtV2hwTT9hB-AKjkQfZHQgE1P1kFs24,12353
-databao-0.3.3.dev2.dist-info/RECORD,,
+databao-0.3.3.dev3.dist-info/METADATA,sha256=kaXR-3XFhEcbGpm_OdpVo2AvjAk4JnD4y_UmfmSKr_4,3374
+databao-0.3.3.dev3.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
+databao-0.3.3.dev3.dist-info/entry_points.txt,sha256=fKKysivJDgZbz5wKQXsbPALm2WTpY1RPHqZmZ_hu1f4,53
+databao-0.3.3.dev3.dist-info/licenses/LICENSE,sha256=_u6l_NE7cSV2FtV2hwTT9hB-AKjkQfZHQgE1P1kFs24,12353
+databao-0.3.3.dev3.dist-info/RECORD,,

databao_cli/__main__.py

@@ -215,8 +215,23 @@ def ask(
     default=False,
     help="Disable all domain-editing operations (init, datasources, build) in the UI",
 )
+@click.option(
+    "--hide-suggested-questions",
+    is_flag=True,
+    default=False,
+    help="Hide the suggested questions on the empty chat screen",
+)
+@click.option(
+    "--hide-build-context-hint",
+    is_flag=True,
+    default=False,
+    help=(
+        "Hide the 'Context isn't built yet' warning on the empty chat screen and "
+        "remove the Build Context step from the setup wizard"
+    ),
+)
 @click.pass_context
-def app(ctx: click.Context, read_only_domain: bool) -> None:
+def app(ctx: click.Context, read_only_domain: bool, hide_suggested_questions: bool, hide_build_context_hint: bool) -> None:
     """Launch the Databao Streamlit web interface.
 
     All additional arguments are passed directly to streamlit run.
@@ -231,6 +246,8 @@ def app(ctx: click.Context, read_only_domain: bool) -> None:
     from databao_cli.commands.app import app_impl
 
     ctx.obj["read_only_domain"] = read_only_domain
+    ctx.obj["hide_suggested_questions"] = hide_suggested_questions
+    ctx.obj["hide_build_context_hint"] = hide_build_context_hint
     app_impl(ctx)
 
 

databao_cli/commands/app.py

@@ -16,6 +16,8 @@ def app_impl(ctx: click.Context) -> None:
             ctx.obj["project_dir"],
             ctx.args,
             read_only_domain=ctx.obj.get("read_only_domain", False),
+            hide_suggested_questions=ctx.obj.get("hide_suggested_questions", False),
+            hide_build_context_hint=ctx.obj.get("hide_build_context_hint", False),
         )
     except subprocess.CalledProcessError as e:
         click.echo(f"Error running Streamlit: {e}", err=True)

databao_cli/ui/app.py

@@ -177,11 +177,28 @@ def _clear_all_chat_threads() -> None:
         chat.thread = None
 
 
+def invalidate_agent(status_message: str = "Reloading project...") -> None:
+    """Reset the in-memory agent and project so they are re-created on the next rerun.
+
+    Call this whenever the set of datasources changes (add, remove, save)
+    to ensure the agent does not reference stale catalogs/schemas.
+    """
+    st.session_state.databao_project = None
+    st.session_state.agent = None
+    _clear_all_chat_threads()
+    set_status(AppStatus.INITIALIZING, status_message)
+
+
 def is_read_only_domain() -> bool:
     """Check whether domain-editing operations are disabled."""
     return cast(bool, st.session_state.get("_read_only_domain", False))
 
 
+def is_hide_build_context_hint() -> bool:
+    """Check whether the build context hint and step are hidden."""
+    return cast(bool, st.session_state.get("_hide_build_context_hint", False))
+
+
 def _is_project_ready(project_dir: Path) -> bool:
     """Check if the Databao project is fully set up and ready for normal use."""
     project = find_project(project_dir)
@@ -503,11 +520,28 @@ def main() -> None:
         default=False,
         help="Disable all domain-editing operations (init, datasources, build)",
     )
+    parser.add_argument(
+        "--hide-suggested-questions",
+        action="store_true",
+        default=False,
+        help="Hide the suggested questions on the empty chat screen",
+    )
+    parser.add_argument(
+        "--hide-build-context-hint",
+        action="store_true",
+        default=False,
+        help="Hide the 'Context isn't built yet' warning on the empty chat screen and skip the Build Context step in the setup wizard",
+    )
     try:
         args = parser.parse_args()
     except SystemExit:
         st.warning("Failed to parse arguments. Using current directory as project directory.")
-        args = argparse.Namespace(project_dir=None, read_only_domain=False)
+        args = argparse.Namespace(
+            project_dir=None,
+            read_only_domain=False,
+            hide_suggested_questions=False,
+            hide_build_context_hint=False,
+        )
 
     project_dir = Path(args.project_dir) if args.project_dir else Path.cwd()
 
@@ -515,6 +549,10 @@ def main() -> None:
 
     if "_read_only_domain" not in st.session_state:
         st.session_state._read_only_domain = args.read_only_domain
+    if "_hide_suggested_questions" not in st.session_state:
+        st.session_state._hide_suggested_questions = args.hide_suggested_questions
+    if "_hide_build_context_hint" not in st.session_state:
+        st.session_state._hide_build_context_hint = args.hide_build_context_hint
 
     if "_setup_mode_active" not in st.session_state:
         project = find_project(project_dir)

databao_cli/ui/cli.py

@@ -21,6 +21,8 @@ def bootstrap_streamlit_app(
     streamlit_args: list[str] | None = None,
     *,
     read_only_domain: bool = False,
+    hide_suggested_questions: bool = False,
+    hide_build_context_hint: bool = False,
 ) -> None:
     """Bootstrap the UI."""
 
@@ -32,6 +34,10 @@ def bootstrap_streamlit_app(
     app_args = ["--project-dir", str(project_path)]
     if read_only_domain:
         app_args.append("--read-only-domain")
+    if hide_suggested_questions:
+        app_args.append("--hide-suggested-questions")
+    if hide_build_context_hint:
+        app_args.append("--hide-build-context-hint")
 
     subprocess.run(
         [sys.executable, "-m", "streamlit", "run", app_path, *streamlit_args, "--", *app_args],

databao_cli/ui/components/chat.py

@@ -13,6 +13,7 @@ from databao_cli.ui.services import (
     is_query_running,
     save_current_chat,
     start_query_execution,
+    stop_query,
 )
 from databao_cli.ui.suggestions import (
     check_suggestions_completion,
@@ -51,6 +52,9 @@ def render_assistant_message(
         elif message.content:
             st.error(message.content)
 
+        if message.metadata.get("stopped"):
+            st.warning("Query was stopped by user")
+
 
 def _truncate_question(question: str, max_len: int = 60) -> tuple[str, bool]:
     """Truncate a question for display, returning (display_text, was_truncated)."""
@@ -73,8 +77,9 @@ def render_welcome_component(chat: "ChatSession") -> None:
     - ready: shows questions with appropriate subtitle
     """
     status = st.session_state.get("suggestions_status", "not_started")
+    hide_suggested_questions = bool(st.session_state.get("_hide_suggested_questions"))
 
-    if status == "not_started":
+    if status == "not_started" and not hide_suggested_questions:
         agent: Agent | None = st.session_state.get("agent")
         if agent is not None:
             start_suggestions_generation(agent)
@@ -94,7 +99,7 @@ def render_welcome_component(chat: "ChatSession") -> None:
         unsafe_allow_html=True,
     )
 
-    if status == "loading":
+    if status == "loading" and not hide_suggested_questions:
         st.markdown(
             "<p style='text-align: center; color: #888; font-size: 0.9em; margin-top: 1em;'>"
             "🔄 Analyzing your data to suggest questions..."
@@ -106,7 +111,7 @@ def render_welcome_component(chat: "ChatSession") -> None:
     questions: list[str] = st.session_state.get("suggested_questions", [])
     is_llm_generated: bool = st.session_state.get("suggestions_are_llm_generated", False)
 
-    if questions:
+    if questions and not hide_suggested_questions:
         if is_llm_generated:
             st.markdown(
                 "<p style='text-align: center; color: #888; font-size: 0.9em; margin-top: 1em;'>"
@@ -360,14 +365,170 @@ def _should_show_welcome(chat: "ChatSession") -> bool:
     return not has_messages and not query_running
 
 
-def render_chat_interface(chat: "ChatSession") -> None:
-    """Render the complete chat interface."""
-    query_running = is_query_running(chat)
+def _has_stopped_exchange(chat: "ChatSession") -> bool:
+    """Check if the last exchange was stopped *without* producing results.
+
+    Returns False when the stop happened during the visualization phase,
+    because the data result is already available and the conversation can
+    continue normally.
+    """
+    if not chat.messages:
+        return False
+    last = chat.messages[-1]
+    return last.role == "assistant" and bool(last.metadata.get("stopped")) and last.result is None
+
+
+def _remove_stopped_exchange(chat: "ChatSession") -> None:
+    """Remove the trailing user + stopped-assistant message pair."""
+    if not chat.messages:
+        return
+    last = chat.messages[-1]
+    if last.role == "assistant" and last.metadata.get("stopped"):
+        chat.messages.pop()
+    if chat.messages and chat.messages[-1].role == "user":
+        chat.messages.pop()
+
+
+@st.dialog("Overwrite previous request?")
+def _confirm_overwrite_dialog() -> None:
+    """Modal dialog asking whether to discard the stopped exchange."""
+    st.markdown("The previous request was stopped. Sending a new request will remove it.")
+    col1, col2 = st.columns(2)
+    with col1:
+        if st.button("OK", use_container_width=True, type="primary"):
+            st.session_state["overwrite_confirmed"] = True
+            st.rerun()
+    with col2:
+        if st.button("Cancel", use_container_width=True):
+            st.session_state.pop("pending_query", None)
+            st.rerun()
+
+
+def _handle_stop_click(chat: "ChatSession") -> None:
+    """Stop the running query and record the partial result.
+
+    If the data phase already produced a message (``viz_pending``), marks
+    it as stopped.  Otherwise creates a new assistant message with
+    whatever thinking text was captured so far.
+    """
+    thinking_text = stop_query(chat)
+    if thinking_text is None:
+        return
 
-    user_input = st.chat_input("Ask a question about your data...", disabled=query_running)
+    pending = _find_pending_viz_message(chat)
+    if pending is not None:
+        pending.viz_pending = False
+        pending.metadata["stopped"] = True
+    else:
+        chat.messages.append(
+            ChatMessage(
+                role="assistant",
+                thinking=thinking_text or None,
+                content="",
+                metadata={"stopped": True},
+            )
+        )
+    save_current_chat()
+
+
+def _render_chat_input_bar(chat: "ChatSession", query_running: bool) -> None:
+    """Render the chat input bar.
+
+    When ``query_running`` is False, shows an enabled text input inside a
+    form (so Enter submits) and a send button.
+
+    When ``query_running`` is True (background query **or** manual plot
+    in progress), the input is disabled and a stop button is shown
+    instead.
+    """
+    if query_running:
+        col1, col2 = st.columns([12, 1], vertical_alignment="bottom")
+        with col1:
+            st.text_input(
+                "Message",
+                placeholder="Query in progress...",
+                disabled=True,
+                label_visibility="collapsed",
+                key="chat_input_disabled",
+            )
+        with col2:
+            if st.button(
+                ":material/stop_circle:",
+                key="stop_btn",
+                use_container_width=True,
+                type="primary",
+                help="Stop the running query",
+            ):
+                _handle_stop_click(chat)
+                st.rerun()
+    else:
+        with st.form("chat_input_form", clear_on_submit=True, border=False):
+            col1, col2 = st.columns([12, 1], vertical_alignment="bottom")
+            with col1:
+                user_input = st.text_input(
+                    "Message",
+                    placeholder="Ask a question about your data...",
+                    label_visibility="collapsed",
+                    key="chat_input",
+                )
+            with col2:
+                submitted = st.form_submit_button(
+                    ":material/send:",
+                    use_container_width=True,
+                )
+
+            if submitted and user_input:
+                if _has_stopped_exchange(chat):
+                    st.session_state["pending_query"] = user_input
+                    _confirm_overwrite_dialog()
+                else:
+                    user_message = ChatMessage(role="user", content=user_input)
+                    chat.messages.append(user_message)
+                    save_current_chat()
+                    start_background_query(chat, user_input)
+                    st.rerun()
+
+
+def _process_pending_overwrite(chat: "ChatSession") -> None:
+    """Process a confirmed overwrite of a stopped exchange.
+
+    Called at the top of ``render_chat_interface`` (before any widgets are
+    rendered), so no ``st.rerun()`` is needed — the updated state is
+    picked up by the rendering that follows in the same script run.
+    """
+    if not st.session_state.pop("overwrite_confirmed", False):
+        return
+
+    pending = st.session_state.pop("pending_query", None)
+    if not pending:
+        return
+
+    _remove_stopped_exchange(chat)
+    user_message = ChatMessage(role="user", content=pending)
+    chat.messages.append(user_message)
+    save_current_chat()
+    start_background_query(chat, pending)
+
+
+def render_chat_interface(chat: "ChatSession") -> None:
+    """Render the complete chat interface.
+
+    Orchestrates, in order:
+    1. Processing a confirmed overwrite of a stopped exchange.
+    2. Chat history, thinking section, and polling fragments.
+    3. The input bar (enabled or disabled).
+    4. A pending manual "Generate Plot" execution (blocking, runs last
+       so the disabled input bar is already visible).
+    """
+    _process_pending_overwrite(chat)
 
     agent: Agent | None = st.session_state.get("agent")
-    if agent is not None and not agent.domain.is_context_built() and not chat.messages:
+    if (
+        agent is not None
+        and not agent.domain.is_context_built()
+        and not chat.messages
+        and not st.session_state.get("_hide_build_context_hint")
+    ):
         st.markdown(
             "⚠️ Context isn't built yet. "
             '<a href="/context-settings#build-context" target="_self">Build context</a> '
@@ -375,15 +536,7 @@ def render_chat_interface(chat: "ChatSession") -> None:
             unsafe_allow_html=True,
         )
 
-    if user_input:
-        user_message = ChatMessage(role="user", content=user_input)
-        chat.messages.append(user_message)
-
-        save_current_chat()
-
-        start_background_query(chat, user_input)
-
-        st.rerun()
+    query_running = is_query_running(chat) or "pending_plot_message_index" in st.session_state
 
     if handle_query_completion(chat):
         st.rerun()
@@ -403,3 +556,11 @@ def render_chat_interface(chat: "ChatSession") -> None:
 
         if query_running:
             _query_polling_fragment()
+
+    st.markdown("<div style='height: 2em'></div>", unsafe_allow_html=True)
+    _render_chat_input_bar(chat, query_running)
+
+    if "pending_plot_message_index" in st.session_state:
+        from databao_cli.ui.components.results import execute_pending_plot
+
+        execute_pending_plot(chat)

databao_cli/ui/components/datasource_manager.py

@@ -11,6 +11,7 @@ from typing import Any, cast
 import streamlit as st
 from databao_context_engine import ConfiguredDatasource, DatasourceConnectionStatus
 
+from databao_cli.ui.app import invalidate_agent
 from databao_cli.ui.components.datasource_form import render_datasource_config_form
 from databao_cli.ui.services.dce_operations import (
     add_datasource,
@@ -216,6 +217,7 @@ def _render_existing_datasource(project_dir: Path, ds: ConfiguredDatasource, idx
                         try:
                             remove_datasource(project_dir, ds_id)
                             st.session_state.pop(f"_confirm_remove_{idx}", None)
+                            invalidate_agent("Datasource removed, reinitializing agent...")
                             st.rerun()
                         except Exception as e:
                             st.error(f"Remove failed: {e}")

databao_cli/ui/components/results.py

@@ -308,13 +308,17 @@ def render_visualization_and_actions(
     *,
     is_latest: bool = False,
 ) -> None:
-    """Fragment that renders visualization and action buttons together.
+    """Fragment that renders visualization status and action buttons.
 
-    This is a fragment so that when action buttons trigger updates (e.g., Generate Plot),
-    only this fragment reruns, showing the new visualization without a full app rerun.
+    Re-reads ``has_visualization``, ``visualization_data``, ``viz_pending``
+    and ``viz_error`` from ``chat.messages[message_index]`` on every run so
+    it reflects the latest state.
 
-    The fragment reads has_visualization and visualization_data from chat.messages
-    so it can see updates made by button click handlers on fragment rerun.
+    When a manual "Generate Plot" is pending (flagged via session state),
+    this fragment only renders a placeholder; the actual blocking
+    ``thread.plot()`` call is executed by ``execute_pending_plot`` in the
+    main script — after the input bar has been rendered — to keep the UI
+    responsive and the input visibly disabled.
     """
     current_chat = _get_current_chat()
     if current_chat is None:
@@ -337,11 +341,15 @@ def render_visualization_and_actions(
         viz_pending = False
         viz_error = None
 
+    if st.session_state.get("pending_plot_message_index") == message_index:
+        st.info("Generating visualization...", icon="📈")
+        return
+
     if viz_pending:
         st.info("Generating visualization...", icon="📈")
     elif viz_error:
         st.error(f"Failed to generate visualization: {viz_error}")
-    elif has_visualization or thread._visualization_result is not None or visualization_data is not None:
+    elif has_visualization or visualization_data is not None or (is_latest and thread._visualization_result is not None):
         render_visualization_section(thread, visualization_data)
 
     if is_latest and not viz_pending and not viz_error:
@@ -354,9 +362,11 @@ def _render_and_handle_action_buttons(
     message_index: int,
     has_visualization: bool,
 ) -> None:
-    """Render action buttons and handle clicks inline.
+    """Render action buttons and handle clicks.
 
-    Called from within the fragment, so button clicks can trigger fragment-scoped reruns.
+    "Generate Plot" sets a session-state flag and triggers a full app
+    rerun so that the input bar is rendered as disabled before the
+    blocking ``thread.plot()`` call begins.
     """
     from databao_cli.ui.services import is_query_running
 
@@ -377,31 +387,38 @@ def _render_and_handle_action_buttons(
         button_key = f"action_generate_plot_{message_index}"
         clicked = st.button("📈 Generate Plot", key=button_key, width="stretch", disabled=is_processing)
         if clicked and not is_processing:
-            _handle_generate_plot(chat, message_index)
+            st.session_state["pending_plot_message_index"] = message_index
+            st.rerun(scope="app")
 
 
-def _handle_generate_plot(chat: "ChatSession", message_index: int) -> None:
-    """Handle Generate Plot button click.
+def execute_pending_plot(chat: "ChatSession") -> None:
+    """Execute a pending manual "Generate Plot" request.
 
-    Called from within a fragment, so st.rerun() will only rerun the fragment.
+    Must be called **after** the input bar has been rendered (disabled)
+    so the user sees a locked UI while the blocking ``thread.plot()``
+    runs.  On completion (or error) triggers a full app rerun.
     """
     thread = chat.thread
     if thread is None:
         return
 
-    with st.spinner("Generating visualization..."):
-        try:
-            thread.plot()
+    message_index: int = st.session_state.pop("pending_plot_message_index")
+
+    try:
+        thread.plot()
 
-            messages = chat.messages
-            if message_index < len(messages):
-                messages[message_index].has_visualization = True
-                messages[message_index].visualization_data = _extract_visualization_data(thread)
-                save_current_chat()
+        messages = chat.messages
+        if message_index < len(messages):
+            messages[message_index].has_visualization = True
+            messages[message_index].visualization_data = _extract_visualization_data(thread)
+            save_current_chat()
+    except Exception as e:
+        logger.exception("Failed to generate visualization")
+        if message_index < len(chat.messages):
+            chat.messages[message_index].metadata["viz_error"] = str(e)
+            save_current_chat()
 
-            st.rerun()
-        except Exception as e:
-            st.error(f"Failed to generate visualization: {e}")
+    st.rerun(scope="app")
 
 
 def render_execution_result(

databao_cli/ui/pages/context_settings.py

@@ -6,10 +6,9 @@ import streamlit as st
 from databao.agent import Agent
 
 from databao_cli.project.layout import ProjectLayout
-from databao_cli.ui.app import _clear_all_chat_threads, is_read_only_domain
+from databao_cli.ui.app import invalidate_agent, is_read_only_domain
 from databao_cli.ui.components.datasource_manager import render_datasource_manager
 from databao_cli.ui.components.icons import get_db_type_and_icon
-from databao_cli.ui.components.status import AppStatus, set_status
 from databao_cli.ui.project_utils import DatabaoProjectStatus, databao_project_status
 from databao_cli.ui.services.build_service import render_build_section
 from databao_cli.ui.services.dce_operations import get_status_info
@@ -50,10 +49,7 @@ def render_context_settings_page() -> None:
         st.info("No Databao project detected. Initialize one from the Setup page.")
 
     if reload_clicked:
-        st.session_state.databao_project = None
-        st.session_state.agent = None
-        _clear_all_chat_threads()
-        set_status(AppStatus.INITIALIZING, "Reloading project...")
+        invalidate_agent("Reloading project...")
         st.rerun()
 
     # ---- Datasource Management section ----

databao_cli/ui/pages/welcome.py

@@ -122,16 +122,18 @@ def render_welcome_page() -> None:
 def render_setup_wizard_page() -> None:
     """Render the setup wizard for first-time project configuration.
 
-    Four sections, all visible, disabled based on prerequisites:
+    The wizard is organized into up to five sections, which are disabled
+    based on prerequisites:
     1. Initialize Project
     2. Configure Datasources
-    3. Build Context
-    4. Ready
+    3. Configure Agent
+    4. Build Context (optional; can be hidden via feature flag)
+    5. Ready
 
     When read-only-domain mode is active, editing sections are disabled with
     an explanation banner.
     """
-    from databao_cli.ui.app import _create_new_chat, is_read_only_domain
+    from databao_cli.ui.app import _create_new_chat, is_hide_build_context_hint, is_read_only_domain
     from databao_cli.ui.components.datasource_manager import render_datasource_manager
     from databao_cli.ui.services.build_service import (
         get_build_status,
@@ -141,6 +143,7 @@ def render_setup_wizard_page() -> None:
 
     project_dir: Path = st.session_state.get("_project_dir", Path.cwd())
     read_only = is_read_only_domain()
+    hide_build_context = is_hide_build_context_hint()
 
     _col1, col2, _col3 = st.columns([1, 3, 1])
 
@@ -253,31 +256,33 @@ def render_setup_wizard_page() -> None:
         st.markdown("---")
 
         # ---- Section 4: Build Context ----
-        _render_section_header(
-            "4",
-            "Build Context (Optional)",
-            completed=build_started_or_done,
-            enabled=has_datasources,
-        )
-
-        if not has_datasources:
-            st.caption("Complete the previous steps first.")
-        elif project is None:
-            st.caption("Add at least one datasource first.")
-        elif read_only:
-            render_build_section(project.root_domain_dir, read_only=True)
-        else:
-            st.markdown(
-                "Building the context indexes your datasources so Databao can better understand "
-                "your data structure and provide higher-quality answers."
+        if not hide_build_context:
+            _render_section_header(
+                "4",
+                "Build Context (Optional)",
+                completed=build_started_or_done,
+                enabled=has_datasources,
             )
-            render_build_section(project.root_domain_dir)
 
-        st.markdown("---")
+            if not has_datasources:
+                st.caption("Complete the previous steps first.")
+            elif project is None:
+                st.caption("Add at least one datasource first.")
+            elif read_only:
+                render_build_section(project.root_domain_dir, read_only=True)
+            else:
+                st.markdown(
+                    "Building the context indexes your datasources so Databao can better understand "
+                    "your data structure and provide higher-quality answers."
+                )
+                render_build_section(project.root_domain_dir)
+
+            st.markdown("---")
 
         # ---- Final Section: Start Using Databao ----
+        final_step = "4" if hide_build_context else "5"
         _render_section_header(
-            "5",
+            final_step,
             "Start Using Databao",
             completed=False,
             enabled=has_datasources,
@@ -286,12 +291,12 @@ def render_setup_wizard_page() -> None:
         if not has_datasources:
             st.caption("Add at least one datasource first.")
         else:
-            if build_status == "running":
+            if not hide_build_context and build_status == "running":
                 st.info(
                     "The build is still in progress, but you can start exploring Databao. "
                     "Some features may not work until the build completes."
                 )
-            elif not build_started_or_done:
+            elif not hide_build_context and not build_started_or_done:
                 st.markdown(
                     "You're ready to start using Databao! Consider building the context above for the best experience."
                 )

databao_cli/ui/services/__init__.py

@@ -18,6 +18,7 @@ from databao_cli.ui.services.query_executor import (
     get_query_phase,
     is_query_running,
     start_query_execution,
+    stop_query,
 )
 from databao_cli.ui.services.settings_persistence import (
     delete_settings,
@@ -55,5 +56,6 @@ __all__ = [
     "save_current_chat",
     "save_settings",
     "start_query_execution",
+    "stop_query",
     "trigger_title_generation",
 ]

databao_cli/ui/services/query_executor.py

@@ -4,6 +4,7 @@ This module provides background execution of queries so that they continue
 running when users switch between chats. The pattern follows suggestions.py.
 """
 
+import ctypes
 import logging
 import threading
 from dataclasses import dataclass
@@ -21,6 +22,8 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger(__name__)
 
+_STOP_TIMEOUT_SECONDS = 5.0
+
 
 @dataclass
 class QueryResult:
@@ -100,6 +103,8 @@ class QueryThread(threading.Thread):
                 visualization_data=visualization_data,
                 error=None,
             )
+        except (KeyboardInterrupt, SystemExit):
+            logger.debug("Query thread stopped")
         except Exception as e:
             logger.exception("Query execution failed")
             thinking_text = self.writer.getvalue() if self.writer else ""
@@ -202,6 +207,95 @@ def check_query_completion(chat: "ChatSession") -> QueryResult | None:
     return None
 
 
+def _raise_in_thread(thread: threading.Thread, exc_type: type) -> bool:
+    """Raise an exception asynchronously in a running thread.
+
+    Uses ``PyThreadState_SetAsyncExc`` to inject an exception into the
+    target thread.  The exception will be delivered the next time the
+    thread executes Python bytecode (it may not interrupt C-level
+    blocking calls immediately).
+
+    Returns True if the exception was successfully scheduled.
+    """
+    if not thread.is_alive():
+        return False
+    thread_id = thread.ident
+    if thread_id is None:
+        return False
+    res = ctypes.pythonapi.PyThreadState_SetAsyncExc(
+        ctypes.c_ulong(thread_id),
+        ctypes.py_object(exc_type),
+    )
+    if res == 0:
+        logger.warning("Thread %s not found for async exception", thread_id)
+        return False
+    if res > 1:
+        ctypes.pythonapi.PyThreadState_SetAsyncExc(ctypes.c_ulong(thread_id), None)
+        logger.error("Multiple thread states modified — reverted")
+        return False
+    return True
+
+
+def _reap_thread(thread: threading.Thread, timeout: float = _STOP_TIMEOUT_SECONDS) -> None:
+    """Wait for *thread* to finish, then force-kill it if necessary.
+
+    Intended to run in a short-lived daemon reaper thread so the UI is
+    not blocked.
+    """
+    thread.join(timeout=timeout)
+    if thread.is_alive():
+        logger.warning("Query thread did not stop after %.1fs — sending SystemExit", timeout)
+        _raise_in_thread(thread, SystemExit)
+
+
+def stop_query(chat: "ChatSession") -> str | None:
+    """Stop a running background query for a chat.
+
+    Termination strategy:
+    1. Capture partial thinking text and disconnect the writer so old
+       output cannot leak into the next query.
+    2. Discard ``chat.thread`` — the databao ``Thread`` object's internal
+       state (``_lazy_mode``, ``_opas_processed_count``) becomes
+       inconsistent after an interrupted ``ask()`` call and cannot be
+       reused without hitting ``AssertionError``.
+    3. Raise ``KeyboardInterrupt`` in the query thread (graceful stop).
+    4. Spawn a short-lived daemon *reaper* that waits up to
+       ``_STOP_TIMEOUT_SECONDS`` and, if the thread is still alive,
+       force-kills it with ``SystemExit``.
+
+    Returns:
+        The captured thinking text (possibly empty ``""``) if a query
+        was stopped, or ``None`` if nothing was running.
+    """
+    if chat.query_status not in ("running", "visualizing"):
+        return None
+
+    thinking_text = ""
+    if chat.writer:
+        thinking_text = chat.writer.getvalue()
+        chat.writer._on_write = None
+
+    query_thread = chat.query_thread
+
+    chat.query_thread = None
+    chat.query_status = "idle"
+    chat.thread = None
+    chat.writer = None
+
+    if query_thread is not None and query_thread.is_alive():
+        _raise_in_thread(query_thread, KeyboardInterrupt)
+        reaper = threading.Thread(
+            target=_reap_thread,
+            args=(query_thread,),
+            name="query_reaper",
+            daemon=True,
+        )
+        reaper.start()
+
+    logger.info(f"Stopped query for chat {chat.id}")
+    return thinking_text
+
+
 def is_query_running(chat: "ChatSession") -> bool:
     """Check if a chat has a query currently running (data or visualization phase).
 

databao_cli/ui/suggestions.py

@@ -6,6 +6,7 @@ from concurrent.futures import TimeoutError as FuturesTimeoutError
 from typing import TYPE_CHECKING
 
 import streamlit as st
+from databao.agent.duckdb import describe_duckdb_schema
 from langchain_core.messages import HumanMessage, SystemMessage
 from pydantic import BaseModel, Field
 
@@ -40,10 +41,45 @@ class SuggestedQuestions(BaseModel):
     )
 
 
+def _ensure_sources_initialized(agent: "Agent") -> None:
+    """Ensure the executor has data sources attached to DuckDB.
+
+    The executor lazily attaches data sources on first query execution.
+    We need them earlier for schema introspection during suggestions.
+    """
+    executor = agent.executor
+    init_fn = getattr(executor, "_init_sources_from_domain", None)
+    if init_fn is None:
+        return
+    try:
+        init_fn(agent.domain)
+    except Exception:
+        logger.debug("Failed to initialize sources for suggestions", exc_info=True)
+
+
+def _get_duckdb_schema(agent: "Agent") -> str | None:
+    """Try to extract the DuckDB schema from the agent's executor."""
+    conn = getattr(agent.executor, "_duckdb_connection", None)
+    if conn is None:
+        return None
+    _ensure_sources_initialized(agent)
+    try:
+        schema: str = describe_duckdb_schema(conn)
+        if schema and schema != "(no tables found)":
+            return schema
+    except Exception:
+        logger.debug("Failed to describe DuckDB schema for suggestions", exc_info=True)
+    return None
+
+
 def _build_context_from_sources(agent: "Agent") -> str:
     """Extract schema/context information from agent sources."""
     context_parts: list[str] = []
 
+    db_schema = _get_duckdb_schema(agent)
+    if db_schema:
+        context_parts.append(f"Database schema:\n{db_schema}")
+
     for name, db_source in agent.dbs.items():
         if db_source.description:
             context_parts.append(f"Database '{name}':\n{db_source.description}")