Make Native Chat Handlers Overridable via Entry Points

[Darshan808]

Closes #507

Description:

This PR allows overriding native chat handlers (/ask, /generate, /learn, etc.) through entry points, enabling external extensions to customize behavior.

Changes Implemented:

• Registered native chat handlers (default, /ask, /generate, /learn) via entry points.
• Ensured only the native /ask command receives a Retriever.
• Allowed external handlers to override existing slash commands.
• Improved logging to warn when handlers are overridden.

Why This Change?

As discussed in #398, native chat handlers are hardcoded, making it impossible for external extensions to override them. If a consumer of this plugin does want to override the /ask, /clear, /generate, etc. handlers defined natively here, It is possible now.

Edge Cases Handled:

🔹 If multiple extensions define the same slash command, the last loaded one takes precedence.
🔹 If an external /ask handler wont receive Retriever.
🔹 Invalid or duplicate slash commands trigger warnings/logs.

Testing & Validation:

• Successfully registered and overridden /ask, /generate, and /learn using external extensions.

Next Steps

🔹 Document the process for adding custom chat handlers via entry points.

[Darshan808]

Request for Feedback

Looking for feedback on the approach taken to make native chat handlers overridable via entry points. Are there any improvements or edge cases I might have missed? Also, let me know if there's a better way to handle the /ask retriever logic. Suggestions for implementation improvements are welcome!
CC: @krassowski @dlqqq @Zsailer @3coins

[dlqqq]

@Darshan808 Great work, thank you for contributing this! 🤗 Left some feedback for you below.

Note that I have pretty much a full day of meetings tomorrow, so it's likely that I won't be able to review this again until Thursday. 😓

[dlqqq]

@Darshan808 Thank you for addressing my feedback! ❤️

I gave this PR a lot of thought because while I think this is a good change, I'm concerned about the complexity of the code being introduced. There exists a lot of special handling for /ask, /generate, and /learn. This causes the code to be difficult to read and is prone to subtle bugs. This isn't your fault, this is due to our technical debt. 😅

I recommend that this PR should be paused until we can simplify the way chat handlers are initialized. Ideally, every chat handler should be initialized like:

chat_handlers[command_name] = chat_handler(**chat_handler_kwargs)

I'll open a new issue for this and include guidance on how to do so. Would you be open to fixing this issue first before switching back to this PR? Simplifying how chat commands are initialized will make this PR much cleaner.

[dlqqq]

@Darshan808 Here is the new issue: #1256. Let me know if you'd like to work on it. If so, I'll mark this PR as a draft and move it back to "Active" on the project board.

[Darshan808]

Here is the new issue: #1256. Let me know if you'd like to work on it.

Sure, I'll raise a PR for this. Once it's merged, we can revisit this one.

[dlqqq]

@Darshan808 Thank you so much for helping clean up the code in #1257 and #1268. These changes will make your PR more simple & robust after rebasing & refactoring. Let me know when you're ready, and I can give this PR another review.

I can help release these changes in JAI v2 on Thursday. 👍

[Darshan808]

@dlqqq
I've updated the docs too. Please have a look and let me know if any improvements are needed.

[dlqqq]

@Darshan808 I noticed an issue while testing this PR. Otherwise this looks great, thanks for driving this!

[Darshan808]

Caught the problem!
The issue with the '/' completer is that it retrieves chat handlers from a hard-coded dictionary, meaning it doesn't dynamically update when new handlers are added or existing ones are disabled. As a result, it always displays the same list, regardless of changes.

jupyter-ai/packages/jupyter-ai/jupyter_ai/handlers.py

Lines 35 to 43 in 45ab90f

	# TODO v3: unify loading of chat handlers in a single place, then read
	# from that instead of this hard-coded dict.
	CHAT_HANDLER_DICT = {
	"default": DefaultChatHandler,
	"/ask": AskChatHandler,
	"/learn": LearnChatHandler,
	"/generate": GenerateChatHandler,
	"/help": HelpChatHandler,
	}

[Darshan808]

One way to solve this is by passing ai_extension as a parameter to the handlers, as:

(r"api/ai/chats/autocomplete_options?", AutocompleteOptionsHandler,  
    {"ai_extension": self}  
),

Then, initializing it inside the handler:

def initialize(self, ai_extension, *args, **kwargs) -> None:  
    self.ai_extension = ai_extension  
    super().initialize(*args, **kwargs)  

This way, we can include the room ID in the request and dynamically fetch chat_handlers using:

self.ai_extension.chat_handlers_by_room[room_id]

I've seen this approach used in another extension. What do you think about this ?

[paulrutter]

@Darshan808 following this thread as i want to disable certain commands like learn, ask and generate because we use a custom provider where we currently haven't implemented these features.

Will this PR allow disabling commands as well? Thus removing them from the help as well as the UI? Thanks!

From the updated docs, it seems possible. It would be helpful to provide a code snippet on how to disable this from within a custom provider.

[dlqqq]

@Darshan808 I think there may be a misunderstanding about what disabled = True should do. The feature desired by @krassowski is that a developer should be able to define a new chat handler with disabled = True and its ID set to ask. When this is provided through an entry point, it should ignore this chat handler and disable the default /ask. I believe your example is manually setting disabled = True on Jupyter AI's native AskChatHandler, which other developers cannot do after installing the package.

Please correct me if I'm wrong. Have to drop now, but will give this another look tomorrow. Looking forward to getting this released next week!

[Darshan808]

Let me clarify this: the entry points are loaded, and duplicate handlers are removed here:

        # Override native chat handlers if duplicates are present
        sorted_eps = sorted(
            all_chat_handler_eps, key=lambda ep: ep.dist.name != "jupyter_ai"
        )
        seen = {}
        for ep in sorted_eps:
            seen[ep.name] = ep
        chat_handler_eps = list(seen.values())

This means that the default /ask handler will automatically be removed if another handler with the same name (ask) is present when duplicates are filtered out.

Additionally, when initializing and storing handlers, if a handler has disabled=True, we skip initializing and storing it. I tested this by creating a test extension with id=ask and disabled=True, and it successfully disabled the ask handler. Now, when the /ask slash command is used, it redirects to the default chat handler. However, the UI still displays it. I’d appreciate feedback on how we can remove the hardcoded chat_handlers from auto-completions.

That said, @krassowski @dlqqq I have a question for clarification: Suppose a user disables the ask handler and then installs a custom ask handler. Should the initial disabling of the ask handler also disable the newly installed custom handler? Currently, it does not.

If we want to enforce this behavior, we would need to modify the approach slightly. Let me know your thoughts!

[Darshan808]

Suppose a user disables the ask handler and then installs a custom ask handler.

I think this approach works only for disabling native chat handlers. In the above-mentioned case, whether the custom handler will be disabled or not depends on which extension was loaded first whether disabled ask or custom ask when loading the entry point.

I can think of a way to fix this, but I need your answers on this first: Should the initial disabling of the ask handler also disable the newly installed custom handler?

[gogakoreli]

btw, for disabling slash commands I thought we had this capability in the llm provider using: unsupported_slash_commands = {"/generate"}, or is the scope of this change to be able to disable the specific slash command globally for any of the llms in the dropdown?

[gogakoreli]

I have a suggestion, global command disabling in Jupyter-AI should be controlled centrally rather than by individual extensions to ensure consistent behavior and prevent conflicts.

• Move slash command disabling from extensions to JupyterLab Settings Editor for global control
• Keep provider-level unsupported_slash_commands for granular control within specific LLM providers
• This prevents extension conflicts while maintaining flexibility through proper separation of responsibilities

[dlqqq]

@Darshan808 Thanks for clarifying & being patient with me. I now realize my mistake. To disable a chat command, the entry point name has to exactly match the entry point we provide. This is what I added to jupyter-ai-test to disable /ask:

[project.entry-points."jupyter_ai.chat_handlers"]
ask-disabled = "jupyter_ai_test.test_slash_commands:DisabledAskSlashCommand"

This does not work, since this just adds a ask-disabled key to the chat_handlers dict, instead of overriding the ask key. This is the correct way:

[project.entry-points."jupyter_ai.chat_handlers"]
ask = "jupyter_ai_test.test_slash_commands:DisabledAskSlashCommand"

@Darshan808 We just need to clarify this in the documentation, and this PR is good to go! Thanks for your help thus far. 🤗

Suggested change

	## Disabling a built-in slash command
	You can disable a built-in slash command globally by providing a mostly-empty chat handler with `disabled = True`. For example, to disable the default `ask` chat handler of Jupyter AI, define:
	```python
	class AskChatHandler:
	disabled = True
	```
	## Overriding or disabling a built-in slash command
	You can define a custom implementation of a built-in slash command by following the steps above on building a custom slash command. This will involve creating and installing a new package. Then, to override a chat handler with this custom implementation, provide an entry point with a name matching the ID of the chat handler to override.
	For example, to override `/ask` with a `CustomAskChatHandler` class, add the following to `pyproject.toml` and re-install the new package:
	```python
	[project.entry-points."jupyter_ai.chat_handlers"]
	ask = "<module-path>:CustomAskChatHandler"
	```
	You can also disable a built-in slash command by providing a mostly-empty chat handler with `disabled = True`. For example, to disable the default `ask` chat handler of Jupyter AI, define a new `DisabledAskChatHandler`:
	```python
	class DisabledAskChatHandler:
	id = 'ask'
	disabled = True
	```
	Then, provide this as an entry point in your custom package:
	```python
	[project.entry-points."jupyter_ai.chat_handlers"]
	ask = "<module-path>:DisabledAskChatHandler"
	```
	Finally, re-install your custom package. After starting JupyterLab, the `/ask` command should now be disabled.
	:::{warning}
	:name: entry-point-name
	To override or disable a built-in slash command via an entry point, the name of the entry point (left of the `=` symbol) must match the chat handler ID exactly.
	:::

[dlqqq]

Let me answer everybody's questions:

@paulrutter

Will this PR allow disabling commands as well? Thus removing them from the help as well as the UI?

Yes, that is the goal.

@Darshan808

I’d appreciate feedback on how we can remove the hardcoded chat_handlers from auto-completions.

Don't worry about solving this issue on the main branch, since main is still in v3.0.0 alpha (not ready). We only need to make sure that the commands menu is correct when backporting this PR to 2.x, since that is what users are installing today. The command menu frontend will be refactored to use the new chat commands API that I implemented in Jupyter Chat: jupyterlab/jupyter-chat#161

This work is still pending and I haven't opened issues for this yet. Let me know if you'd like to contribute.

Should the initial disabling of the ask handler also disable the newly installed custom handler?

I don't think we need to define what happens when there are more than 2 providers of the same entry point. Regardless of what strategy we choose, only 1 can be used, so users shouldn't be installing multiple packages that provide the same entry points (with the exception of overriding/disabling Jupyter AI's entry points).

@gogakoreli

is the scope of this change to be able to disable the specific slash command globally for any of the llms in the dropdown?

Yes.

I have a suggestion, global command disabling in Jupyter-AI should be controlled centrally rather than by individual extensions to ensure consistent behavior and prevent conflicts.

I think we need to have a broader discussion about how developers want to disable slash commands first. We have at least 2 different use-cases:

1. Disable a set of commands globally (this PR)
2. Disable a set of commands on specific models (currently unsupported_slash_commands)

Before we try to unify everything, it's important that we first make sure that the unified design will actually cover all the use-cases. If you can help open a new issue for this, we can collaborate there.

[gogakoreli]

@dlqqq on a separate note, quick question, when overriding default slash command for example /ask command, is there a way to specify that I am overriding for the specific llm? The reason Im interested in this feature is because if I override default /ask command and make it more efficient for certain LLM provider, then rest of the providers will suffer. Ideally I would like to have /ask customized for specific LLM provider but use the default /ask command for rest of the available LLM providers. Do you think this is in-scope or out-of-scope for the current PR?

[Darshan808]

The primary goal of this PR as the title suggests, was to make the native chat handler overridable via entry points, which has been successfully addressed. As you mentioned, discussing ways to disable chat handlers or slash commands would be better suited for a separate issue.