Remove config schema and legacy v2 persona references (#1420)

* move config models to config/config_models.py

* replace config schema validation with pydantic

* override config w/ traitlets defaults on init, fix all tests

* remove old v2 persona files

* fix jupyternaut route

Author: dlqqq

packages/jupyter-ai-magics/jupyter_ai_magics/__init__.py

@@ -13,7 +13,6 @@
     )
     from .exception import store_exception
     from .magics import AiMagics
-    from .models.persona import JupyternautPersona, Persona
     from .providers import (
         AI21Provider,
         BaseProvider,
@@ -24,7 +23,6 @@
     )
 else:
     _exports_by_module = {
-        # expose embedding model providers on the package root
         "embedding_providers": [
             "BaseEmbeddingsProvider",
             "GPT4AllEmbeddingsProvider",
@@ -33,9 +31,6 @@
         ],
         "exception": ["store_exception"],
         "magics": ["AiMagics"],
-        # expose JupyternautPersona on the package root
-        # required by `jupyter-ai`.
-        "models.persona": ["JupyternautPersona", "Persona"],
         # expose model providers on the package root
         "providers": [
             "AI21Provider",
@@ -80,8 +75,6 @@ def unload_ipython_extension(ipython):
     "QianfanEmbeddingsEndpointProvider",
     "store_exception",
     "AiMagics",
-    "JupyternautPersona",
-    "Persona",
     "AI21Provider",
     "BaseProvider",
     "GPT4AllProvider",

packages/jupyter-ai-magics/jupyter_ai_magics/base_provider.py

@@ -32,7 +32,6 @@
     InlineCompletionRequest,
     InlineCompletionStreamChunk,
 )
-from .models.persona import Persona
 
 CHAT_SYSTEM_PROMPT = """
 You are Jupyternaut, a conversational assistant living in JupyterLab to help users.
@@ -218,18 +217,6 @@ class BaseProvider(BaseModel):
     set to `True`, Jupyter AI will not pass the chat history to this provider
     when invoked."""
 
-    persona: ClassVar[Optional[Persona]] = None
-    """
-    The **persona** of this provider, a struct that defines the name and avatar
-    shown on agent replies in the chat UI. When set to `None`, `jupyter-ai` will
-    choose a default persona when rendering agent messages by this provider.
-
-    Because this field is set to `None` by default, `jupyter-ai` will render a
-    default persona for all providers that are included natively with the
-    `jupyter-ai` package. This field is reserved for Jupyter AI modules that
-    serve a custom provider and want to distinguish it in the chat UI.
-    """
-
     unsupported_slash_commands: ClassVar[set] = set()
     """
     A set of slash commands unsupported by this provider. Unsupported slash

packages/jupyter-ai-magics/jupyter_ai_magics/models/persona.py

@@ -0,0 +1 @@
+from .config_models import JaiConfig, UpdateConfigRequest, DescribeConfigResponse

packages/jupyter-ai/jupyter_ai/config/__init__.py

@@ -0,0 +1,91 @@
+from pydantic import BaseModel, field_validator
+from typing import Any, Optional
+
+class JaiConfig(BaseModel):
+    """
+    Pydantic model that serializes and validates the Jupyter AI config.
+    """
+
+    model_provider_id: Optional[str] = None
+    """
+    Model ID of the chat model.
+    """
+
+    embeddings_provider_id: Optional[str] = None
+    """
+    Model ID of the embedding model.
+    """
+
+    completions_model_provider_id: Optional[str] = None
+    """
+    Model ID of the completions model.
+    """
+
+    api_keys: dict[str, str] = {}
+    """
+    Dictionary of API keys. The name of each key should correspond to the
+    environment variable expected by the underlying client library, e.g.
+    "OPENAI_API_KEY".
+    """
+
+    send_with_shift_enter: bool = False
+    """
+    Whether the "Enter" key should create a new line instead of sending the
+    message.
+    """
+
+    fields: dict[str, dict[str, Any]] = {}
+    """
+    Dictionary that defines custom fields for each chat model.
+    Key: chat model ID.
+    Value: Dictionary of keyword arguments.
+    """
+
+    embeddings_fields: dict[str, dict[str, Any]] = {}
+    completions_fields: dict[str, dict[str, Any]] = {}
+
+
+class DescribeConfigResponse(BaseModel):
+    model_provider_id: Optional[str] = None
+    embeddings_provider_id: Optional[str] = None
+    send_with_shift_enter: bool
+    fields: dict[str, dict[str, Any]]
+
+    api_keys: list[str]
+    """
+    List of the names of the API keys. This deliberately does not include the
+    value of each API key in the interest of security.
+    """
+
+    last_read: int
+    """
+    Timestamp indicating when the configuration file was last read. Should be
+    passed to the subsequent UpdateConfig request if an update is made.
+    """
+
+    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
+    # 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
+    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
+    def ensure_not_none_if_passed(cls, field_val: Any) -> Any:
+        """
+        Field validator ensuring that certain fields are never `None` if set.
+        """
+        assert field_val is not None, "size may not be None"
+        return field_val