gRPC Worker Plugins (Python)

1
# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
# SPDX-License-Identifier: Apache-2.0
3
4
"""Example Python gRPC worker plugin package."""

1
# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
# SPDX-License-Identifier: Apache-2.0
3
4
"""Example Python worker plugin using the nemo-relay-plugin SDK."""
5
6
from __future__ import annotations
7
8
from typing import Any
9
10
from nemo_relay_plugin import ConfigDiagnostic, DiagnosticLevel, Json, PluginContext, WorkerPlugin, serve_plugin
11
12
13
class ExamplePythonWorker(WorkerPlugin):
14
    """Small worker plugin that tags tool request JSON and emits a host mark."""
15
16
    plugin_id = "examples.python_grpc_worker"
17
18
    def validate(self, config: Json) -> list[ConfigDiagnostic | dict[str, Any]]:
19
        if not isinstance(config, dict):
20
            return [
21
                ConfigDiagnostic(
22
                    level=DiagnosticLevel.ERROR,
23
                    code="examples.python_grpc_worker.invalid_config",
24
                    component=self.plugin_id,
25
                    message="plugin config must be a JSON object",
26
                )
27
            ]
28
        if config.get("reject") is True:
29
            return [
30
                ConfigDiagnostic(
31
                    level=DiagnosticLevel.ERROR,
32
                    code="examples.python_grpc_worker.rejected",
33
                    component=self.plugin_id,
34
                    field="reject",
35
                    message="Python gRPC worker rejection requested",
36
                )
37
            ]
38
        if "tag" in config and not isinstance(config["tag"], str):
39
            return [
40
                ConfigDiagnostic(
41
                    level=DiagnosticLevel.ERROR,
42
                    code="examples.python_grpc_worker.invalid_tag",
43
                    component=self.plugin_id,
44
                    field="tag",
45
                    message="tag must be a string",
46
                )
47
            ]
48
        return []
49
50
    def register(self, ctx: PluginContext, config: Json) -> None:
51
        if not isinstance(config, dict):
52
            raise TypeError("plugin config must be a JSON object")
53
        if config.get("reject") is True:
54
            raise ValueError("Python gRPC worker rejection requested")
55
        tag = config.get("tag", "python_grpc_worker")
56
        if not isinstance(tag, str):
57
            raise TypeError("tag must be a string")
58
59
        async def tag_tool_request(tool_name: str, args: Json) -> Json:
60
            tagged_args = _tag_json(args, tag)
61
            await ctx.runtime.emit_mark(
62
                "examples.python_grpc_worker.tool_request",
63
                {"tool_name": tool_name, "source": "python-grpc-worker", "tag": tag},
64
            )
65
            return tagged_args
66
67
        ctx.register_tool_request_intercept("tag_tool_request", tag_tool_request)
68
69
70
def _tag_json(value: Json, tag: str) -> Json:
71
    if not isinstance(value, dict):
72
        return value
73
    metadata = value.get("_nemo_relay_plugin")
74
    if metadata is None:
75
        metadata = {}
76
    elif not isinstance(metadata, dict):
77
        return value
78
    return {
79
        **value,
80
        "_nemo_relay_plugin": {**metadata, "tag": tag},
81
    }
82
83
84
async def main() -> None:
85
    """Entrypoint referenced by relay-plugin.toml."""
86
    await serve_plugin(ExamplePythonWorker())
87
88
89
if __name__ == "__main__":
90
    import asyncio
91
92
    asyncio.run(main())