remove 'avatar_path' from 'Persona', drop 'PersonaDescription'

jupyterlab · Apr 4, 2024 · 5ecce2e · 5ecce2e
1 parent 0ed6644
commit 5ecce2e
Show file tree

Hide file tree

Showing 7 changed files with 17 additions and 40 deletions.
diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/models/persona.py b/packages/jupyter-ai-magics/jupyter_ai_magics/models/persona.py
@@ -1,14 +1,5 @@
-import os
-from typing import ClassVar
-
 from langchain.pydantic_v1 import BaseModel
 
-JUPYTERNAUT_AVATAR_PATH = str(
-    os.path.join(os.path.dirname(__file__), "..", "static", "jupyternaut.svg")
-)
-JUPYTERNAUT_AVATAR_ROUTE = "api/ai/static/jupyternaut.svg"
-
-
 class Persona(BaseModel):
     """
     Model of an **agent persona**, a struct that includes the name & avatar
@@ -17,26 +8,18 @@ class Persona(BaseModel):
     Each persona is specific to a single provider, set on the `persona` field.
     """
 
-    name: ClassVar[str] = ...
+    name: str = ...
     """
     Name of the persona, e.g. "Jupyternaut". This is used to render the name
     shown on agent replies in the chat UI.
     """
 
-    avatar_route: ClassVar[str] = ...
+    avatar_route: str = ...
     """
     The server route that should be used the avatar of this persona. This is
     used to render the avatar shown on agent replies in the chat UI.
     """
 
-    avatar_path: ClassVar[str] = ...
-    """
-    The path to the avatar SVG file on the server filesystem. The server should
-    serve the file at this path on the route specified by `avatar_route`.
-    """
 
-
-class JupyternautPersona(Persona):
-    name: ClassVar[str] = "Jupyternaut"
-    avatar_route: ClassVar[str] = JUPYTERNAUT_AVATAR_ROUTE
-    avatar_path: ClassVar[str] = JUPYTERNAUT_AVATAR_PATH
+JUPYTERNAUT_AVATAR_ROUTE = "api/ai/static/jupyternaut.svg"
+JupyternautPersona = Persona(name="Jupyternaut", avatar_route=JUPYTERNAUT_AVATAR_ROUTE)
diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py
@@ -20,8 +20,8 @@
     AgentChatMessage,
     ChatMessage,
     HumanChatMessage,
-    PersonaDescription,
 )
+from jupyter_ai_magics import Persona
 from jupyter_ai_magics.providers import BaseProvider
 from langchain.pydantic_v1 import BaseModel
 
@@ -186,7 +186,7 @@ def reply(self, response: str, human_msg: Optional[HumanChatMessage] = None):
             time=time.time(),
             body=response,
             reply_to=human_msg.id if human_msg else "",
-            persona=PersonaDescription(
+            persona=Persona(
                 name=persona.name, avatar_route=persona.avatar_route
             ),
         )

diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/help.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/help.py
@@ -2,7 +2,7 @@
 from typing import Dict
 from uuid import uuid4
 
-from jupyter_ai.models import AgentChatMessage, HumanChatMessage, PersonaDescription
+from jupyter_ai.models import AgentChatMessage, HumanChatMessage
 from jupyter_ai_magics import Persona
 
 from .base import BaseChatHandler, SlashCommandRoutingType
@@ -45,7 +45,7 @@ def build_help_message(
         time=time.time(),
         body=_format_help_message(chat_handlers, persona, unsupported_slash_commands),
         reply_to="",
-        persona=PersonaDescription(
+        persona=Persona(
             name=persona.name, avatar_route=persona.avatar_route
         ),
     )

diff --git a/packages/jupyter-ai/jupyter_ai/extension.py b/packages/jupyter-ai/jupyter_ai/extension.py
@@ -1,6 +1,6 @@
-import logging
 import re
 import time
+import os
 
 from dask.distributed import Client as DaskClient
 from importlib_metadata import entry_points
@@ -32,9 +32,11 @@
     RootChatHandler,
 )
 
-JUPYTERNAUT_AVATAR_PATH = JupyternautPersona.avatar_path
-JUPYTERNAUT_AVATAR_ROUTE = JupyternautPersona.avatar_route
 
+JUPYTERNAUT_AVATAR_ROUTE = JupyternautPersona.avatar_route
+JUPYTERNAUT_AVATAR_PATH = str(
+    os.path.join(os.path.dirname(__file__), "static", "jupyternaut.svg")
+)
 
 class AiExtension(ExtensionApp):
     name = "jupyter_ai"

diff --git a/packages/jupyter-ai/jupyter_ai/models.py b/packages/jupyter-ai/jupyter_ai/models.py
@@ -1,6 +1,7 @@
 from typing import Any, Dict, List, Literal, Optional, Union
 
 from jupyter_ai_magics.providers import AuthStrategy, Field
+from jupyter_ai_magics import Persona
 from langchain.pydantic_v1 import BaseModel, validator
 
 DEFAULT_CHUNK_SIZE = 2000
@@ -29,15 +30,6 @@ class ChatClient(ChatUser):
     id: str
 
 
-class PersonaDescription(BaseModel):
-    """
-    Description of a persona to a chat client.
-    """
-
-    name: str
-    avatar_route: str
-
-
 class AgentChatMessage(BaseModel):
     type: Literal["agent"] = "agent"
     id: str
@@ -50,7 +42,7 @@ class AgentChatMessage(BaseModel):
     string if not applicable.
     """
 
-    persona: PersonaDescription
+    persona: Persona
     """
     The persona of the selected provider. If the selected provider is `None`,
     this defaults to a description of `JupyternautPersona`.

diff --git a/.../jupyter_ai_magics/static/jupyternaut.svg → ...yter-ai/jupyter_ai/static/jupyternaut.svg b/.../jupyter_ai_magics/static/jupyternaut.svg → ...yter-ai/jupyter_ai/static/jupyternaut.svg
diff --git a/packages/jupyter-ai/src/handler.ts b/packages/jupyter-ai/src/handler.ts
@@ -68,7 +68,7 @@ export namespace AiService {
     id: string;
   };
 
-  export type PersonaDescription = {
+  export type Persona = {
     name: string;
     avatar_route: string;
   };
@@ -79,7 +79,7 @@ export namespace AiService {
     time: number;
     body: string;
     reply_to: string;
-    persona: PersonaDescription;
+    persona: Persona;
   };
 
   export type HumanChatMessage = {