Skip to content

Commit 1667062

Browse files
cpsievertclaude
andauthored
feat: Allow deferred chat client initialization (#207)
* feat(py): Allow deferred chat client initialization (#205) When using Posit Connect managed OAuth credentials, chat client connections need access to HTTP headers in the Shiny session object, requiring creation inside the server function. Changes: - Defer client initialization when data_source=None and client=None - Add chat_client property getter/setter for setting client after init - Add client parameter to server() method for deferred pattern - Add _require_client() method for runtime checks - Update methods (client, console, generate_greeting) to require client Closes #205 Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> * fix(py): require client before Shiny app setup * refactor(py): Defer client initialization and support lazy defaults Client is no longer eagerly normalized at __init__ when not explicitly provided. Instead, _require_client(default=) resolves the client lazily: self._client takes priority, then the caller-provided default, then raises. Methods like .app() and .client() pass default=None (global default), while .server() passes the user's explicit arg (MISSING if omitted, which errors). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * refactor(py): Make normalize_client always return a fresh Chat normalize_client now deepcopies and clears turns internally, so callers no longer need to repeat the copy.deepcopy + set_turns([]) sequence. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * refactor(py): Rename normalize_client to create_client Better reflects that the function always returns a new, independent Chat instance (deepcopied, with turns cleared) rather than just coercing a value into canonical form. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * refactor(py): Rename _require_client to _ensure_client Better reflects that the method resolves and stores the client (side effect) rather than just validating it exists. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * refactor(py): Use create_client for forking in client() and generate_greeting() Eliminates remaining inline copy.deepcopy + set_turns([]) sequences by reusing create_client, which already handles that. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * refactor(py): Replace _client with _client_spec, eliminate state matrix Store the client specification without eager resolution. A single _create_session_client() method resolves spec -> fresh Chat -> system prompt -> tools at point-of-use. - Replace dual-role self._client with self._client_spec (just stores spec) - Delete _ensure_client, replace with _create_session_client - Rename chat_client property to client_spec (no naming collision with client()) - System prompt assigned in exactly one place (_create_session_client) - Remove MISSING sentinel from server() signature - Pass bound method to mod_server instead of lambda/bound-method-as-callable Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * fix(py): Move _require_data_source check to public methods Users calling client() now see "data_source must be set before calling client()" instead of a reference to the internal _create_session_client. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * docs(py): Add caller-guard comment to _create_session_client Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * fix(py): restore deferred client behavior * test(py): lock down shiny session-local client semantics * test(py): refine deferred shiny client semantics tests * fix(py): keep shiny server client overrides session-local * fix(py): clarify server client requirement * fix(py): align server docs and types * test(py): cover deferred shiny client resolution helpers * test(py): verify session-local shiny client fix * style(py): satisfy ruff for deferred client fix * refactor(py): remove _require_client_spec, allow None as valid default Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(py): remove stale docstring about client=None, improve readability Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * test(py): update tests for None-as-valid-default client behavior Remove TestRequireClientSpec class and two tests that asserted None is an invalid client argument, since _require_client_spec no longer exists and None now resolves to the QUERYCHAT_CLIENT env var or "openai". Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * docs(py): revert unnecessary explicit client= in examples and docs Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * style(py): fix ruff formatting in test_deferred_client Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * refactor(py): collapse _create_session_client_from_spec into _create_session_client Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * fix(py): restore _require_data_source check in console() for error clarity Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * docs(py): remove unnecessary client= from server() docstring example Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * docs(py): add deferred client feature to changelog Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * refactor(r): defer client resolution, add create_session_client() Store the client spec (NULL/string/Chat) lazily instead of eagerly resolving it in initialize(). Add private create_session_client() as the single resolution point, and delegate $client() and $generate_greeting() to it. This enables constructing QueryChat with NULL data_source and no API key, deferring all resolution until the client is actually needed. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * feat(r): add client parameter to $server() for session-local overrides The client parameter allows per-session chat client overrides without mutating the shared client spec. This enables Posit Connect managed OAuth credentials that require session scope. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * style(r): format deferred client changes with air Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * fix(r): clone Chat in create_session_client() to prevent mutation When the client spec is a Chat object, as_querychat_client() returns it as-is. Without cloning, subsequent calls to $client() would mutate the same object (duplicate tool registrations, shared turn history). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * fix(py): keep deferred client resolution lazy * `air format` (GitHub Actions) * `devtools::document()` (GitHub Actions) * docs(r): add deferred client feature to NEWS.md Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * docs(r): mention deferred client in build vignette Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com> Co-authored-by: cpsievert <cpsievert@users.noreply.github.com>
1 parent ac25ab5 commit 1667062

15 files changed

Lines changed: 537 additions & 113 deletions

pkg-py/CHANGELOG.md

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

88
## [Unreleased]
99

10+
### New features
11+
12+
* `QueryChat()` now supports deferred chat client initialization. Pass `client=` to `server()` to provide a session-scoped chat client, enabling use cases where API credentials are only available at session time (e.g., Posit Connect managed OAuth tokens). When no `client` is specified anywhere, querychat resolves a sensible default from the `QUERYCHAT_CLIENT` environment variable (or `"openai"`). (#205)
13+
1014
### Improvements
1115

1216
* When a custom `prompt_template` is provided that doesn't contain Mustache references to `{{schema}}`, the expensive `get_schema()` call is now skipped entirely. This allows users with large databases to avoid slow startup by providing their own prompt that includes schema information inline (or omits it). (#208)

pkg-py/docs/build.qmd

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -154,6 +154,8 @@ app = App(app_ui, server)
154154

155155
</details>
156156

157+
If your chat client also depends on session-scoped credentials, you can defer that too by passing it to `qc.server(client=...)` alongside the `data_source`.
158+
157159
:::
158160

159161
:::

pkg-py/src/querychat/_querychat_base.py

Lines changed: 51 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -81,11 +81,7 @@ def __init__(
8181
self._extra_instructions = extra_instructions
8282
self._categorical_threshold = categorical_threshold
8383

84-
# Normalize and initialize client (doesn't need data_source)
85-
client = normalize_client(client)
86-
self._client = copy.deepcopy(client)
87-
self._client.set_turns([])
88-
84+
self._client_spec: str | chatlas.Chat | None = client
8985
self._client_console = None
9086

9187
# Initialize data source (may be None for deferred pattern)
@@ -114,7 +110,6 @@ def _build_system_prompt(self) -> None:
114110
extra_instructions=self._extra_instructions,
115111
categorical_threshold=self._categorical_threshold,
116112
)
117-
self._client.system_prompt = self._system_prompt.render(self.tools)
118113

119114
def _require_data_source(self, method_name: str) -> DataSource[IntoFrameT]:
120115
"""Raise if data_source is not set, otherwise return it for type narrowing."""
@@ -126,6 +121,39 @@ def _require_data_source(self, method_name: str) -> DataSource[IntoFrameT]:
126121
)
127122
return self._data_source
128123

124+
def _create_session_client(
125+
self,
126+
*,
127+
client_spec: str | chatlas.Chat | None | MISSING_TYPE = MISSING,
128+
tools: TOOL_GROUPS | tuple[TOOL_GROUPS, ...] | None | MISSING_TYPE = MISSING,
129+
update_dashboard: Callable[[UpdateDashboardData], None] | None = None,
130+
reset_dashboard: Callable[[], None] | None = None,
131+
) -> chatlas.Chat:
132+
"""Create a fresh, fully-configured Chat."""
133+
spec = self._client_spec if isinstance(client_spec, MISSING_TYPE) else client_spec
134+
chat = create_client(spec)
135+
136+
resolved_tools = normalize_tools(tools, default=self.tools)
137+
138+
if self._system_prompt is not None:
139+
chat.system_prompt = self._system_prompt.render(resolved_tools)
140+
141+
if resolved_tools is None:
142+
return chat
143+
144+
data_source = self._require_data_source("_create_session_client")
145+
146+
if "update" in resolved_tools:
147+
update_fn = update_dashboard or (lambda _: None)
148+
reset_fn = reset_dashboard or (lambda: None)
149+
chat.register_tool(tool_update_dashboard(data_source, update_fn))
150+
chat.register_tool(tool_reset_dashboard(reset_fn))
151+
152+
if "query" in resolved_tools:
153+
chat.register_tool(tool_query(data_source))
154+
155+
return chat
156+
129157
def client(
130158
self,
131159
*,
@@ -151,35 +179,20 @@ def client(
151179
A configured chat client.
152180
153181
"""
154-
data_source = self._require_data_source("client")
155-
if self._system_prompt is None:
156-
raise RuntimeError("System prompt not initialized")
157-
tools = normalize_tools(tools, default=self.tools)
158-
159-
chat = copy.deepcopy(self._client)
160-
chat.set_turns([])
161-
chat.system_prompt = self._system_prompt.render(tools)
162-
163-
if tools is None:
164-
return chat
165-
166-
if "update" in tools:
167-
update_fn = update_dashboard or (lambda _: None)
168-
reset_fn = reset_dashboard or (lambda: None)
169-
chat.register_tool(tool_update_dashboard(data_source, update_fn))
170-
chat.register_tool(tool_reset_dashboard(reset_fn))
171-
172-
if "query" in tools:
173-
chat.register_tool(tool_query(data_source))
174-
175-
return chat
182+
self._require_data_source("client")
183+
return self._create_session_client(
184+
tools=tools,
185+
update_dashboard=update_dashboard,
186+
reset_dashboard=reset_dashboard,
187+
)
176188

177189
def generate_greeting(self, *, echo: Literal["none", "output"] = "none") -> str:
178190
"""Generate a welcome greeting for the chat."""
179191
self._require_data_source("generate_greeting")
180-
client = copy.deepcopy(self._client)
181-
client.set_turns([])
182-
return str(client.chat(GREETING_PROMPT, echo=echo))
192+
chat = create_client(self._client_spec)
193+
if self._system_prompt is not None:
194+
chat.system_prompt = self._system_prompt.render(self.tools)
195+
return str(chat.chat(GREETING_PROMPT, echo=echo))
183196

184197
def console(
185198
self,
@@ -190,8 +203,6 @@ def console(
190203
) -> None:
191204
"""Launch an interactive console chat with the data."""
192205
self._require_data_source("console")
193-
tools = normalize_tools(tools, default=("query",))
194-
195206
if new or self._client_console is None:
196207
self._client_console = self.client(tools=tools, **kwargs)
197208

@@ -262,17 +273,21 @@ def normalize_data_source(
262273
)
263274

264275

265-
def normalize_client(client: str | chatlas.Chat | None) -> chatlas.Chat:
276+
def create_client(client: str | chatlas.Chat | None) -> chatlas.Chat:
277+
"""Resolve a client spec into a fresh Chat with no conversation history."""
266278
if client is None:
267279
client = os.getenv("QUERYCHAT_CLIENT", None)
268280

269281
if client is None:
270282
client = "openai"
271283

272284
if isinstance(client, chatlas.Chat):
273-
return client
285+
chat = copy.deepcopy(client)
286+
else:
287+
chat = chatlas.ChatAuto(provider_model=client)
274288

275-
return chatlas.ChatAuto(provider_model=client)
289+
chat.set_turns([])
290+
return chat
276291

277292

278293
def normalize_tools(

pkg-py/src/querychat/_shiny.py

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
from ._icons import bs_icon
1313
from ._querychat_base import TOOL_GROUPS, QueryChatBase
1414
from ._shiny_module import ServerValues, mod_server, mod_ui
15-
from ._utils import as_narwhals
15+
from ._utils import MISSING, MISSING_TYPE, as_narwhals
1616

1717
if TYPE_CHECKING:
1818
from pathlib import Path
@@ -299,7 +299,7 @@ def app_server(input: Inputs, output: Outputs, session: Session):
299299
self.id,
300300
data_source=data_source,
301301
greeting=self.greeting,
302-
client=self._client,
302+
client=self._create_session_client,
303303
enable_bookmarking=enable_bookmarking,
304304
)
305305

@@ -405,6 +405,7 @@ def server(
405405
self,
406406
*,
407407
data_source: Optional[IntoFrame | sqlalchemy.Engine | ibis.Table] = None,
408+
client: str | chatlas.Chat | MISSING_TYPE = MISSING,
408409
enable_bookmarking: bool = False,
409410
id: Optional[str] = None,
410411
) -> ServerValues[IntoFrameT]:
@@ -422,6 +423,12 @@ def server(
422423
Optional data source to use. If provided, sets the data_source property
423424
before initializing server logic. This is useful for the deferred pattern
424425
where data_source is not known at initialization time.
426+
client
427+
Optional chat client to use for this session. If provided, overrides
428+
any client set at initialization time for this call only. This is useful
429+
for the deferred pattern where the client cannot be created at
430+
initialization time (e.g., when using Posit Connect managed OAuth
431+
credentials that require session access).
425432
enable_bookmarking
426433
Whether to enable bookmarking for the querychat module.
427434
id
@@ -486,12 +493,18 @@ def title():
486493
self.data_source = data_source
487494

488495
resolved_data_source = self._require_data_source("server")
496+
resolved_client_spec = self._client_spec if isinstance(client, MISSING_TYPE) else client
497+
498+
def create_session_client(**kwargs) -> chatlas.Chat:
499+
return self._create_session_client(
500+
client_spec=resolved_client_spec, **kwargs
501+
)
489502

490503
return mod_server(
491504
id or self.id,
492505
data_source=resolved_data_source,
493506
greeting=self.greeting,
494-
client=self.client,
507+
client=create_session_client,
495508
enable_bookmarking=enable_bookmarking,
496509
)
497510

@@ -728,7 +741,7 @@ def __init__(
728741
self.id,
729742
data_source=self._data_source,
730743
greeting=self.greeting,
731-
client=self._client,
744+
client=self._create_session_client,
732745
enable_bookmarking=enable,
733746
)
734747

pkg-py/src/querychat/_shiny_module.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -78,14 +78,16 @@ class ServerValues(Generic[IntoFrameT]):
7878
client
7979
The session-specific chat client instance. This is a deep copy of the
8080
base client configured for this specific session, containing the chat
81-
history and tool registrations for this session only.
81+
history and tool registrations for this session only. This may be
82+
`None` during stub sessions when the client depends on deferred,
83+
session-scoped state.
8284
8385
"""
8486

8587
df: Callable[[], IntoFrameT]
8688
sql: ReactiveStringOrNone
8789
title: ReactiveStringOrNone
88-
client: chatlas.Chat
90+
client: chatlas.Chat | None
8991

9092

9193
@module.server
@@ -115,7 +117,7 @@ def _stub_df():
115117
df=_stub_df,
116118
sql=sql,
117119
title=title,
118-
client=client if isinstance(client, chatlas.Chat) else client(),
120+
client=client if isinstance(client, chatlas.Chat) else None,
119121
)
120122

121123
# Real session requires data_source

pkg-py/tests/test_base.py

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
from querychat._datasource import DataFrameSource, SQLAlchemySource
1111
from querychat._querychat_base import (
1212
QueryChatBase,
13-
normalize_client,
13+
create_client,
1414
normalize_data_source,
1515
normalize_tools,
1616
)
@@ -110,26 +110,26 @@ def test_with_special_column_names(self):
110110

111111
class TestNormalizeClient:
112112
def test_with_none_uses_default(self):
113-
result = normalize_client(None)
113+
result = create_client(None)
114114
assert isinstance(result, chatlas.Chat)
115115

116116
def test_with_string_provider(self):
117-
result = normalize_client("openai")
117+
result = create_client("openai")
118118
assert isinstance(result, chatlas.Chat)
119119

120120
def test_with_chat_instance(self):
121121
chat = chatlas.ChatOpenAI()
122-
result = normalize_client(chat)
122+
result = create_client(chat)
123123
assert isinstance(result, chatlas.Chat)
124124

125125
def test_respects_env_variable(self, monkeypatch):
126126
monkeypatch.setenv("QUERYCHAT_CLIENT", "openai")
127-
result = normalize_client(None)
127+
result = create_client(None)
128128
assert isinstance(result, chatlas.Chat)
129129

130130
def test_with_invalid_provider_raises(self):
131131
with pytest.raises(ValueError, match="is not a known chatlas provider"):
132-
normalize_client("not_a_real_provider_xyz123")
132+
create_client("not_a_real_provider_xyz123")
133133

134134

135135
class TestNormalizeTools:
@@ -207,9 +207,15 @@ def test_client_with_callbacks(self, sample_df):
207207
update_called = []
208208
reset_called = []
209209

210+
def update_dashboard(data):
211+
update_called.append(data)
212+
213+
def reset_dashboard():
214+
reset_called.append(True)
215+
210216
client = qc.client(
211-
update_dashboard=update_called.append,
212-
reset_dashboard=lambda: reset_called.append(True),
217+
update_dashboard=update_dashboard,
218+
reset_dashboard=reset_dashboard,
213219
)
214220
assert isinstance(client, chatlas.Chat)
215221

0 commit comments

Comments
 (0)