Client needs to be session-specific

Author: cpsievert

pkg-py/src/querychat/_querychat.py

@@ -136,10 +136,11 @@ def __init__(
             prompt_template=prompt_template,
         )
 
+        # Fork and empty chat now so the per-session forks are fast
         client = normalize_client(client)
-        self.client = copy.deepcopy(client)
-        self.client.set_turns([])
-        self.client.system_prompt = prompt
+        self._client = copy.deepcopy(client)
+        self._client.set_turns([])
+        self._client.system_prompt = prompt
 
         # Populated when ._server() gets called (in an active session)
         self._server_values: ModServerResult | None = None
@@ -313,7 +314,7 @@ def _server(self, *, enable_bookmarking: bool = False) -> None:
             self.id,
             data_source=self._data_source,
             greeting=self.greeting,
-            client=self.client,
+            client=self._client,
             enable_bookmarking=enable_bookmarking,
         )
 
@@ -444,11 +445,27 @@ def generate_greeting(self, *, echo: Literal["none", "output"] = "none"):
             The greeting string (in Markdown format).
 
         """
-        client = copy.deepcopy(self.client)
+        client = copy.deepcopy(self._client)
         client.set_turns([])
         prompt = "Please give me a friendly greeting. Include a few sample prompts in a two-level bulleted list."
         return str(client.chat(prompt, echo=echo))
 
+    @property
+    def client(self):
+        """
+        Get the (session-specific) chat client.
+
+        Returns
+        -------
+        :
+            The current chat client.
+
+        """
+        vals = self._server_values
+        if vals is None:
+            raise RuntimeError("Must call .server() before accessing .client")
+        return vals.client
+
     @property
     def data_source(self):
         """

pkg-py/src/querychat/_querychat_module.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import copy
 from dataclasses import dataclass
 from pathlib import Path
 from typing import TYPE_CHECKING, Callable, Union
@@ -47,6 +48,7 @@ class ModServerResult:
     df: Callable[[], pd.DataFrame]
     sql: ReactiveString
     title: ReactiveStringOrNone
+    client: chatlas.Chat
 
 
 @module.server
@@ -65,15 +67,18 @@ def mod_server(
     title = ReactiveStringOrNone(None)
     has_greeted = reactive.value[bool](False)  # noqa: FBT003
 
+    # Set up the chat object for this session
+    chat = copy.deepcopy(client)
+
     # Create the tool functions
     update_dashboard_tool = tool_update_dashboard(data_source, sql, title)
     reset_dashboard_tool = tool_reset_dashboard(sql, title)
     query_tool = tool_query(data_source)
 
     # Register tools with annotations for the UI
-    client.register_tool(update_dashboard_tool)
-    client.register_tool(query_tool)
-    client.register_tool(reset_dashboard_tool)
+    chat.register_tool(update_dashboard_tool)
+    chat.register_tool(query_tool)
+    chat.register_tool(reset_dashboard_tool)
 
     # Execute query when SQL changes
     @reactive.calc
@@ -89,7 +94,7 @@ def filtered_df():
     # Handle user input
     @chat_ui.on_user_submit
     async def _(user_input: str):
-        stream = await client.stream_async(user_input, echo="none", content="all")
+        stream = await chat.stream_async(user_input, echo="none", content="all")
         await chat_ui.append_message_stream(stream)
 
     @reactive.effect
@@ -100,7 +105,7 @@ async def greet_on_startup():
         if greeting:
             await chat_ui.append_message(greeting)
         elif greeting is None:
-            stream = await client.stream_async(
+            stream = await chat.stream_async(
                 "Please give me a friendly greeting. Include a few sample prompts in a two-level bulleted list.",
                 echo="none",
             )
@@ -145,4 +150,4 @@ def _on_restore(x: RestoreState) -> None:
             if "querychat_has_greeted" in vals:
                 has_greeted.set(vals["querychat_has_greeted"])
 
-    return ModServerResult(df=filtered_df, sql=sql, title=title)
+    return ModServerResult(df=filtered_df, sql=sql, title=title, client=chat)