Skip to content

Commit 9cc73d9

Browse files
authored
feat: Allow QueryChat usage outside of Shiny app (#168)
* chore(pkg-r): Tools check their inputs * chore(pkg-r): Consolidate update function into single callback * feat(pkg-r): Add a `$client()` method * chore(pkg-r): Make `tools` configurable in `$client()` * feat(pkg-r): Add `$console()` method * chore(pkg-r): mod_server takes a client-creating function, or a client instance * feat(pkg-r): Allow choosing which tools can be used * refactor(pkg-r): Assemble the system prompt when creating the forked client * feat(pkg-r): Conditionally include tool instructions in system prompt * chore(pkg-r): Minor changes * fix(pkg-r): Fix a couple of typos in $client * tests(pkg-r): `QueryChat$client()` * chore: make r-format * feat(pkg-py): Port changes from R * chore(pkg-py): Types and documentation completeness * tests(pkg-py): Add tests for new querychat features and fix a few small issues along the way * refactor(pkg-r): mock_ellmer_chat_client() * tests(pkg-r): Test console method * chore: Add changelog items * tests(pkg-r): Skip sqlite tests if RSQLite is not available * refactor(pkg-py): Improve types and handling of `tools` * chore(pkg-py): Add missing types * refactor(pkg-py): Factor out `QueryChatSystemPrompt` * refactor(pkg-r): Factor out `QueryChatSystemPrompt` utility class * docs(pkg-py): Minor tweak * `air format` (GitHub Actions) * docs: simplify example again * chore(pkg-py): Update `tools` typing * chore(pkg-r): restore use of `read_utf8()` * tests(pkg-r): Make resilient to whitespace differences * fix(pkg-r): typo in tests
1 parent 93dcaf8 commit 9cc73d9

27 files changed

+2287
-439
lines changed

pkg-py/CHANGELOG.md

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
77

88
## [UNRELEASED]
99

10+
### New features
11+
12+
* `QueryChat.client()` can now create standalone querychat-enabled chat clients with configurable tools and callbacks, enabling use outside of Shiny applications. (#168)
13+
14+
* `QueryChat.console()` was added to launch interactive console-based chat sessions with your data source, with persistent conversation state across invocations. (#168)
1015

16+
* The tools used in a `QueryChat` chatbot are now configurable. Use the new `tools` parameter of `QueryChat()` to select either or both `"query"` or `"update"` tools. Choose `tools=["update"]` if you only want QueryChat to be able to update the dashboard (useful when you want to be 100% certain that the LLM will not see _any_ raw data). (#168)
1117

1218
## [0.3.0] - 2025-12-10
1319

pkg-py/src/querychat/_querychat.py

Lines changed: 236 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,23 @@
1717
from ._datasource import DataFrameSource, DataSource, SQLAlchemySource
1818
from ._icons import bs_icon
1919
from ._querychat_module import GREETING_PROMPT, ServerValues, mod_server, mod_ui
20+
from ._system_prompt import QueryChatSystemPrompt
21+
from ._utils import MISSING, MISSING_TYPE
22+
from .tools import (
23+
UpdateDashboardData,
24+
tool_query,
25+
tool_reset_dashboard,
26+
tool_update_dashboard,
27+
)
2028

2129
if TYPE_CHECKING:
30+
from collections.abc import Callable
31+
2232
import pandas as pd
2333
from narwhals.stable.v1.typing import IntoFrame
2434

35+
TOOL_GROUPS = Literal["update", "query"]
36+
2537

2638
class QueryChatBase:
2739
def __init__(
@@ -32,6 +44,7 @@ def __init__(
3244
id: Optional[str] = None,
3345
greeting: Optional[str | Path] = None,
3446
client: Optional[str | chatlas.Chat] = None,
47+
tools: TOOL_GROUPS | tuple[TOOL_GROUPS, ...] | None = ("update", "query"),
3548
data_description: Optional[str | Path] = None,
3649
categorical_threshold: int = 20,
3750
extra_instructions: Optional[str | Path] = None,
@@ -47,21 +60,28 @@ def __init__(
4760

4861
self.id = id or table_name
4962

63+
self.tools = normalize_tools(tools, default=("update", "query"))
5064
self.greeting = greeting.read_text() if isinstance(greeting, Path) else greeting
5165

52-
prompt = assemble_system_prompt(
53-
self._data_source,
66+
# Store prompt components for lazy assembly
67+
if prompt_template is None:
68+
prompt_template = Path(__file__).parent / "prompts" / "prompt.md"
69+
70+
self._system_prompt = QueryChatSystemPrompt(
71+
prompt_template=prompt_template,
72+
data_source=self._data_source,
5473
data_description=data_description,
5574
extra_instructions=extra_instructions,
5675
categorical_threshold=categorical_threshold,
57-
prompt_template=prompt_template,
5876
)
5977

6078
# Fork and empty chat now so the per-session forks are fast
6179
client = as_querychat_client(client)
6280
self._client = copy.deepcopy(client)
6381
self._client.set_turns([])
64-
self._client.system_prompt = prompt
82+
83+
# Storage for console client
84+
self._client_console = None
6585

6686
def app(
6787
self, *, bookmark_store: Literal["url", "server", "disable"] = "url"
@@ -241,6 +261,158 @@ def generate_greeting(self, *, echo: Literal["none", "output"] = "none"):
241261
client.set_turns([])
242262
return str(client.chat(GREETING_PROMPT, echo=echo))
243263

264+
def client(
265+
self,
266+
*,
267+
tools: TOOL_GROUPS | tuple[TOOL_GROUPS, ...] | None | MISSING_TYPE = MISSING,
268+
update_dashboard: Callable[[UpdateDashboardData], None] | None = None,
269+
reset_dashboard: Callable[[], None] | None = None,
270+
) -> chatlas.Chat:
271+
"""
272+
Create a chat client with registered tools.
273+
274+
This method creates a standalone chat client configured with the
275+
specified tools and callbacks. Each call returns an independent client
276+
instance with its own conversation state.
277+
278+
Parameters
279+
----------
280+
tools
281+
Which tools to include: `"update"`, `"query"`, or both. Can be:
282+
- A single tool string: `"update"` or `"query"`
283+
- A tuple of tools: `("update", "query")`
284+
- `None` or `()` to skip adding any tools
285+
- If not provided (default), uses the tools specified during initialization
286+
update_dashboard
287+
Optional callback function to call when the update_dashboard tool
288+
succeeds. Takes a dict with `"query"` and `"title"` keys. Only used
289+
if `"update"` is in tools.
290+
reset_dashboard
291+
Optional callback function to call when the `tool_reset_dashboard`
292+
is invoked. Takes no arguments. Only used if `"update"` is in tools.
293+
294+
Returns
295+
-------
296+
chatlas.Chat
297+
A configured chat client with tools registered based on the tools parameter.
298+
299+
Examples
300+
--------
301+
```python
302+
from querychat import QueryChat
303+
import pandas as pd
304+
305+
df = pd.DataFrame({"a": [1, 2, 3]})
306+
qc = QueryChat(df, "my_data")
307+
308+
# Create client with all tools (default)
309+
client = qc.client()
310+
response = client.chat("What's the average of column a?")
311+
312+
# Create client with only query tool (single string)
313+
client = qc.client(tools="query")
314+
315+
# Create client with only query tool (tuple)
316+
client = qc.client(tools=("query",))
317+
318+
# Create client with custom callbacks
319+
from querychat import UpdateDashboardData
320+
321+
322+
def my_update(data: UpdateDashboardData):
323+
print(f"Query: {data['query']}, Title: {data['title']}")
324+
325+
326+
client = qc.client(update_dashboard=my_update)
327+
```
328+
329+
"""
330+
tools = normalize_tools(tools, default=self.tools)
331+
332+
chat = copy.deepcopy(self._client)
333+
chat.set_turns([])
334+
335+
chat.system_prompt = self._system_prompt.render(tools)
336+
337+
if tools is None:
338+
return chat
339+
340+
if "update" in tools:
341+
# Default callbacks that do nothing
342+
update_fn = update_dashboard or (lambda _: None)
343+
reset_fn = reset_dashboard or (lambda: None)
344+
345+
chat.register_tool(tool_update_dashboard(self._data_source, update_fn))
346+
chat.register_tool(tool_reset_dashboard(reset_fn))
347+
348+
if "query" in tools:
349+
chat.register_tool(tool_query(self._data_source))
350+
351+
return chat
352+
353+
def console(
354+
self,
355+
*,
356+
new: bool = False,
357+
tools: TOOL_GROUPS | tuple[TOOL_GROUPS, ...] | None = "query",
358+
**kwargs,
359+
) -> None:
360+
"""
361+
Launch an interactive console chat with the data.
362+
363+
This method provides a REPL (Read-Eval-Print Loop) interface for
364+
chatting with your data from the command line. The console session
365+
persists by default, so you can exit and return to continue your
366+
conversation.
367+
368+
Parameters
369+
----------
370+
new
371+
If True, creates a new chat client and starts a fresh conversation.
372+
If False (default), continues the conversation from the previous
373+
console session.
374+
tools
375+
Which tools to include: "update", "query", or both. Can be:
376+
- A single tool string: `"update"` or `"query"`
377+
- A tuple of tools: `("update", "query")`
378+
- `None` or `()` to skip adding any tools
379+
- If not provided (default), defaults to `("query",)` only for
380+
privacy (prevents the LLM from accessing data values)
381+
Ignored if `new=False` and a console session already exists.
382+
**kwargs
383+
Additional arguments passed to the `client()` method when creating a
384+
new client.
385+
386+
Examples
387+
--------
388+
```python
389+
from querychat import QueryChat
390+
import pandas as pd
391+
392+
df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
393+
qc = QueryChat(df, "my_data")
394+
395+
# Start console (query tool only by default)
396+
qc.console()
397+
398+
# Start fresh console with all tools (using tuple)
399+
qc.console(new=True, tools=("update", "query"))
400+
401+
# Start fresh console with all tools (using single string for one tool)
402+
qc.console(new=True, tools="query")
403+
404+
# Continue previous console session
405+
qc.console() # picks up where you left off
406+
```
407+
408+
"""
409+
tools = normalize_tools(tools, default=("query",))
410+
411+
if new or self._client_console is None:
412+
self._client_console = self.client(tools=tools, **kwargs)
413+
414+
self._client_console.console()
415+
244416
@property
245417
def system_prompt(self) -> str:
246418
"""
@@ -252,7 +424,7 @@ def system_prompt(self) -> str:
252424
The system prompt string.
253425
254426
"""
255-
return self._client.system_prompt or ""
427+
return self._system_prompt.render(self.tools)
256428

257429
@property
258430
def data_source(self):
@@ -286,15 +458,43 @@ class QueryChat(QueryChatBase):
286458
"""
287459
Create a QueryChat instance.
288460
461+
QueryChat enables natural language interaction with your data through an
462+
LLM-powered chat interface. It can be used in Shiny applications, as a
463+
standalone chat client, or in an interactive console.
464+
289465
Examples
290466
--------
467+
**Basic Shiny app:**
291468
```python
292469
from querychat import QueryChat
293470
294471
qc = QueryChat(my_dataframe, "my_data")
295472
qc.app()
296473
```
297474
475+
**Standalone chat client:**
476+
```python
477+
from querychat import QueryChat
478+
import pandas as pd
479+
480+
df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
481+
qc = QueryChat(df, "my_data")
482+
483+
# Get a chat client with all tools
484+
client = qc.client()
485+
response = client.chat("What's the average of column a?")
486+
487+
# Start an interactive console chat
488+
qc.console()
489+
```
490+
491+
**Privacy-focused mode:** Only allow dashboard filtering, ensuring the LLM
492+
can't see any raw data.
493+
```python
494+
qc = QueryChat(df, "my_data", tools="update")
495+
qc.app()
496+
```
497+
298498
Parameters
299499
----------
300500
data_source
@@ -324,6 +524,19 @@ class QueryChat(QueryChatBase):
324524
If `client` is not provided, querychat consults the
325525
`QUERYCHAT_CLIENT` environment variable. If that is not set, it
326526
defaults to `"openai"`.
527+
tools
528+
Which querychat tools to include in the chat client by default. Can be:
529+
- A single tool string: `"update"` or `"query"`
530+
- A tuple of tools: `("update", "query")`
531+
- `None` or `()` to disable all tools
532+
533+
Default is `("update", "query")` (both tools enabled).
534+
535+
Set to `"update"` to prevent the LLM from accessing data values, only
536+
allowing dashboard filtering without answering questions.
537+
538+
The tools can be overridden per-client by passing a different `tools`
539+
parameter to the `.client()` method.
327540
data_description
328541
Description of the data in plain text or Markdown. If a pathlib.Path
329542
object is passed, querychat will read the contents of the path into a
@@ -419,7 +632,7 @@ def title():
419632
self.id,
420633
data_source=self._data_source,
421634
greeting=self.greeting,
422-
client=self._client,
635+
client=self.client,
423636
enable_bookmarking=enable_bookmarking,
424637
)
425638

@@ -648,19 +861,6 @@ def title(self, value: Optional[str] = None) -> str | None | bool:
648861
else:
649862
return self._vals.title.set(value)
650863

651-
@property
652-
def client(self):
653-
"""
654-
Get the (session-specific) chat client.
655-
656-
Returns
657-
-------
658-
:
659-
The current chat client.
660-
661-
"""
662-
return self._vals.client
663-
664864

665865
def normalize_data_source(
666866
data_source: IntoFrame | sqlalchemy.Engine | DataSource,
@@ -731,3 +931,20 @@ def assemble_system_prompt(
731931
"extra_instructions": extra_instructions_str,
732932
},
733933
)
934+
935+
936+
def normalize_tools(
937+
tools: TOOL_GROUPS | tuple[TOOL_GROUPS, ...] | None | MISSING_TYPE,
938+
default: tuple[TOOL_GROUPS, ...] | None,
939+
) -> tuple[TOOL_GROUPS, ...] | None:
940+
if tools is None or tools == ():
941+
return None
942+
elif isinstance(tools, MISSING_TYPE):
943+
return default
944+
elif isinstance(tools, str):
945+
return (tools,)
946+
elif isinstance(tools, tuple):
947+
return tools
948+
else:
949+
# Convert any other sequence to tuple
950+
return tuple(tools)

0 commit comments

Comments
 (0)