diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/models/persona.py b/packages/jupyter-ai-magics/jupyter_ai_magics/models/persona.py index 54d45b7ce..fe25397b0 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/models/persona.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/models/persona.py @@ -1,5 +1,6 @@ from langchain.pydantic_v1 import BaseModel + class Persona(BaseModel): """ Model of an **agent persona**, a struct that includes the name & avatar diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py index 973972c67..1ae80c5c5 100644 --- a/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py +++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py @@ -16,11 +16,7 @@ from dask.distributed import Client as DaskClient from jupyter_ai.config_manager import ConfigManager, Logger -from jupyter_ai.models import ( - AgentChatMessage, - ChatMessage, - HumanChatMessage, -) +from jupyter_ai.models import AgentChatMessage, ChatMessage, HumanChatMessage from jupyter_ai_magics import Persona from jupyter_ai_magics.providers import BaseProvider from langchain.pydantic_v1 import BaseModel @@ -186,9 +182,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=Persona( - name=persona.name, avatar_route=persona.avatar_route - ), + persona=Persona(name=persona.name, avatar_route=persona.avatar_route), ) for handler in self._root_chat_handlers.values(): diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/help.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/help.py index d79c32fd7..e46038da5 100644 --- a/packages/jupyter-ai/jupyter_ai/chat_handlers/help.py +++ b/packages/jupyter-ai/jupyter_ai/chat_handlers/help.py @@ -45,9 +45,7 @@ def build_help_message( time=time.time(), body=_format_help_message(chat_handlers, persona, unsupported_slash_commands), reply_to="", - persona=Persona( - name=persona.name, avatar_route=persona.avatar_route - ), + 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 index 3dea2af84..9abbf5e50 100644 --- a/packages/jupyter-ai/jupyter_ai/extension.py +++ b/packages/jupyter-ai/jupyter_ai/extension.py @@ -1,6 +1,6 @@ +import os import re import time -import os from dask.distributed import Client as DaskClient from importlib_metadata import entry_points @@ -31,12 +31,12 @@ RootChatHandler, ) - 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" handlers = [ diff --git a/packages/jupyter-ai/jupyter_ai/models.py b/packages/jupyter-ai/jupyter_ai/models.py index ff711a9cc..32353a694 100644 --- a/packages/jupyter-ai/jupyter_ai/models.py +++ b/packages/jupyter-ai/jupyter_ai/models.py @@ -1,7 +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 jupyter_ai_magics.providers import AuthStrategy, Field from langchain.pydantic_v1 import BaseModel, validator DEFAULT_CHUNK_SIZE = 2000