Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Make Native Chat Handlers Overridable via Entry Points #1249

Open
wants to merge 6 commits into
base: main
Choose a base branch
from

Conversation

Darshan808
Copy link
Member

@Darshan808 Darshan808 commented Feb 18, 2025

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
Copy link
Member Author

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

@Darshan808 Darshan808 self-assigned this Feb 18, 2025
@Darshan808 Darshan808 added the enhancement New feature or request label Feb 18, 2025
Copy link
Member

@dlqqq dlqqq left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@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. 😓

@@ -5,10 +5,10 @@
from functools import partial
from typing import Dict

import jupyter_ydoc # must be imported before YChat
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Did you need this import for the branch to work? It's fine to import this package, but I'm curious as to why.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would guess this would be related to:

# The following import is to make sure jupyter_ydoc is imported before
# jupyterlab_chat, otherwise it leads to circular import because of the
# YChat relying on YBaseDoc, and jupyter_ydoc registering YChat from the entry point.
import jupyter_ydoc

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would guess this would be related to:

Yes it is!
I was getting circular import error after I moved the import of Retriever class from top level to below. When I looked into the codebase I saw the above comment explaining the error

Comment on lines 537 to 538
# Lazy import Retriever
from jupyter_ai.chat_handlers.learn import Retriever
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can this be moved to a top-level import? I don't see an obvious benefit to doing the import here.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since this is one of the heavy imports, deployments where speed really matters could disable this entry point and avoid triggering this import, right?

Comment on lines +46 to +52
[project.entry-points."jupyter_ai.chat_handlers"]
default = "jupyter_ai.chat_handlers.default:DefaultChatHandler"
ask = "jupyter_ai.chat_handlers.ask:AskChatHandler"
generate = "jupyter_ai.chat_handlers.generate:GenerateChatHandler"
learn = "jupyter_ai.chat_handlers.learn:LearnChatHandler"
help = "jupyter_ai.chat_handlers.help:HelpChatHandler"

Copy link
Member

@dlqqq dlqqq Feb 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think that providing our own entry points to "jupyter_ai.chat_handlers" may not allow other extensions to override our entry points. Suppose another server extension wants to override Jupyter AI's /help command. Then their pyproject.toml file contains:

[project.entry-points."jupyter_ai.chat_handlers"]
help = "example_extension.chat_handlers:CustomHelpChatHandler"

In the current implementation, the chat handler used for /help is defined by the order the entry points are loaded in, because both are setting the same key in the chat_handlers dictionary.

For example, if Jupyter AI gets loaded first, then example_extension overrides /help with their custom implementation (which is desired). However, if example_extension is loaded first, then the custom /help is overridden by the Jupyter AI's default /help (not desired).

  1. Can you find documentation on what order entry points are loaded in if multiple packages are providing the same entry point?
  2. Can you make sure that other extensions can always override our entry points?

Feel free to add an example in the packages/jupyter-ai-test package, which is a local Python package used to verify that entry points work.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great point. As per https://packaging.python.org/en/latest/specifications/entry-points/#entry-points

Within a distribution, entry point names should be unique. If different distributions provide the same name, the consumer decides how to handle such conflicts.

My understanding is that we need to loop the output of entry_points() (chat_handler_eps) twice or thrice:

  • sort so that jupyter-ai is first
  • reduce so that duplicates later in the list win
  • only then load the entry points

I think we would ideally have a way to specify an entry point which will skip adding a chat handler (say to disable /ask if downstream has its own RAG which is not compatible with Retriver class). This could already be done by providing one that will error out but this is of course suboptimal.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Indeed, this will be useful. Should we let extension authors add a 'disabled: true' field in the entry point classes to skip adding the entry point?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Allow developers to override built-in chat commands
3 participants