@@ -295,6 +295,39 @@ def register(
295
295
message_normalizer_registry = NormalizerRegistry ()
296
296
297
297
298
+ def register_custom_normalizer (
299
+ provider : str , normalizer : BaseMessageNormalizer , force : bool = False
300
+ ) -> None :
301
+ """
302
+ Register a custom normalizer for handling specific message types.
303
+
304
+ Parameters
305
+ ----------
306
+ provider : str
307
+ A unique identifier for this normalizer in the registry
308
+ normalizer : BaseMessageNormalizer
309
+ A normalizer instance that can handle your specific message type
310
+ force : bool, optional
311
+ Whether to override an existing normalizer with the same provider name,
312
+ by default False
313
+
314
+ Examples
315
+ --------
316
+ >>> class MyCustomMessage:
317
+ ... def __init__(self, content):
318
+ ... self.content = content
319
+ ...
320
+ >>> class MyCustomNormalizer(StringNormalizer):
321
+ ... def normalize(self, message):
322
+ ... return ChatMessage(content=message.content, role="assistant")
323
+ ... def can_normalize(self, message):
324
+ ... return isinstance(message, MyCustomMessage)
325
+ ...
326
+ >>> register_custom_normalizer("my_provider", MyCustomNormalizer())
327
+ """
328
+ message_normalizer_registry .register (provider , normalizer , force )
329
+
330
+
298
331
def normalize_message (message : Any ) -> ChatMessage :
299
332
strategies = message_normalizer_registry ._strategies
300
333
for strategy in strategies .values ():
@@ -313,5 +346,5 @@ def normalize_message_chunk(chunk: Any) -> ChatMessage:
313
346
return strategy .normalize_chunk (chunk )
314
347
raise ValueError (
315
348
f"Could not find a normalizer for message chunk of type { type (chunk )} : { chunk } . "
316
- "Consider registering a custom normalizer via shiny.ui._chat_types.registry.register ()"
349
+ "Consider registering a custom normalizer via shiny.ui._chat_normalize.register_custom_normalizer ()"
317
350
)
0 commit comments