[v3-beta] EPIC: Define new API for models & model providers

[dlqqq]

Problem

In Jupyter AI v1 & v2, our team was heavily focused on moving quickly & adding new features that users wanted. This was great, but some of these changes required "hacky" solutions to our model provider classes. These experiences have taught us that we need to explore defining a new API for models & model providers in v3.

This new epic documents multiple issues that are difficult to address separately, since they each require breaking API changes. The main issues are listed below in this description. Subissues on this epic represent other issues reported by users which would be fixed in v3, were this epic to be resolved.

Specifically, the issues with the v2 API are:

1. Providers were intended to abstract a model service, but there may be >1 provider per model service. In other words, there is a confusing mismatch between what a provider is & what a provider should be.

Show details (click to expand)

• Example: OpenAI models are actually served through 2 different providers in Jupyter AI: openai and openai-chat. This is fine for chat, but causes confusion for magic command users, since a user has to type out the full model ID in the %ai magic command.

• A provider should be defined as a class abstraction that represents a model service, a service served either locally (e.g. via Ollama) or remotely (e.g. via Amazon Bedrock). A provider is said to provide one or more models, hence its name.

• However, in v2, a provider is a LangChain class. Specifically, each provider is defined as a class that inherits from 2 classes: the BaseProvider (from jupyter_ai_magics) and either a LLM or ChatModel subclass from LangChain.

• We did this since it seemed simple at the time. This was before LangChain separated LLM and ChatModel classes, which forced duplication of some providers.

• If 2 LangChain classes are offered for the same service (e.g. OpenAI and ChatOpenAI), then we are forced to vendor those models through 2 different providers (e.g. OpenAIProvider and ChatOpenAIProvider).

2. The definition of a provider forces multiple inheritance. This introduces ways for non-breaking changes in LangChain to break Jupyter AI, and has caused issues in the past.

3. Users have no way of updating a model provider's list of models, if it is defined. By contrast, "registry providers" (e.g. ollama) allow a user to type in a model ID, but do not allow users to select a previous model used with a registry provider. We should unify these 2 capabilities in a new experience: let users pick & choose from a list of models, but also let users add new ones in case ours don't work.

4. A provider provides every chat model as a completion model by default. Since we seek to improve code completions by preferring FITM (fill-in-the-middle) models, we should make this configurable to allow changing our defaults without introducing a breaking change.

5. A provider may define the unsupported_slash_commands and persona fields. These were added as quick solutions for v2 users, but these cannot continue to exist in v3. Persona will be a new abstraction, and we will likely want to replace unsupported_slash_commands with configurable "model rules" in the settings UI.

• Personas in v3: [v3-beta] Personas #1313
• TODO: open a new issue about model rules and explain what they are.

6. Providers are hidden from the AI Settings menu when their dependencies are not installed. We should be able to at least show these models anyways & remind the user that they must install that dependency to access that model.

7. All providers are defined in the jupyter_ai_magics package. While leaving it there is fine for users, it is slightly confusing for contributors, since it isn't obvious why models are defined in jupyter_ai_magics and not jupyter_ai. We should explore providing a package which provides a Python API for managing models.

Proposed Solution

WIP.

• Create a new jupyter_ai_models package that provides APIs that allow one to:
  • Fetch the list of models (even if not all model dependencies are installed).
  • Fetch the list of model providers.
  • Add a custom model to a provider, along with user-entered fields.
  • Delete a previously-added custom model.
  • Update a previously-added custom model with new fields.
  • Add model rules to customize certain models, e.g. disabling slash commands on models matching amazon-bedrock:*.
• Identify a way for one provider to specify multiple LangChain classes as clients to the model service. For example, there should be exactly one OpenAI provider, which calls OpenAI or ChatOpenAI from LangChain, based on the model ID.

Additional context

This issue is still a WIP.

[krassowski]

👍 I agree with all points 1-7. On point 7 I will note that IPython reuses jupyter_ai_magics for LLM inline completions but this leads to a circular dependency and makes startup slower than it should be, so jupyter_ai_models would be great indeed.

I think we can summarise a lot of the above by "there is no 1-1 mapping between personas, models, and providers, and all three have n-m relation" (although if a model is provided by multiple providers, not all providers will support the same parameters).

[krassowski]

Slightly related, I opened:

• Performance: move BaseProvider to a separate python module #1337

to highlight the need for the BaseProvider to be defined separately from implementations shipped by default. In terms of this issue I would propose to make sure that jupyter_ai_models is split into jupyter_ai_models.implementations and jupyter_ai_models.protocols or jupyter_ai_models.definitions (I am not set on using this exact naming). While this issue tracks 3.0 work, I would like to make the change proposed in #1337 already in 2.x.