Move plugin hook management to a class

Author: ncoghlan

src/lmstudio/plugin/__init__.py

@@ -35,7 +35,7 @@
 # * [DONE] add a global plugin config to control the in-place status update demo
 # * handle "Abort" requests from server (including sending "Aborted" responses)
 # * [DONE] catch hook invocation failures and send "Error" responses
-# * [partial] this includes adding runtime checks for the hook returning the wrong type
+# * [DONE] this includes adding runtime checks for the hook returning the wrong type
 #
 # Token generator hook
 # * add an example plugin for this (probably proxying a remote LM Studio instance)

src/lmstudio/plugin/hooks/common.py

@@ -5,13 +5,15 @@
 from random import randrange
 from typing import (
     Any,
+    Awaitable,
     Callable,
     Generic,
     TypeAlias,
     TypeVar,
 )
 
 from ...async_api import AsyncSession
+from ...schemas import DictObject
 from ..._sdk_models import (
     # TODO: Define aliases at schema generation time
     PluginsChannelSetGeneratorToClientPacketGenerate as TokenGenerationRequest,
@@ -43,6 +45,7 @@ class AsyncSessionPlugins(AsyncSession):
 TPluginConfigSchema = TypeVar("TPluginConfigSchema", bound=BaseConfigSchema)
 TGlobalConfigSchema = TypeVar("TGlobalConfigSchema", bound=BaseConfigSchema)
 TConfig = TypeVar("TConfig", bound=BaseConfigSchema)
+SendMessageCallback: TypeAlias = Callable[[DictObject], Awaitable[Any]]
 
 
 class HookController(Generic[TPluginRequest, TPluginConfigSchema, TGlobalConfigSchema]):
@@ -57,6 +60,7 @@ def __init__(
     ) -> None:
         """Initialize common hook controller settings."""
         self.session = session
+        self.request = request
         self.plugin_config = self._parse_config(
             request.plugin_config, plugin_config_schema
         )

src/lmstudio/plugin/hooks/prompt_preprocessor.py

@@ -1,11 +1,12 @@
 """Invoking and supporting prompt preprocessor hook implementations."""
 
+from dataclasses import dataclass
 from traceback import format_tb
-
 from typing import (
     Any,
     Awaitable,
     Callable,
+    Generic,
     Iterable,
     TypeAlias,
     assert_never,
@@ -39,6 +40,7 @@
 from .common import (
     AsyncSessionPlugins,
     HookController,
+    SendMessageCallback,
     StatusBlockController,
     TPluginConfigSchema,
     TGlobalConfigSchema,
@@ -125,6 +127,7 @@ def __init__(
     ) -> None:
         """Initialize prompt preprocessor hook controller."""
         super().__init__(session, request, plugin_config_schema, global_config_schema)
+        self.task_id = request.task_id
         self.pci = request.pci
         self.token = request.token
 
@@ -182,6 +185,124 @@ async def notify_done(self, message: str) -> None:
 ]
 
 
+# TODO: Define a common "PluginHookHandler" base class
+@dataclass()
+class PromptPreprocessor(Generic[TPluginConfigSchema, TGlobalConfigSchema]):
+    """Handle accepting prompt preprocessing requests."""
+
+    plugin_name: str
+    hook_impl: PromptPreprocessorHook
+    plugin_config_schema: type[TPluginConfigSchema]
+    global_config_schema: type[TGlobalConfigSchema]
+
+    def __post_init__(self) -> None:
+        self._logger = logger = new_logger(__name__)
+        logger.update_context(plugin_name=self.plugin_name)
+
+    async def process_requests(
+        self, session: AsyncSessionPlugins, notify_ready: Callable[[], Any]
+    ) -> None:
+        logger = self._logger
+        endpoint = PromptPreprocessingEndpoint()
+        async with session._create_channel(endpoint) as channel:
+            notify_ready()
+            logger.info("Opened channel to receive prompt preprocessing requests...")
+            send_cb = channel.send_message
+            async with create_task_group() as tg:
+                logger.debug("Waiting for prompt preprocessing requests...")
+                async for contents in channel.rx_stream():
+                    logger.debug(
+                        f"Handling prompt preprocessing channel message: {contents}"
+                    )
+                    for event in endpoint.iter_message_events(contents):
+                        logger.debug("Handling prompt preprocessing channel event")
+                        endpoint.handle_rx_event(event)
+                        match event:
+                            case PromptPreprocessingRequestEvent():
+                                logger.debug(
+                                    "Running prompt preprocessing request hook"
+                                )
+                                ctl = PromptPreprocessorController(
+                                    session,
+                                    event.arg,
+                                    self.plugin_config_schema,
+                                    self.global_config_schema,
+                                )
+                                tg.start_soon(self._invoke_hook, ctl, send_cb)
+                    if endpoint.is_finished:
+                        break
+
+    async def _invoke_hook(
+        self,
+        ctl: PromptPreprocessorController[TPluginConfigSchema, TGlobalConfigSchema],
+        send_response: SendMessageCallback,
+    ) -> None:
+        logger = self._logger
+        request = ctl.request
+        message = request.input
+        expected_cls = UserMessage
+        if not isinstance(message, expected_cls):
+            logger.error(
+                f"Received {type(message).__name__!r} ({expected_cls.__name__!r} expected)"
+            )
+            return
+        error_details: SerializedLMSExtendedErrorDict | None = None
+        response_dict: UserMessageDict
+        try:
+            response = await self.hook_impl(ctl, message)
+        except Exception as exc:
+            err_msg = "Error calling prompt preprocessing hook"
+            logger.error(err_msg, exc_info=True, exc=repr(exc))
+            # TODO: Determine if it's worth sending the stack trace to the server
+            ui_cause = f"{err_msg}\n({type(exc).__name__}: {exc})"
+            error_details = SerializedLMSExtendedErrorDict(
+                cause=ui_cause, stack="\n".join(format_tb(exc.__traceback__))
+            )
+        else:
+            if response is None:
+                # No change to message
+                response_dict = message.to_dict()
+            else:
+                logger.debug(
+                    "Validating prompt preprocessing response", response=response
+                )
+                if isinstance(response, dict):
+                    try:
+                        parsed_response = load_struct(response, expected_cls)
+                    except ValidationError as exc:
+                        err_msg = f"Failed to parse prompt preprocessing response as {expected_cls.__name__}\n({exc})"
+                        logger.error(err_msg)
+                        error_details = SerializedLMSExtendedErrorDict(cause=err_msg)
+                    else:
+                        response_dict = parsed_response.to_dict()
+                elif isinstance(response, UserMessage):
+                    response_dict = response.to_dict()
+                else:
+                    err_msg = f"Prompt preprocessing hook returned {type(response).__name__!r} ({expected_cls.__name__!r} expected)"
+                    logger.error(err_msg)
+                    error_details = SerializedLMSExtendedErrorDict(cause=err_msg)
+        channel_message: DictObject
+        if error_details is not None:
+            error_title = f"Prompt preprocessing error in plugin {self.plugin_name!r}"
+            common_error_args: SerializedLMSExtendedErrorDict = {
+                "title": error_title,
+                "rootTitle": error_title,
+            }
+            error_details.update(common_error_args)
+            channel_message = PromptPreprocessingErrorDict(
+                type="error",
+                taskId=request.task_id,
+                error=error_details,
+            )
+        else:
+            channel_message = PromptPreprocessingCompleteDict(
+                type="complete",
+                taskId=request.task_id,
+                processed=response_dict,
+            )
+        await send_response(channel_message)
+
+
 async def run_prompt_preprocessor(
     plugin_name: str,
     hook_impl: PromptPreprocessorHook,
@@ -191,94 +312,7 @@ async def run_prompt_preprocessor(
     notify_ready: Callable[[], Any],
 ) -> None:
     """Accept prompt preprocessing requests."""
-    logger = new_logger(__name__)
-    logger.update_context(plugin_name=plugin_name)
-    endpoint = PromptPreprocessingEndpoint()
-    async with session._create_channel(endpoint) as channel:
-        notify_ready()
-        logger.info("Opened channel to receive prompt preprocessing requests...")
-
-        async def _invoke_hook(request: PromptPreprocessingRequest) -> None:
-            message = request.input
-            expected_cls = UserMessage
-            if not isinstance(message, expected_cls):
-                logger.error(
-                    f"Received {type(message).__name__!r} ({expected_cls.__name__!r} expected)"
-                )
-                return
-            hook_controller = PromptPreprocessorController(
-                session, request, plugin_config_schema, global_config_schema
-            )
-            error_details: SerializedLMSExtendedErrorDict | None = None
-            response_dict: UserMessageDict
-            try:
-                response = await hook_impl(hook_controller, message)
-            except Exception as exc:
-                err_msg = "Error calling prompt preprocessing hook"
-                logger.error(err_msg, exc_info=True, exc=repr(exc))
-                # TODO: Determine if it's worth sending the stack trace to the server
-                ui_cause = f"{err_msg}\n({type(exc).__name__}: {exc})"
-                error_details = SerializedLMSExtendedErrorDict(
-                    cause=ui_cause, stack="\n".join(format_tb(exc.__traceback__))
-                )
-            else:
-                if response is None:
-                    # No change to message
-                    response_dict = message.to_dict()
-                else:
-                    logger.debug(
-                        "Validating prompt preprocessing response", response=response
-                    )
-                    if isinstance(response, dict):
-                        try:
-                            parsed_response = load_struct(response, expected_cls)
-                        except ValidationError as exc:
-                            err_msg = f"Failed to parse prompt preprocessing response as {expected_cls.__name__}\n({exc})"
-                            logger.error(err_msg)
-                            error_details = SerializedLMSExtendedErrorDict(
-                                cause=err_msg
-                            )
-                        else:
-                            response_dict = parsed_response.to_dict()
-                    elif isinstance(response, UserMessage):
-                        response_dict = response.to_dict()
-                    else:
-                        err_msg = f"Prompt preprocessing hook returned {type(response).__name__!r} ({expected_cls.__name__!r} expected)"
-                        logger.error(err_msg)
-                        error_details = SerializedLMSExtendedErrorDict(cause=err_msg)
-            channel_message: DictObject
-            if error_details is not None:
-                error_title = f"Prompt preprocessing error in plugin {plugin_name!r}"
-                common_error_args: SerializedLMSExtendedErrorDict = {
-                    "title": error_title,
-                    "rootTitle": error_title,
-                }
-                error_details.update(common_error_args)
-                channel_message = PromptPreprocessingErrorDict(
-                    type="error",
-                    taskId=request.task_id,
-                    error=error_details,
-                )
-            else:
-                channel_message = PromptPreprocessingCompleteDict(
-                    type="complete",
-                    taskId=request.task_id,
-                    processed=response_dict,
-                )
-            await channel.send_message(channel_message)
-
-        async with create_task_group() as tg:
-            logger.debug("Waiting for prompt preprocessing requests...")
-            async for contents in channel.rx_stream():
-                logger.debug(
-                    f"Handling prompt preprocessing channel message: {contents}"
-                )
-                for event in endpoint.iter_message_events(contents):
-                    logger.debug("Handling prompt preprocessing channel event")
-                    endpoint.handle_rx_event(event)
-                    match event:
-                        case PromptPreprocessingRequestEvent():
-                            logger.debug("Running prompt preprocessing request hook")
-                            tg.start_soon(_invoke_hook, event.arg)
-                if endpoint.is_finished:
-                    break
+    prompt_preprocessor = PromptPreprocessor(
+        plugin_name, hook_impl, plugin_config_schema, global_config_schema
+    )
+    await prompt_preprocessor.process_requests(session, notify_ready)