Allow embedding model fields, fix coupled model fields, add custom OpenAI provider (#1264)

* Simplifying the OpenAI provider to use multiple model providers

* Update openrouter.md

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* openai general interface added

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* embedding

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Updated settings to take OpenAI generic embedding models

* added openai generic embeddings screenshot

* Fixed Issue 1261

* bump version floor on jupyter server

* linter

* adding embedding model fields

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update test_config_manager

* Update pyproject.toml

* Update pyproject.toml

* Update pyproject.toml

* pyproject.toml fixes

* pyproject.toml updates

* Update pyproject.toml

* Update pyproject.toml

* Make Native Chat Handlers Overridable via Entry Points (#1249)

* make native chat handlers customizable

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* remove-ci-error

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add-disabled-check-and-sort-entrypoints

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Refactor Chat Handlers to Simplify Initialization (#1257)

* simplify-entrypoints-loading

* fix-lint

* fix-tests

* add-retriever-typing

* remove-retriever-from-base

* fix-circular-import(ydoc-import)

* fix-tests

* fix-type-check-failure

* refactor-retriever-init

* Allow chat handlers to be initialized in any order (#1268)

* lazy-initialize-retriever

* add-retriever-property

* rebase-into-main

* update-docs

* update-documentation

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

* pyproject toml files

* pyproject toml updates

* update snapshot

* writing config file correctly

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* tsx lint

* Update use-server-info.ts

* Update pyproject.toml

* adds embedding_models attribute

* Fixed display of Base url for embeddings and completions

* removed embedding_models

* Added help fields

* Update chat-settings.tsx

* minor reversions moved to new issue

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Darshan Poudel <[email protected]>

Author: 3 people

docs/source/_static/openai-chat-deepseek.png

@@ -346,11 +346,13 @@ Jupyter AI enables use of language models hosted on [Amazon Bedrock](https://aws
 For details on enabling model access in your AWS account, using cross-region inference, or invoking custom/provisioned models, please see our dedicated documentation page on [using Amazon Bedrock in Jupyter AI](bedrock.md).
 
 
-### OpenRouter Usage
+### OpenRouter and OpenAI Interface Usage
 
 Jupyter AI enables use of language models accessible through [OpenRouter](https://openrouter.ai)'s unified interface. Examples of models that may be accessed via OpenRouter are: [Deepseek](https://openrouter.ai/deepseek/deepseek-chat), [Qwen](https://openrouter.ai/qwen/), [mistral](https://openrouter.ai/mistralai/), etc. OpenRouter enables usage of any model conforming to the OpenAI API.
 
-For details on enabling model access via the AI Settings and using models via OpenRouter, please see the dedicated documentation page on using [OpenRouter in Jupyter AI](openrouter.md).
+Likewise, for many models, you may directly choose the OpenAI provider in Jupyter AI instead of OpenRouter in the same way.
+
+For details on enabling model access via the AI Settings and using models via OpenRouter or OpenAI, please see the dedicated documentation page on using [OpenRouter and OpenAI providers in Jupyter AI](openrouter.md).
 
 
 ### SageMaker endpoints usage

docs/source/_static/openai-chat-openai.png

@@ -1,4 +1,4 @@
-# Using OpenRouter in Jupyter AI
+# Using OpenRouter or OpenAI Interfaces in Jupyter AI
 
 [(Return to the Chat Interface page)](index.md#openrouter-usage)
 
@@ -33,4 +33,36 @@ You should now be able to use Deepseek! An example of usage is shown next:
     alt='Screenshot of chat using Deepseek via the OpenRouter provider.'
     class="screenshot" />
 
+In a similar manner, models may also be invoked directly using the OpenAI provider interface in Jupyter AI. First, you can choose the OpenAI provider and then enter in the model ID, as shown on the OpenAI [models page](https://platform.openai.com/docs/models). An example is shown below:
+
+<img src="../_static/openai-chat-openai.png"
+    width="75%"
+    alt='Screenshot of chat using gpt-4o via the OpenAI provider.'
+    class="screenshot" />
+
+DeepSeek models may be used via the same interface, if the base API url is provided:
+
+<img src="../_static/openai-chat-deepseek.png"
+    width="75%"
+    alt='Screenshot of chat using deepseek via the OpenAI provider.'
+    class="screenshot" />
+
+For DeepSeek models, enter the DeepSeek API for the OpenAI API key.
+
+Models deployed using vLLM may be used in a similar manner:
+
+<img src="../_static/openai-chat-vllm.png"
+    width="75%"
+    alt='Screenshot of chat using vllm via the OpenAI provider.'
+    class="screenshot" />
+
+Usage of models using vLLM and their deployment is discussed [here](vllm.md).
+
+For embedding models from OpenAI, you can generically choose them using the AI Settings interface as well:
+
+<img src="../_static/openai-embeddings.png"
+    width="75%"
+    alt='Screenshot of embedding use via the OpenAI provider.'
+    class="screenshot" />
+
 [(Return to the Chat Interface page)](index.md#openrouter-usage)

docs/source/_static/openai-chat-vllm.png

@@ -24,12 +24,7 @@ class OllamaEmbeddingsProvider(BaseEmbeddingsProvider, OllamaEmbeddings):
     name = "Ollama"
     # source: https://ollama.com/library
     model_id_key = "model"
-    models = [
-        "nomic-embed-text",
-        "mxbai-embed-large",
-        "all-minilm",
-        "snowflake-arctic-embed",
-    ]
+    models = ["*"]
     registry = True
     fields = [
         TextField(key="base_url", label="Base API URL (optional)", format="text"),

docs/source/_static/openai-embeddings.png

@@ -76,6 +76,27 @@ def is_api_key_exc(cls, e: Exception):
         return False
 
 
+class ChatOpenAICustomProvider(BaseProvider, ChatOpenAI):
+    id = "openai-chat-custom"
+    name = "OpenAI (general interface)"
+    models = ["*"]
+    model_id_key = "model_name"
+    model_id_label = "Model ID"
+    pypi_package_deps = ["langchain_openai"]
+    auth_strategy = EnvAuthStrategy(name="OPENAI_API_KEY")
+    fields = [
+        TextField(
+            key="openai_api_base", label="Base API URL (optional)", format="text"
+        ),
+        TextField(
+            key="openai_organization", label="Organization (optional)", format="text"
+        ),
+        TextField(key="openai_proxy", label="Proxy (optional)", format="text"),
+    ]
+    help = "Supports non-OpenAI model that use the OpenAI API interface. Replace the OpenAI API key with the API key for the chosen provider."
+    registry = True
+
+
 class AzureChatOpenAIProvider(BaseProvider, AzureChatOpenAI):
     id = "azure-chat-openai"
     name = "Azure OpenAI"
@@ -107,6 +128,15 @@ class OpenAIEmbeddingsProvider(BaseEmbeddingsProvider, OpenAIEmbeddings):
     model_id_key = "model"
     pypi_package_deps = ["langchain_openai"]
     auth_strategy = EnvAuthStrategy(name="OPENAI_API_KEY")
+
+
+class OpenAIEmbeddingsCustomProvider(BaseEmbeddingsProvider, OpenAIEmbeddings):
+    id = "openai-custom"
+    name = "OpenAI (general interface)"
+    models = ["*"]
+    model_id_key = "model"
+    pypi_package_deps = ["langchain_openai"]
+    auth_strategy = EnvAuthStrategy(name="OPENAI_API_KEY")
     registry = True
     fields = [
         TextField(
@@ -128,7 +158,6 @@ class AzureOpenAIEmbeddingsProvider(BaseEmbeddingsProvider, AzureOpenAIEmbedding
     auth_strategy = EnvAuthStrategy(
         name="AZURE_OPENAI_API_KEY", keyword_param="openai_api_key"
     )
-    registry = True
     fields = [
         TextField(key="azure_endpoint", label="Base API URL (optional)", format="text"),
     ]

docs/source/users/index.md

@@ -66,6 +66,7 @@ huggingface_hub = "jupyter_ai_magics:HfHubProvider"
 ollama = "jupyter_ai_magics.partner_providers.ollama:OllamaProvider"
 openai = "jupyter_ai_magics.partner_providers.openai:OpenAIProvider"
 openai-chat = "jupyter_ai_magics.partner_providers.openai:ChatOpenAIProvider"
+openai-chat-custom = "jupyter_ai_magics.partner_providers.openai:ChatOpenAICustomProvider"
 azure-chat-openai = "jupyter_ai_magics.partner_providers.openai:AzureChatOpenAIProvider"
 sagemaker-endpoint = "jupyter_ai_magics.partner_providers.aws:SmEndpointProvider"
 amazon-bedrock = "jupyter_ai_magics.partner_providers.aws:BedrockProvider"
@@ -87,6 +88,7 @@ gpt4all = "jupyter_ai_magics:GPT4AllEmbeddingsProvider"
 huggingface_hub = "jupyter_ai_magics:HfHubEmbeddingsProvider"
 ollama = "jupyter_ai_magics.partner_providers.ollama:OllamaEmbeddingsProvider"
 openai = "jupyter_ai_magics.partner_providers.openai:OpenAIEmbeddingsProvider"
+openai-custom = "jupyter_ai_magics.partner_providers.openai:OpenAIEmbeddingsCustomProvider"
 qianfan = "jupyter_ai_magics:QianfanEmbeddingsEndpointProvider"
 
 [tool.hatch.version]

docs/source/users/openrouter.md

@@ -44,6 +44,17 @@
       },
       "additionalProperties": false
     },
+    "embeddings_fields": {
+      "$comment": "Dictionary of model-specific fields, mapping LM GIDs to sub-dictionaries of field key-value pairs for embeddings.",
+      "type": "object",
+      "default": {},
+      "patternProperties": {
+        "^.*$": {
+          "anyOf": [{ "type": "object" }]
+        }
+      },
+      "additionalProperties": false
+    },
     "completions_fields": {
       "$comment": "Dictionary of model-specific fields, mapping LM GIDs to sub-dictionaries of field key-value pairs for completions.",
       "type": "object",

packages/jupyter-ai-magics/jupyter_ai_magics/partner_providers/ollama.py

@@ -184,6 +184,7 @@ def _process_existing_config(self, default_config):
     def _validate_model_ids(self, config):
         lm_provider_keys = ["model_provider_id", "completions_model_provider_id"]
         em_provider_keys = ["embeddings_provider_id"]
+        clm_provider_keys = ["completions_model_provider_id"]
 
         # if the currently selected language or embedding model are
         # forbidden, set them to `None` and log a warning.
@@ -201,6 +202,13 @@ def _validate_model_ids(self, config):
                     f"Embedding model {em_id} is forbidden by current allow/blocklists. Setting to None."
                 )
                 setattr(config, em_key, None)
+        for clm_key in clm_provider_keys:
+            clm_id = getattr(config, clm_key)
+            if clm_id is not None and not self._validate_model(clm_id, raise_exc=False):
+                self.log.warning(
+                    f"Completion model {clm_id} is forbidden by current allow/blocklists. Setting to None."
+                )
+                setattr(config, clm_key, None)
 
         # if the currently selected language or embedding model ids are
         # not associated with models, set them to `None` and log a warning.
@@ -218,6 +226,16 @@ def _validate_model_ids(self, config):
                     f"No embedding model is associated with '{em_id}'. Setting to None."
                 )
                 setattr(config, em_key, None)
+        for clm_key in clm_provider_keys:
+            clm_id = getattr(config, clm_key)
+            if (
+                clm_id is not None
+                and not get_lm_provider(clm_id, self._lm_providers)[1]
+            ):
+                self.log.warning(
+                    f"No completion model is associated with '{clm_id}'. Setting to None."
+                )
+                setattr(config, clm_key, None)
 
         return config
 
@@ -228,7 +246,8 @@ def _init_defaults(self):
         config_keys = GlobalConfig.model_fields.keys()
         schema_properties = self.validator.schema.get("properties", {})
         default_config = {
-            field: schema_properties.get(field).get("default") for field in config_keys
+            field: schema_properties.get(field, {}).get("default")
+            for field in config_keys
         }
         if self._defaults is None:
             return default_config
@@ -283,6 +302,36 @@ def _validate_config(self, config: GlobalConfig):
             # verify model is authenticated
             _validate_provider_authn(config, lm_provider)
 
+            # verify fields exist for this model if needed
+            if lm_provider.fields and config.model_provider_id not in config.fields:
+                config.fields[config.model_provider_id] = {}
+
+        # validate completions model config
+        if config.completions_model_provider_id:
+            _, completions_provider = get_lm_provider(
+                config.completions_model_provider_id, self._lm_providers
+            )
+
+            # verify model is declared by some provider
+            if not completions_provider:
+                raise ValueError(
+                    f"No language model is associated with '{config.completions_model_provider_id}'."
+                )
+
+            # verify model is not blocked
+            self._validate_model(config.completions_model_provider_id)
+
+            # verify model is authenticated
+            _validate_provider_authn(config, completions_provider)
+
+            # verify completions fields exist for this model if needed
+            if (
+                completions_provider.fields
+                and config.completions_model_provider_id
+                not in config.completions_fields
+            ):
+                config.completions_fields[config.completions_model_provider_id] = {}
+
         # validate embedding model config
         if config.embeddings_provider_id:
             _, em_provider = get_em_provider(
@@ -301,6 +350,13 @@ def _validate_config(self, config: GlobalConfig):
             # verify model is authenticated
             _validate_provider_authn(config, em_provider)
 
+            # verify embedding fields exist for this model if needed
+            if (
+                em_provider.fields
+                and config.embeddings_provider_id not in config.embeddings_fields
+            ):
+                config.embeddings_fields[config.embeddings_provider_id] = {}
+
     def _validate_model(self, model_id: str, raise_exc=True):
         """
         Validates a model against the set of allow/blocklists specified by the
@@ -349,6 +405,9 @@ def _write_config(self, new_config: GlobalConfig):
         new_config.completions_fields = {
             k: v for k, v in new_config.completions_fields.items() if v
         }
+        new_config.embeddings_fields = {
+            k: v for k, v in new_config.embeddings_fields.items() if v
+        }
 
         self._validate_config(new_config)
         with open(self.config_path, "w") as f:
@@ -462,18 +521,25 @@ def _provider_params(self, key, listing, completions: bool = False):
         # get config fields (e.g. base API URL, etc.)
         if completions:
             fields = config.completions_fields.get(model_uid, {})
+        elif key == "embeddings_provider_id":
+            fields = config.embeddings_fields.get(model_uid, {})
         else:
             fields = config.fields.get(model_uid, {})
 
         # exclude empty fields
         # TODO: modify the config manager to never save empty fields in the
         # first place.
-        for field_key in fields:
-            if isinstance(fields[field_key], str) and not len(fields[field_key]):
-                fields[field_key] = None
+        fields = {
+            k: None if isinstance(v, str) and not len(v) else v
+            for k, v in fields.items()
+        }
 
         # get authn fields
-        _, Provider = get_em_provider(model_uid, listing)
+        _, Provider = (
+            get_em_provider(model_uid, listing)
+            if key == "embeddings_provider_id"
+            else get_lm_provider(model_uid, listing)
+        )
         authn_fields = {}
         if Provider.auth_strategy and Provider.auth_strategy.type == "env":
             keyword_param = (

packages/jupyter-ai-magics/jupyter_ai_magics/partner_providers/openai.py

@@ -53,19 +53,21 @@ class DescribeConfigResponse(BaseModel):
     last_read: int
     completions_model_provider_id: Optional[str] = None
     completions_fields: Dict[str, Dict[str, Any]]
+    embeddings_fields: Dict[str, Dict[str, Any]]
 
 
 class UpdateConfigRequest(BaseModel):
     model_provider_id: Optional[str] = None
     embeddings_provider_id: Optional[str] = None
+    completions_model_provider_id: Optional[str] = None
     send_with_shift_enter: Optional[bool] = None
     api_keys: Optional[Dict[str, str]] = None
-    fields: Optional[Dict[str, Dict[str, Any]]] = None
     # if passed, this will raise an Error if the config was written to after the
     # time specified by `last_read` to prevent write-write conflicts.
     last_read: Optional[int] = None
-    completions_model_provider_id: Optional[str] = None
+    fields: Optional[Dict[str, Dict[str, Any]]] = None
     completions_fields: Optional[Dict[str, Dict[str, Any]]] = None
+    embeddings_fields: Optional[Dict[str, Dict[str, Any]]] = None
 
     @field_validator("send_with_shift_enter", "api_keys", "fields", mode="before")
     @classmethod
@@ -88,6 +90,7 @@ class GlobalConfig(BaseModel):
     api_keys: Dict[str, str]
     completions_model_provider_id: Optional[str] = None
     completions_fields: Dict[str, Dict[str, Any]]
+    embeddings_fields: Dict[str, Dict[str, Any]]
 
 
 class ListSlashCommandsEntry(BaseModel):

packages/jupyter-ai-magics/pyproject.toml

@@ -6,6 +6,8 @@
     'completions_fields': dict({
     }),
     'completions_model_provider_id': None,
+    'embeddings_fields': dict({
+    }),
     'embeddings_provider_id': None,
     'fields': dict({
     }),