1010import chevron
1111import shinychat
1212import sqlalchemy
13- from shiny import Inputs , Outputs , Session , module , reactive , ui
13+ from shiny import App , Inputs , Outputs , Session , module , reactive , render , req , ui
14+ from shinychat import output_markdown_stream
1415
16+ from ._icons import bs_icon
1517from ._utils import normalize_client
18+ from .datasource import DataFrameSource , DataSource , SQLAlchemySource
1619from .tools import tool_query , tool_reset_dashboard , tool_update_dashboard
1720
1821if TYPE_CHECKING :
1922 import chatlas
2023 import pandas as pd
2124 from narwhals .stable .v1 .typing import IntoFrame
22-
23- from .datasource import DataFrameSource , DataSource , SQLAlchemySource
25+ from shiny .bookmark import BookmarkState , RestoreState
2426
2527
2628@dataclass
@@ -128,6 +130,97 @@ def __init__(
128130 self .greeting = config .greeting
129131 self .client = config .client
130132
133+ def app (self , bookmark_store : Literal ["url" , "server" , "disable" ] = "url" ) -> App :
134+ """
135+ Quickly chat with a dataset.
136+
137+ Creates a Shiny app with a chat sidebar and data table view -- providing a
138+ quick-and-easy way to start chatting with your data.
139+
140+ Parameters
141+ ----------
142+ bookmark_store
143+ The bookmarking store to use for the Shiny app. Options are:
144+ - `"url"`: Store bookmarks in the URL (default).
145+ - `"server"`: Store bookmarks on the server.
146+ - `"disable"`: Disable bookmarking.
147+
148+ Returns
149+ -------
150+ :
151+ A Shiny App object that can be run with `app.run()` or served with `shiny run`.
152+
153+ """
154+ enable_bookmarking = bookmark_store != "disable"
155+ table_name = self .data_source .table_name
156+
157+ def app_ui (request ):
158+ return ui .page_sidebar (
159+ self .sidebar ("chat" ),
160+ ui .card (
161+ ui .card_header (
162+ ui .div (
163+ ui .div (
164+ bs_icon ("terminal-fill" ),
165+ ui .output_text ("query_title" , inline = True ),
166+ class_ = "d-flex align-items-center gap-2" ,
167+ ),
168+ ui .output_ui ("ui_reset" , inline = True ),
169+ class_ = "hstack gap-3" ,
170+ ),
171+ ),
172+ ui .output_ui ("sql_output" ),
173+ fill = False ,
174+ style = "max-height: 33%;" ,
175+ ),
176+ ui .card (
177+ ui .card_header (bs_icon ("table" ), " Data" ),
178+ ui .output_data_frame ("dt" ),
179+ ),
180+ title = ui .span ("querychat with " , ui .code (table_name )),
181+ class_ = "bslib-page-dashboard" ,
182+ fillable = True ,
183+ )
184+
185+ def app_server (input : Inputs , output : Outputs , session : Session ):
186+ qc = self .server ("chat" , enable_bookmarking = enable_bookmarking )
187+
188+ @render .text
189+ def query_title ():
190+ return qc .title () or "SQL Query"
191+
192+ @render .ui
193+ def ui_reset ():
194+ req (qc .sql ())
195+ return ui .input_action_button (
196+ "reset_query" ,
197+ "Reset Query" ,
198+ class_ = "btn btn-outline-danger btn-sm lh-1 ms-auto" ,
199+ )
200+
201+ @reactive .effect
202+ @reactive .event (input .reset_query )
203+ def _ ():
204+ qc .sql ("" )
205+ qc .title (None )
206+
207+ @render .data_frame
208+ def dt ():
209+ return qc .df ()
210+
211+ @render .ui
212+ def sql_output ():
213+ sql = qc .sql () or f"SELECT * FROM { table_name } "
214+ sql_code = f"```sql\n { sql } \n ```"
215+ return output_markdown_stream (
216+ "sql_code" ,
217+ content = sql_code ,
218+ auto_scroll = False ,
219+ width = "100%" ,
220+ )
221+
222+ return App (app_ui , app_server , bookmark_store = bookmark_store )
223+
131224 def sidebar (
132225 self ,
133226 id : str ,
@@ -182,14 +275,17 @@ def ui(self, id: str, **kwargs):
182275 """
183276 return mod_ui (id , ** kwargs )
184277
185- def server (self , id : str ) :
278+ def server (self , id : str , * , enable_bookmarking : bool = True ) -> QueriedValues :
186279 """
187280 Initialize the querychat server logic.
188281
189282 Parameters
190283 ----------
191284 id
192285 An ID corresponding to the UI component.
286+ enable_bookmarking
287+ Whether to enable bookmarking for this chat session. For this to take
288+ effect, the Shiny app must also have a `bookmark_store` configured.
193289
194290 Returns
195291 -------
@@ -204,7 +300,11 @@ def server(self, id: str):
204300 greeting = self .greeting ,
205301 client = self .client ,
206302 )
207- return mod_server (id , querychat_config = config )
303+ return mod_server (
304+ id ,
305+ querychat_config = config ,
306+ enable_bookmarking = enable_bookmarking ,
307+ )
208308
209309 def generate_greeting (self , * , echo : Literal ["none" , "text" ] = "none" ):
210310 """
@@ -663,6 +763,8 @@ def mod_server(
663763 output : Outputs ,
664764 session : Session ,
665765 querychat_config : QueryChatConfig ,
766+ * ,
767+ enable_bookmarking : bool = True ,
666768) -> QueriedValues :
667769 """
668770 Initialize the querychat server.
@@ -678,6 +780,7 @@ def mod_server(
678780 # Reactive values to store state
679781 current_title = ReactiveStringOrNone (None )
680782 current_query = ReactiveString ("" )
783+ has_greeted = reactive .value [bool ](False ) # noqa: FBT003
681784
682785 @reactive .calc
683786 def filtered_df ():
@@ -735,6 +838,9 @@ def _():
735838
736839 @reactive .effect
737840 async def greet_on_startup ():
841+ if has_greeted ():
842+ return
843+
738844 if querychat_config .greeting :
739845 await chat_ui .append_message (greeting )
740846 elif querychat_config .greeting is None :
@@ -744,5 +850,29 @@ async def greet_on_startup():
744850 )
745851 await chat_ui .append_message_stream (stream )
746852
853+ has_greeted .set (True )
854+
855+ if enable_bookmarking :
856+ chat_ui .enable_bookmarking (client )
857+
858+ def _on_bookmark (x : BookmarkState ) -> None :
859+ vals = x .values # noqa: PD011
860+ vals ["querychat_current_query" ] = current_query .get ()
861+ vals ["querychat_current_title" ] = current_title .get ()
862+ vals ["querychat_has_greeted" ] = has_greeted .get ()
863+
864+ session .bookmark .on_bookmark (_on_bookmark )
865+
866+ def _on_restore (x : RestoreState ) -> None :
867+ vals = x .values # noqa: PD011
868+ if "querychat_current_query" in vals :
869+ current_query .set (vals ["querychat_current_query" ])
870+ if "querychat_current_title" in vals :
871+ current_title .set (vals ["querychat_current_title" ])
872+ if "querychat_has_greeted" in vals :
873+ has_greeted .set (vals ["querychat_has_greeted" ])
874+
875+ session .bookmark .on_restore (_on_restore )
876+
747877 # Return the interface for other components to use
748878 return QueriedValues (filtered_df , current_query , current_title , chat )
0 commit comments