Skip to content

Commit 4e996ad

Browse files
authored
Add experimental Python plugin development API (#122)
* Added machinery for developing and executing prompt preprocessing plugins * Added initial skeletons for token generation and tools provision plugins * Started a plugin examples folder * Started a tests subfolder for intentionally invalid plugin definitions
1 parent f0fd18b commit 4e996ad

38 files changed

+1857
-80
lines changed

.github/workflows/test.yml

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,8 @@ jobs:
4141
max-parallel: 8
4242
matrix:
4343
python-version: ["3.10", "3.11", "3.12", "3.13"]
44-
# There's no platform specific SDK code, but explicitly check Windows anyway
44+
# There's no platform specific SDK code, but explicitly check Windows
45+
# to ensure there aren't any inadvertent POSIX-only assumptions
4546
os: [ubuntu-22.04, windows-2022]
4647

4748
# Check https://github.com/actions/action-versions/tree/main/config/actions
Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
# `lmstudio/pydice`
2+
3+
TODO: Example Python tools provider plugin
Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
{
2+
"type": "plugin",
3+
"runner": "python",
4+
"owner": "lmstudio",
5+
"name": "py-dice-tool",
6+
"revision": 2
7+
}
Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
"""Example plugin that provide dice rolling tools."""
2+
3+
# Not yet implemented, currently used to check plugins with no hooks defined
Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
# `lmstudio/pyprompt`
2+
3+
Python prompt preprocessing plugin example
4+
5+
Note: there's no `python` runner in LM Studio yet, so use
6+
`python -m lmstudio.plugin --dev path/to/plugin` to run a dev instance
Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
{
2+
"type": "plugin",
3+
"runner": "python",
4+
"owner": "lmstudio",
5+
"name": "prompt-prefix",
6+
"revision": 1
7+
}
Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
"""Example plugin that adds a prefix to all user prompts."""
2+
3+
import asyncio
4+
5+
from lmstudio.plugin import BaseConfigSchema, PromptPreprocessorController, config_field
6+
from lmstudio import UserMessage, UserMessageDict, TextDataDict
7+
8+
9+
# Assigning ConfigSchema = SomeOtherSchemaClass also works
10+
class ConfigSchema(BaseConfigSchema):
11+
"""The name 'ConfigSchema' implicitly registers this as the per-chat plugin config schema."""
12+
13+
prefix: str = config_field(
14+
label="Prefix to insert",
15+
hint="This text will be inserted at the start of all user prompts",
16+
default="And now for something completely different: ",
17+
)
18+
19+
20+
# Assigning GlobalConfigSchema = SomeOtherGlobalSchemaClass also works
21+
class GlobalConfigSchema(BaseConfigSchema):
22+
"""The name 'GlobalConfigSchema' implicitly registers this as the global plugin config schema."""
23+
24+
enable_inplace_status_demo: bool = config_field(
25+
label="Enable in-place status demo",
26+
hint="The plugin will run an in-place task status updating demo when invoked",
27+
default=True,
28+
)
29+
inplace_status_duration: float = config_field(
30+
label="In-place status total duration (s)",
31+
hint="The number of seconds to spend displaying the in-place task status update",
32+
default=5.0,
33+
)
34+
35+
36+
# Assigning preprocess_prompt = some_other_callable also works
37+
async def preprocess_prompt(
38+
ctl: PromptPreprocessorController[ConfigSchema, GlobalConfigSchema],
39+
message: UserMessage,
40+
) -> UserMessageDict | None:
41+
"""Naming the function 'preprocess_prompt' implicitly registers it."""
42+
if ctl.global_config.enable_inplace_status_demo:
43+
# Run an in-place status prompt update demonstration
44+
status_block = await ctl.notify_start("Starting task (shows a static icon).")
45+
status_updates = (
46+
(status_block.notify_working, "Task in progress (shows a dynamic icon)."),
47+
(status_block.notify_waiting, "Task is blocked (shows a static icon)."),
48+
(status_block.notify_error, "Reporting an error status."),
49+
(status_block.notify_canceled, "Reporting cancellation."),
50+
(
51+
status_block.notify_done,
52+
"In-place status update demonstration completed.",
53+
),
54+
)
55+
status_duration = ctl.global_config.inplace_status_duration / len(
56+
status_updates
57+
)
58+
async with status_block.notify_aborted("Task genuinely cancelled."):
59+
for notification, status_text in status_updates:
60+
await asyncio.sleep(status_duration)
61+
await notification(status_text)
62+
63+
modified_message = message.to_dict()
64+
# Add a prefix to all user messages
65+
prefix_text = ctl.plugin_config.prefix
66+
prefix: TextDataDict = {
67+
"type": "text",
68+
"text": prefix_text,
69+
}
70+
modified_message["content"] = [prefix, *modified_message["content"]]
71+
# Demonstrate simple completion status reporting for non-blocking operations
72+
await ctl.notify_done(f"Added prefix {prefix_text!r} to user message.")
73+
return modified_message
74+
75+
76+
print(f"{__name__} initialized from {__file__}")

sdk-schema/sync-sdk-schema.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -363,6 +363,18 @@ def _infer_schema_unions() -> None:
363363
"LlmChannelPredictCreationParameterDict": "PredictionChannelRequestDict",
364364
"RepositoryChannelDownloadModelCreationParameter": "DownloadModelChannelRequest",
365365
"RepositoryChannelDownloadModelCreationParameterDict": "DownloadModelChannelRequestDict",
366+
# Prettier plugin channel message names
367+
"PluginsChannelSetPromptPreprocessorToClientPacketPreprocess": "PromptPreprocessingRequest",
368+
"PluginsChannelSetPromptPreprocessorToClientPacketPreprocessDict": "PromptPreprocessingRequestDict",
369+
"PluginsChannelSetPromptPreprocessorToServerPacketAborted": "PromptPreprocessingAborted",
370+
"PluginsChannelSetPromptPreprocessorToServerPacketAbortedDict": "PromptPreprocessingAbortedDict",
371+
"PluginsChannelSetPromptPreprocessorToServerPacketComplete": "PromptPreprocessingComplete",
372+
"PluginsChannelSetPromptPreprocessorToServerPacketCompleteDict": "PromptPreprocessingCompleteDict",
373+
"PluginsChannelSetPromptPreprocessorToServerPacketError": "PromptPreprocessingError",
374+
"PluginsChannelSetPromptPreprocessorToServerPacketErrorDict": "PromptPreprocessingErrorDict",
375+
# Prettier config handling type names
376+
"LlmRpcGetLoadConfigReturns": "SerializedKVConfigSettings",
377+
"LlmRpcGetLoadConfigReturnsDict": "SerializedKVConfigSettingsDict",
366378
}
367379

368380

src/lmstudio/_kv_config.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@
4040
LlmSplitStrategy,
4141
LlmStructuredPredictionSetting,
4242
LlmStructuredPredictionSettingDict,
43+
SerializedKVConfigSettings,
4344
)
4445

4546

@@ -324,7 +325,7 @@ def _invert_config_keymap(from_server: FromServerKeymap) -> ToServerKeymap:
324325
)
325326

326327

327-
def dict_from_kvconfig(config: KvConfig) -> DictObject:
328+
def dict_from_kvconfig(config: KvConfig | SerializedKVConfigSettings) -> DictObject:
328329
return {kv.key: kv.value for kv in config.fields}
329330

330331

0 commit comments

Comments
 (0)