diff --git a/pkg-py/CHANGELOG.md b/pkg-py/CHANGELOG.md index e8be2fdc..8bca2840 100644 --- a/pkg-py/CHANGELOG.md +++ b/pkg-py/CHANGELOG.md @@ -31,7 +31,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 qc.server(data_source=conn.table("my_table"), client=chat_client) ``` - Registering a table this way is no longer blocked by an earlier session having already registered one, and no longer tears down a still-in-use data source from an earlier session (replacing it via `server(data_source=)` leaves the replaced source's cleanup to whoever created it). + Registering a table this way is no longer blocked by an earlier session having already registered one, no longer tears down a still-in-use data source from an earlier session (replacing it via `server(data_source=)` leaves the replaced source's cleanup to whoever created it), and each session's auto-generated greeting reflects its own table even if a later session registers a different one before that greeting is generated. ### Improvements diff --git a/pkg-py/src/querychat/_querychat_base.py b/pkg-py/src/querychat/_querychat_base.py index b0a7fb8a..a1433f6a 100644 --- a/pkg-py/src/querychat/_querychat_base.py +++ b/pkg-py/src/querychat/_querychat_base.py @@ -383,10 +383,14 @@ def client_factory( tables: list[str], prompt: str | Path, base: chatlas.Chat | None = None, + *, + data_sources: dict[str, DataSource] | None = None, ) -> chatlas.Chat: sp = QueryChatSystemPrompt( prompt_template=prompt, - data_sources=self._data_sources, + data_sources=( + self._data_sources if data_sources is None else data_sources + ), data_description=self._data_description, extra_instructions=None, categorical_threshold=self._categorical_threshold, @@ -510,20 +514,16 @@ def _add_or_replace_table( """ Stage a table and rebuild the system prompt/executor cache. - This is the guard-free core of :meth:`add_table`. It's also called - directly by ``.server(data_source=...)`` so that each session can - register (or replace) its own table even after an earlier session's - ``.server()`` call has already set ``_server_initialized``. - - ``cleanup_replaced=False`` must be used for that per-session - replacement: a table replaced here may still be in active use by an - earlier, already-running session (e.g. its own - ``DataSourceExecutor`` holds a live reference to it), so closing/ - disposing it here would pull the resource out from under that - session. Cleaning it up is then the caller's own responsibility - (e.g. via ``session.on_ended()`` in the code that created it). The - default (``True``) preserves :meth:`add_table`'s existing behavior, - where a config-time replacement has exactly one owner. + Guard-free core of :meth:`add_table`, also called directly by + ``.server(data_source=...)`` so each session can register its own + table even after an earlier session's ``.server()`` call has set + ``_server_initialized``. + + ``cleanup_replaced=False`` is for that per-session path: the + replaced table may still be in active use by an earlier, + still-running session, so cleaning it up here would pull the + resource out from under it. Cleanup becomes the caller's + responsibility (e.g. via ``session.on_ended()``). """ if not isinstance(include_in_greeting, bool): raise TypeError( diff --git a/pkg-py/src/querychat/_querychat_greeter.py b/pkg-py/src/querychat/_querychat_greeter.py index ef95551e..f5d4422d 100644 --- a/pkg-py/src/querychat/_querychat_greeter.py +++ b/pkg-py/src/querychat/_querychat_greeter.py @@ -12,6 +12,8 @@ import chatlas + from ._datasource import DataSource + class QueryChatGreeter: """Controls greeting generation for a QueryChat instance. Access via ``qc.greeter``.""" @@ -68,3 +70,27 @@ async def generate_async(self, *, base: chatlas.Chat | None = None): """Stream a greeting response from the greeting client.""" client = self.build_client(base) return await client.stream_async(GREETING_PROMPT, echo="none") + + async def _generate_async_snapshot( + self, + *, + base: chatlas.Chat | None, + tables: list[str] | None, + data_sources: dict[str, DataSource], + ): + """ + Stream a greeting response from an explicit session snapshot. + + Internal counterpart to :meth:`generate_async`, used by + ``mod_server()``. The snapshot matters because greeting generation + is scheduled lazily: by the time it runs, a later session's + ``.server(data_source=...)`` call may have already mutated the + shared live state. + """ + client = self._client_factory( + self._tables if tables is None else tables, + self._prompt, + base, + data_sources=data_sources, + ) + return await client.stream_async(GREETING_PROMPT, echo="none") diff --git a/pkg-py/src/querychat/_shiny.py b/pkg-py/src/querychat/_shiny.py index 0e2d1b6e..f394b5ac 100644 --- a/pkg-py/src/querychat/_shiny.py +++ b/pkg-py/src/querychat/_shiny.py @@ -420,6 +420,7 @@ def app_server(input: Inputs, output: Outputs, session: Session): tools=self.tools, greeter=self.greeter, greeting_base=None, + greeting_tables=list(self.greeter.tables), ) @reactive.calc @@ -768,6 +769,7 @@ def create_session_client(**kwargs) -> chatlas.Chat: tools=self.tools, greeter=self.greeter, greeting_base=resolved_client, + greeting_tables=list(self.greeter.tables), ) @@ -1055,6 +1057,7 @@ def _ensure_server_started(self) -> None: tools=self.tools, greeter=self.greeter, greeting_base=None, + greeting_tables=list(self.greeter.tables), ) def sidebar( diff --git a/pkg-py/src/querychat/_shiny_module.py b/pkg-py/src/querychat/_shiny_module.py index 16f2d963..b238b419 100644 --- a/pkg-py/src/querychat/_shiny_module.py +++ b/pkg-py/src/querychat/_shiny_module.py @@ -248,6 +248,7 @@ def mod_server( tools: set[str] | None = None, greeter: QueryChatGreeter, greeting_base: chatlas.Chat | None = None, + greeting_tables: list[str] | None = None, ) -> ServerValues[IntoFrameT]: if not callable(client): raise TypeError("mod_server() requires a callable client factory.") @@ -344,7 +345,11 @@ async def _make_greeting(): GreetWarning, stacklevel=1, ) - stream = await greeter.generate_async(base=greeting_base) + stream = await greeter._generate_async_snapshot( + base=greeting_base, + tables=greeting_tables, + data_sources=data_sources, + ) return shinychat.chat_greeting(stream, persistent=True) greeting_arg = ( diff --git a/pkg-py/tests/test_querychat.py b/pkg-py/tests/test_querychat.py index 5dcbdcb8..6f3e356c 100644 --- a/pkg-py/tests/test_querychat.py +++ b/pkg-py/tests/test_querychat.py @@ -1,14 +1,21 @@ +import asyncio import os import tempfile from pathlib import Path from unittest.mock import patch import ibis +import narwhals.stable.v1 as nw import pandas as pd import polars as pl import pytest from querychat import QueryChat -from querychat._datasource import IbisSource, PolarsLazySource +from querychat._datasource import ( + DataFrameSource, + DataSource, + IbisSource, + PolarsLazySource, +) from sqlalchemy import create_engine, text @@ -429,3 +436,62 @@ def test_remove_table_prunes_greeter_tables(sqlite_engine): qc.remove_table("orders") assert "orders" not in qc.greeter.tables assert "customers" in qc.greeter.tables + + +class TestGreeterSnapshotOverrides: + """ + _generate_async_snapshot() renders from an explicit tables/data_sources + snapshot instead of live shared state, which a later Shiny session may + have mutated before an earlier session's async greeting runs. The + public build_client()/generate()/generate_async() API is unaffected. + """ + + def test_build_client_uses_live_state(self, sample_df): + qc = QueryChat(sample_df, "test_table") + prompt = qc.greeter.build_client().system_prompt + assert prompt is not None + assert "test_table" in prompt + + def test_snapshot_tables_override_ignores_live_greeter_tables(self, sample_df): + qc = QueryChat(sample_df, "test_table") + qc.greeter.tables = [] # live state says "no tables" + seen: dict[str, str | None] = {} + + async def fake_stream_async(self, *args, **kwargs): + seen["system_prompt"] = self.system_prompt + return "stream" + + with patch("chatlas.Chat.stream_async", fake_stream_async): + asyncio.run( + qc.greeter._generate_async_snapshot( + base=None, tables=["test_table"], data_sources=qc._data_sources + ) + ) + + assert seen["system_prompt"] is not None + assert "test_table" in seen["system_prompt"] + + def test_snapshot_data_sources_override_ignores_live_data_sources(self, sample_df): + qc = QueryChat(sample_df, "test_table") + other_df = pd.DataFrame({"z": [1, 2, 3]}) + snapshot: dict[str, DataSource] = { + "other_table": DataFrameSource( + nw.from_native(other_df, eager_only=True), "other_table" + ) + } + seen: dict[str, str | None] = {} + + async def fake_stream_async(self, *args, **kwargs): + seen["system_prompt"] = self.system_prompt + return "stream" + + with patch("chatlas.Chat.stream_async", fake_stream_async): + asyncio.run( + qc.greeter._generate_async_snapshot( + base=None, tables=["other_table"], data_sources=snapshot + ) + ) + + assert seen["system_prompt"] is not None + assert "other_table" in seen["system_prompt"] + assert "test_table" not in seen["system_prompt"] diff --git a/pkg-py/tests/test_server_data_source.py b/pkg-py/tests/test_server_data_source.py index 125eb96a..9261d428 100644 --- a/pkg-py/tests/test_server_data_source.py +++ b/pkg-py/tests/test_server_data_source.py @@ -76,11 +76,7 @@ def test_missing_table_name_raises(self, users_df, captured_mod_server): def test_empty_explicit_table_name_raises_instead_of_falling_back( self, users_df, captured_mod_server ): - """ - An explicit but invalid table_name="" must be validated and rejected, - not silently treated as omitted and fall back to the deferred/first - table name. - """ + """An explicit table_name="" must be rejected, not treated as omitted.""" qc = shiny_mod.QueryChat(None, table_name="users") with pytest.raises(ValueError, match="must begin with a letter"): qc.server(data_source=users_df, table_name="") @@ -101,12 +97,6 @@ def test_no_data_source_leaves_tables_unchanged( class TestServerDataSourceSurvivesSecondSession: - """ - server(data_source=...) must not be blocked by an earlier session having - already registered a table -- unlike the public add_table(), which still - guards against changes after server initialization. - """ - def test_second_session_does_not_raise( self, users_df, other_users_df, captured_mod_server ): @@ -136,10 +126,8 @@ def test_second_session_does_not_clean_up_first_sessions_source( self, users_df, other_users_df, captured_mod_server ): """ - A second session's server(data_source=...) call must not tear down - the DataSource object an earlier, still-running session's own - DataSourceExecutor holds a live reference to (e.g. closing a DuckDB - connection or disposing a SQLAlchemy engine out from under it). + An earlier, still-running session's executor holds a live + reference to the source a later session's registration replaces. """ qc = shiny_mod.QueryChat(None, table_name="users") @@ -154,9 +142,8 @@ def test_public_add_table_replace_still_cleans_up_old_source( self, users_df, other_users_df ): """ - Config-time add_table(replace=True) (before any session starts) has - exactly one owner for the replaced table, so its existing - cleanup-on-replace behavior must be unchanged. + Config-time replacement has a single owner, so cleanup-on-replace + is unchanged on the public path. """ qc = shiny_mod.QueryChat(users_df, "users") first_source = qc._data_sources["users"] @@ -169,10 +156,8 @@ def test_second_session_does_not_clean_up_first_sessions_query_executor( self, users_df, other_users_df, captured_mod_server ): """ - A second session's server(data_source=...) call must not close the - cached QueryExecutor an earlier, still-running session's chat has - already captured (e.g. via _create_session_client) and is actively - querying through. + An earlier, still-running session's chat has already captured the + cached executor and may be querying through it. """ qc = shiny_mod.QueryChat(None, table_name="users") @@ -187,9 +172,8 @@ def test_public_add_table_replace_still_cleans_up_old_query_executor( self, users_df, other_users_df ): """ - Config-time add_table(replace=True) has exactly one owner, so its - existing cleanup-on-replace behavior for the cached executor must be - unchanged. + Config-time replacement has a single owner, so executor cleanup + is unchanged on the public path. """ qc = shiny_mod.QueryChat(users_df, "users") first_executor = qc._require_query_executor("test") @@ -197,3 +181,79 @@ def test_public_add_table_replace_still_cleans_up_old_query_executor( with patch.object(first_executor, "cleanup") as mock_cleanup: qc.add_table(other_users_df, "users", replace=True) mock_cleanup.assert_called_once() + + +class TestServerDataSourceGreetingSnapshot: + def test_server_passes_greeting_tables_snapshot_to_mod_server( + self, users_df, captured_mod_server + ): + """ + Greeting generation runs lazily, after a later session may have + mutated the live greeter.tables -- hence the call-time snapshot. + """ + qc = shiny_mod.QueryChat(None, table_name="users") + qc.server(data_source=users_df) + + assert captured_mod_server[0]["greeting_tables"] == ["users"] + + +class TestServerDataSourceMixedWithConfigTimeAddTable: + def test_unnamed_registration_replaces_config_time_table( + self, users_df, other_users_df, captured_mod_server + ): + qc = shiny_mod.QueryChat() + qc.add_table(users_df, "orders") + + qc.server(data_source=other_users_df) + + # Same table name, but the session's data replaces the config-time data + sources = captured_mod_server[0]["data_sources"] + assert list(sources.keys()) == ["orders"] + assert sources["orders"].get_data()["id"].tolist() == [4, 5] + + def test_replacing_config_time_table_does_not_clean_it_up( + self, users_df, other_users_df, captured_mod_server + ): + """ + Consistent with per-session replacement: the replaced source's + cleanup is left to whoever created it. + """ + qc = shiny_mod.QueryChat() + qc.add_table(users_df, "orders") + config_source = qc._data_sources["orders"] + + with patch.object(config_source, "cleanup") as mock_cleanup: + qc.server(data_source=other_users_df) + mock_cleanup.assert_not_called() + + def test_explicit_table_name_adds_alongside_config_time_table( + self, users_df, other_users_df, captured_mod_server + ): + qc = shiny_mod.QueryChat() + qc.add_table(users_df, "orders") + + qc.server(data_source=other_users_df, table_name="returns") + + sources = captured_mod_server[0]["data_sources"] + assert list(sources.keys()) == ["orders", "returns"] + # The config-time table's own data is untouched + assert sources["orders"].get_data()["id"].tolist() == [1, 2, 3] + assert sources["returns"].get_data()["id"].tolist() == [4, 5] + + def test_later_session_snapshot_includes_earlier_sessions_table( + self, users_df, other_users_df, captured_mod_server + ): + """The registry is shared and cumulative across sessions.""" + qc = shiny_mod.QueryChat() + qc.add_table(users_df, "orders") + + # Session 1 adds its own table alongside the config-time one + qc.server(data_source=other_users_df, table_name="returns") + # Session 2 replaces "orders" only -- but still sees session 1's table + third_df = pd.DataFrame({"id": [7, 8, 9]}) + qc.server(data_source=third_df, table_name="orders") + + sources = captured_mod_server[1]["data_sources"] + assert list(sources.keys()) == ["orders", "returns"] + assert sources["orders"].get_data()["id"].tolist() == [7, 8, 9] + assert sources["returns"].get_data()["id"].tolist() == [4, 5] diff --git a/pkg-py/tests/test_shiny_module.py b/pkg-py/tests/test_shiny_module.py index 45b3252d..8bc884f3 100644 --- a/pkg-py/tests/test_shiny_module.py +++ b/pkg-py/tests/test_shiny_module.py @@ -156,6 +156,77 @@ def fake_chat_constructor( assert handoff_server_mock.call_args.kwargs["executor"] is fake_executor +def test_mod_server_generates_greeting_from_session_snapshot_not_live_state(): + """ + _make_greeting() must render from this session's own snapshot, not the + shared live state, which a later session's server(data_source=...) call + may have mutated before an earlier session's greeting streams. + """ + import asyncio + from unittest.mock import AsyncMock, MagicMock, patch + + from querychat._shiny_module import mod_server + + captured = {} + + def fake_chat_constructor( + id, *, client=None, greeting=None, history=None, **kwargs + ): + captured["greeting"] = greeting + return MagicMock() + + fake_source = MagicMock() + fake_source.get_data.return_value = [] + fake_executor = MagicMock() + fake_executor.execute_query.return_value = [] + + def client_factory(**kwargs): + return MagicMock(spec=["stream_async"]) + + fake_greeter = MagicMock() + fake_greeter._generate_async_snapshot = AsyncMock(return_value=MagicMock()) + fake_greeting_base = MagicMock() + + inner_fn = _unwrap_module_server(mod_server) + + fake_input = MagicMock() + fake_input.__getitem__ = MagicMock(return_value=MagicMock()) + fake_session = MagicMock() + fake_session.is_stub_session.return_value = False + + with ( + patch( + "querychat._shiny_module.shinychat.Chat", side_effect=fake_chat_constructor + ), + patch("querychat._shiny_module.has_viz_tool", return_value=False), + patch( + "querychat._shiny_module.shinychat.chat_greeting", return_value=MagicMock() + ), + ): + inner_fn( + fake_input, + MagicMock(), + fake_session, + data_sources={"t": fake_source}, + executor=fake_executor, + greeting=None, + client=client_factory, + history=True, + tools=None, + greeter=fake_greeter, + greeting_base=fake_greeting_base, + greeting_tables=["t"], + ) + + asyncio.run(captured["greeting"]()) + + fake_greeter._generate_async_snapshot.assert_called_once_with( + base=fake_greeting_base, + tables=["t"], + data_sources={"t": fake_source}, + ) + + def test_mod_server_registers_chat_bookmarking_with_no_auto_trigger_when_history_not_bookmark_mode(): """ Chat.enable_bookmarking() is called when `history` isn't bookmark mode, so @@ -348,9 +419,7 @@ def test_shinychat_chat_contract_used_by_mod_server(): mock_session.app = None with session_context(mock_session): - chat = shinychat.Chat( - "chat", client=MagicMock(), greeting=None, history=True - ) + chat = shinychat.Chat("chat", client=MagicMock(), greeting=None, history=True) @chat.history.on_save def _on_save(values):