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
1618from .datasource import DataFrameSource , DataSource , SQLAlchemySource
1719from .tools import tool_query , tool_reset_dashboard , tool_update_dashboard
2022 import chatlas
2123 import pandas as pd
2224 from narwhals .stable .v1 .typing import IntoFrame
25+ from shiny .bookmark import BookmarkState , RestoreState
2326
2427
2528@dataclass
@@ -127,6 +130,97 @@ def __init__(
127130 self .greeting = config .greeting
128131 self .client = config .client
129132
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+
130224 def sidebar (
131225 self ,
132226 id : str ,
@@ -187,14 +281,17 @@ def _ui_wrapper(**ui_kwargs):
187281
188282 return _ui_wrapper (id , ** kwargs )
189283
190- def server (self , id : str ) :
284+ def server (self , id : str , * , enable_bookmarking : bool = True ) -> QueriedValues :
191285 """
192286 Initialize the querychat server logic.
193287
194288 Parameters
195289 ----------
196290 id
197291 An ID corresponding to the UI component.
292+ enable_bookmarking
293+ Whether to enable bookmarking for this chat session. For this to take
294+ effect, the Shiny app must also have a `bookmark_store` configured.
198295
199296 Returns
200297 -------
@@ -214,13 +311,14 @@ def server(self, id: str):
214311 def mod_server_wrapper (
215312 input : Inputs ,
216313 output : Outputs ,
217- session : Session ,
314+ session : Session
218315 ):
219316 return _server_impl (
220317 input ,
221318 output ,
222319 session ,
223320 querychat_config = config ,
321+ enable_bookmarking = enable_bookmarking ,
224322 )
225323
226324 return mod_server_wrapper (id )
@@ -644,6 +742,8 @@ def _server_impl(
644742 output : Outputs ,
645743 session : Session ,
646744 querychat_config : QueryChatConfig ,
745+ * ,
746+ enable_bookmarking : bool = True ,
647747) -> QueriedValues :
648748 data_source = querychat_config .data_source
649749 system_prompt = querychat_config .system_prompt
@@ -653,6 +753,7 @@ def _server_impl(
653753 # Reactive values to store state
654754 current_title = ReactiveStringOrNone (None )
655755 current_query = ReactiveString ("" )
756+ has_greeted = reactive .value [bool ](False ) # noqa: FBT003
656757
657758 @reactive .calc
658759 def filtered_df ():
@@ -710,6 +811,9 @@ def _():
710811
711812 @reactive .effect
712813 async def greet_on_startup ():
814+ if has_greeted ():
815+ return
816+
713817 if querychat_config .greeting :
714818 await chat_ui .append_message (greeting )
715819 elif querychat_config .greeting is None :
@@ -719,5 +823,29 @@ async def greet_on_startup():
719823 )
720824 await chat_ui .append_message_stream (stream )
721825
826+ has_greeted .set (True )
827+
828+ if enable_bookmarking :
829+ chat_ui .enable_bookmarking (client )
830+
831+ def _on_bookmark (x : BookmarkState ) -> None :
832+ vals = x .values # noqa: PD011
833+ vals ["querychat_current_query" ] = current_query .get ()
834+ vals ["querychat_current_title" ] = current_title .get ()
835+ vals ["querychat_has_greeted" ] = has_greeted .get ()
836+
837+ session .bookmark .on_bookmark (_on_bookmark )
838+
839+ def _on_restore (x : RestoreState ) -> None :
840+ vals = x .values # noqa: PD011
841+ if "querychat_current_query" in vals :
842+ current_query .set (vals ["querychat_current_query" ])
843+ if "querychat_current_title" in vals :
844+ current_title .set (vals ["querychat_current_title" ])
845+ if "querychat_has_greeted" in vals :
846+ has_greeted .set (vals ["querychat_has_greeted" ])
847+
848+ session .bookmark .on_restore (_on_restore )
849+
722850 # Return the interface for other components to use
723851 return QueriedValues (filtered_df , current_query , current_title , chat )
0 commit comments