Add per-message input callbacks

pkg-py/CHANGELOG.md

@@ -18,6 +18,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### New features
 
+* `Chat.transform_user_input()` lets an app add per-message context before content reaches the chat client. The function receives the submitted contents (the arguments for the client's `stream_async()` method) and returns the contents to send, and can be used as a direct decorator or called with arguments. The chat UI and `Chat.user_input()` continue to show the user's original message. (#376)
+
 * You can now edit and resend a message after sending it. Editing forks the conversation from that point — the original branch is kept as a sibling, and `‹ 1 / 2 ›` controls let you switch between versions at any time, including after reloading the page or returning from the history drawer. Requires history to be enabled (the default when using `client=`). (#269)
 
 * Conversation history gained a stable, reactive conversation ID via `chat.history.conversation_id()` (also emitted on OpenTelemetry spans as `gen_ai.conversation.id`, so telemetry can group model work by conversation) and a programmatic `chat.history.save()` for saving the active conversation on demand. (#307, #328)

pkg-py/src/shinychat/_chat.py

@@ -116,8 +116,6 @@
 
 # TODO: UserInput might need to be a list of dicts if we want to support multiple
 # user input content types
-TransformUserInput = Callable[[str], Union[str, None]]
-TransformUserInputAsync = Callable[[str], Awaitable[Union[str, None]]]
 TransformAssistantResponse = Callable[[str], Union[str, HTML, None]]
 TransformAssistantResponseAsync = Callable[
     [str], Awaitable[Union[str, HTML, None]]
@@ -151,6 +149,11 @@
     UserSubmitFunction1,
     UserSubmitFunction2,
 ]
+UserInputContents = list[Any]
+TransformUserInputFn = Callable[
+    [UserInputContents],
+    Union[UserInputContents, None, Awaitable[Union[UserInputContents, None]]],
+]
 
 
 @dataclass(frozen=True)
@@ -347,7 +350,6 @@ def __init__(
         self.id = resolve_id(id)
         self.user_input_id = ResolvedId(f"{self.id}_user_input")
         self._slash_command_id = ResolvedId(f"{self.id}_slash_command")
-        self._transform_user: TransformUserInputAsync | None = None
         self._transform_assistant: (
             TransformAssistantResponseChunkAsync | None
         ) = None
@@ -375,6 +377,7 @@ def __init__(
 
         # Keep track of effects so we can destroy them when the chat is destroyed
         self._effects: list["Effect_"] = []
+        self._transform_user_input_fns: list[TransformUserInputFn] = []
         history_config = (
             history if isinstance(history, HistoryOptions) else None
         )
@@ -537,7 +540,12 @@ def _setup_client(
             async def _on_user_submit(
                 user_input: str, attachments: list[Attachment]
             ) -> None:
-                contents = [attachment_to_content(a) for a in attachments]
+                contents = await self._run_transform_user_input(
+                    [
+                        user_input,
+                        *[attachment_to_content(a) for a in attachments],
+                    ]
+                )
 
                 # Resolve the ID before model work begins: later history
                 # switches, new-chat actions, or client swaps must not
@@ -557,7 +565,6 @@ async def _on_user_submit(
                     client.conversation_id = conversation_id
 
                 response = await chat_client.value.stream_async(
-                    user_input,
                     *contents,
                     content="all",
                     controller=controller,
@@ -687,8 +694,78 @@ async def handle_user_input():
 
         if fn is None:
             return create_effect
-        else:
-            return create_effect(fn)
+
+        return create_effect(fn)
+
+    @overload
+    def transform_user_input(
+        self, fn: TransformUserInputFn
+    ) -> TransformUserInputFn: ...
+
+    @overload
+    def transform_user_input(
+        self,
+    ) -> Callable[[TransformUserInputFn], TransformUserInputFn]: ...
+
+    def transform_user_input(
+        self, fn: TransformUserInputFn | None = None
+    ) -> (
+        TransformUserInputFn
+        | Callable[[TransformUserInputFn], TransformUserInputFn]
+    ):
+        """
+        Add a function that transforms the contents sent to a chat client.
+
+        The function receives a list of arguments for the client's ``stream_async()``
+        method and returns the list that the client receives. The list starts with the
+        submitted text; chatlas content objects converted from attachments follow it.
+        Add per-message context before the submitted contents. The chat UI and
+        :meth:`user_input` continue to show the user's original message.
+
+        This method only affects automatic client handling from ``Chat(client=...)``.
+        It does not apply to slash commands: a command's handler owns the
+        transformation for its submissions. Transforms run in registration order.
+        Register transforms before a user submits a message, either directly or
+        with arguments::
+
+            @chat.transform_user_input
+            def add_context(contents):
+                context = retrieve_context(contents)
+                return [context, *contents]
+
+
+            @chat.transform_user_input()
+            def add_instruction(contents):
+                return [*contents, "instruction"]
+
+        A function that returns ``None`` is skipped with a warning; return the
+        modified contents instead.
+        """
+
+        def register(fn: TransformUserInputFn) -> TransformUserInputFn:
+            self._transform_user_input_fns.append(fn)
+            return fn
+
+        if fn is None:
+            return register
+        return register(fn)
+
+    async def _run_transform_user_input(
+        self, contents: UserInputContents
+    ) -> UserInputContents:
+        for fn in self._transform_user_input_fns:
+            result = cast(
+                UserInputContents | None, await _utils.wrap_async(fn)(contents)
+            )
+            if result is None:
+                warnings.warn(
+                    "A `transform_user_input` function returned None; contents are "
+                    "unchanged. Did you forget to return the modified contents?",
+                    stacklevel=2,
+                )
+                continue
+            contents = result
+        return contents
 
     @overload
     def slash_command(
@@ -1579,13 +1656,6 @@ async def _restore_bookmark_message(self, message_dict: Any) -> None:
         self._store_message(stored)
         await self._send_append_message(stored)
 
-    def transform_user_input(self, *args: object, **kwargs: object) -> object:
-        raise TypeError(
-            "`.transform_user_input()` has been removed. "
-            "Instead, transform user input manually before passing it to your "
-            "LLM provider (e.g., chatlas, LangChain)."
-        )
-
     @overload
     def transform_assistant_response(
         self, fn: TransformAssistantResponseFunction

pkg-py/tests/pytest/test_chat.py

@@ -16,7 +16,6 @@
 from shinychat._chat_normalize import message_content, message_content_chunk
 from shinychat._chat_types import (
     ChatMessage,
-    ChatMessageDict,
     Role,
     StoredMessage,
     StoredSegment,
@@ -107,14 +106,6 @@ def test_tokenizer_raises():
             Chat(id="chat", tokenizer=object())  # type: ignore[arg-type]
 
 
-def test_transform_user_input_raises():
-    with session_context(test_session):
-        chat = Chat(id="chat")
-
-        with pytest.raises(TypeError, match="transform_user_input.*removed"):
-            chat.transform_user_input(lambda x: x)
-
-
 def test_stream_replace_discards_stale_html_dependencies():
     with session_context(test_session):
         chat = Chat(id="chat")
@@ -950,9 +941,7 @@ async def _exercise() -> None:
             "type": "message",
             "message": {
                 "role": "assistant",
-                "segments": [
-                    {"content": "queued", "content_type": "markdown"}
-                ],
+                "segments": [{"content": "queued", "content_type": "markdown"}],
             },
         }
 
@@ -1367,7 +1356,7 @@ def test_messages_surfaces_attachments():
 
         # First message: assistant with attachment. No `format=` was passed, so
         # messages() returns ChatMessageDict entries.
-        att_msg = cast(ChatMessageDict, msgs[0])
+        att_msg = msgs[0]
         assert "attachments" in att_msg
         atts = att_msg["attachments"]
         assert len(atts) == 1

pkg-py/tests/pytest/test_chat_auto.py

@@ -139,6 +139,58 @@ def test_client_value_returns_raw_client():
     assert chat.client.value is mock
 
 
+def test_transform_user_input_registers_callback():
+    with session_context(test_session):
+        chat = Chat("test_transform_user_input")
+
+    def add_context(contents: list[Any]) -> list[Any]:
+        return ["context", *contents]
+
+    assert chat.transform_user_input(add_context) is add_context
+    assert chat._transform_user_input_fns == [add_context]
+
+
+def test_transform_user_input_callbacks_transform_contents_in_order():
+    with session_context(test_session):
+        chat = Chat("test_transform_user_input_order")
+
+    @chat.transform_user_input
+    def add_context(contents: list[Any]) -> list[Any]:
+        return ["context", *contents]
+
+    @chat.transform_user_input()
+    async def add_instruction(contents: list[Any]) -> list[Any]:
+        return [*contents, "instruction"]
+
+    async def run() -> None:
+        assert await chat._run_transform_user_input(["question"]) == [
+            "context",
+            "question",
+            "instruction",
+        ]
+
+    _run_async(run)
+
+
+def test_transform_user_input_keeps_contents_when_callback_returns_none():
+    with session_context(test_session):
+        chat = Chat("test_transform_user_input_none")
+
+    @chat.transform_user_input
+    def forget_to_return(contents: list[Any]) -> None:
+        pass
+
+    async def run() -> None:
+        with pytest.warns(
+            UserWarning, match="transform_user_input.*returned None"
+        ):
+            assert await chat._run_transform_user_input(["question"]) == [
+                "question"
+            ]
+
+    _run_async(run)
+
+
 # ---------------------------------------------------------------------------
 # ChatClient._swap_client — sync / no-sync
 # ---------------------------------------------------------------------------

pkg-r/NEWS.md

@@ -13,6 +13,8 @@
 
 ## New features and improvements
 
+* `chat_server()` now returns `transform_user_input()`, which lets an app add per-message context before content reaches the chat client. The callback receives the submitted contents and returns the contents to send, while the chat UI and `last_input()` continue to show the user's original message. (#376)
+
 * `chat_server()` gets multi-conversation history automatically: a drawer for starting new chats and returning to previous ones, with LLM-generated titles, search, rename, and delete. Conversations are persisted per-user (or a custom scope) via a pluggable store (the default `FileConversationStore` works out of the box, including on Posit Connect). Customize with `history = history_options(...)` — e.g. how the active conversation is restored across reloads (browser storage, URL, or Shiny bookmarking) and callbacks to keep app state synced to it — or opt out with `history = FALSE`. The history object also offers programmatic control, including a reactive conversation ID and an explicit `save()` method. For apps that can't use `chat_server()`, wire it up manually with `chat_enable_history()`. (#266, #307, #328)
 
 * You can now edit and resend a message after sending it. Editing forks the conversation from that point — the original branch is kept as a sibling, and `‹ 1 / 2 ›` controls let you switch between versions at any time, including after reloading the page or returning from the history drawer. Requires history to be enabled (the default with `chat_server()`). (#269)

pkg-r/R/chat_app.R

@@ -90,6 +90,19 @@
 #'       when attachments are disabled, a list of ellmer `Content` objects when
 #'       enabled).
 #'     * `last_turn`: A reactive value containing the last assistant turn.
+#'     * `transform_user_input(fn)`: Add a callback that changes the contents
+#'       sent to the chat client. `fn` receives the submitted contents and must
+#'       return the contents to send. Use it to add per-message context without
+#'       changing the message in the chat UI or `last_input`. Callbacks run in
+#'       registration order. Slash commands are not transformed: the command's
+#'       handler owns the transformation for its submissions.
+#'
+#'       ```r
+#'       chat$transform_user_input(function(contents) {
+#'         context <- retrieve_context(contents)
+#'         c(list(context), contents)
+#'       })
+#'       ```
 #'     * `update_user_input()`: A function to update the chat input or submit a
 #'       new user input. Takes the same arguments as [update_chat_user_input()],
 #'       except for `id` and `session`, which are supplied automatically.
@@ -150,7 +163,9 @@
 #'       the `shiny:chat-slash-command` DOM event. A handler that takes one
 #'       argument receives a [ContentSlashCommand] object (not a plain string).
 #'       See [ContentSlashCommand] for details on how to use this object to
-#'       preserve the original command text across bookmarks. `echo` controls
+#'       preserve the original command text across bookmarks. Slash command
+#'       submissions are not sent through `transform_user_input()`: the handler
+#'       owns the transformation for its command. `echo` controls
 #'       whether invoking the command is echoed as a user message and awaits a
 #'       response; it defaults to `TRUE` when a handler is given and `FALSE`
 #'       otherwise (set `echo = FALSE` for a handler that only performs side
@@ -475,6 +490,7 @@ chat_server <- function(
     }
   )
 
+  transform_user_input_fns <- list()
   saved_on_save_fns <- list()
   saved_on_restore_fns <- list()
 
@@ -604,6 +620,10 @@ chat_server <- function(
     {
       user_input <- session$input[[paste0(id, "_user_input")]]
       last_input(user_input)
+      contents <- user_input
+      for (fn in transform_user_input_fns) {
+        contents <- call_transform_user_input(fn, contents)
+      }
 
       # Resolve the active conversation ID before model work begins and set
       # it on the client as a scalar: later history switches, new-chat
@@ -636,7 +656,7 @@ chat_server <- function(
       append_stream_task$invoke(
         client,
         id,
-        user_input,
+        contents,
         controller = ctrl
       )
     }
@@ -958,6 +978,10 @@ chat_server <- function(
   ret$set_greeting <- set_greeting_mod
   ret$set_client <- set_client
   ret$slash_command <- slash_command_method
+  ret$transform_user_input <- function(fn) {
+    transform_user_input_fns <<- c(transform_user_input_fns, list(fn))
+    invisible(fn)
+  }
 
   hist_env <- new.env(parent = emptyenv())
 
@@ -1011,6 +1035,17 @@ chat_server <- function(
   ret
 }
 
+call_transform_user_input <- function(fn, contents) {
+  result <- fn(contents)
+  if (is.null(result)) {
+    rlang::warn(
+      "A `transform_user_input` function returned NULL; contents are unchanged. Did you forget to return the modified contents?"
+    )
+    return(contents)
+  }
+  result
+}
+
 #' @describeIn chat_mod_ui A Shiny module server for chat (deprecated). Use [chat_server()] instead.
 #' @export
 chat_mod_server <- function(