WIP: migrate Jupyternaut code into this package

Author: dlqqq

jupyter_ai_jupyternaut/extension_app.py

@@ -8,6 +8,7 @@
 
 from .config import ConfigManager, ConfigRestAPI
 from .handlers import RouteHandler
+from .jupyternaut import JupyternautPersona
 from .models import ChatModelsRestAPI, ModelParametersRestAPI
 from .secrets import EnvSecretsManager, SecretsRestAPI
 
@@ -176,7 +177,7 @@ def initialize_settings(self):
         }
 
         # Initialize ConfigManager
-        self.settings["jupyternaut.config_manager"] = ConfigManager(
+        config_manager = ConfigManager(
             config=self.config,
             log=self.log,
             allowed_providers=self.allowed_providers,
@@ -185,8 +186,14 @@ def initialize_settings(self):
             blocked_models=self.blocked_models,
             defaults=defaults,
         )
+        
+        # Bind ConfigManager instance to global settings dictionary
+        self.settings["jupyternaut.config_manager"] = config_manager
 
-        # Initialize SecretsManager
+        # Bind ConfigManager instance to Jupyternaut as a class variable
+        JupyternautPersona.config_manager = config_manager
+
+        # Initialize SecretsManager and bind it to global settings dictionary
         self.settings["jupyternaut.secrets_manager"] = EnvSecretsManager(parent=self)
 
     def _link_jupyter_server_extension(self, server_app: ServerApp):

jupyter_ai_jupyternaut/jupyternaut/__init__.py

@@ -0,0 +1 @@
+from .jupyternaut import JupyternautPersona

jupyter_ai_jupyternaut/jupyternaut/jupyternaut.py

@@ -0,0 +1,107 @@
+from typing import Any, Optional
+
+from jupyterlab_chat.models import Message
+from litellm import acompletion
+
+from ..base_persona import BasePersona, PersonaDefaults
+from ..persona_manager import SYSTEM_USERNAME
+from .prompt_template import (
+    JUPYTERNAUT_SYSTEM_PROMPT_TEMPLATE,
+    JupyternautSystemPromptArgs,
+)
+
+
+class JupyternautPersona(BasePersona):
+    """
+    The Jupyternaut persona, the main persona provided by Jupyter AI.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    @property
+    def defaults(self):
+        return PersonaDefaults(
+            name="Jupyternaut",
+            avatar_path="/api/ai/static/jupyternaut.svg",
+            description="The standard agent provided by JupyterLab. Currently has no tools.",
+            system_prompt="...",
+        )
+
+    async def process_message(self, message: Message) -> None:
+        if not hasattr(self, 'config_manager'):
+            self.send_message(
+                "Jupyternaut requires the `jupyter_ai_jupyternaut` server extension package.\n\n",
+                "Please make sure to first install that package in your environment & restart the server."
+            )
+        if not self.config_manager.chat_model:
+            self.send_message(
+                "No chat model is configured.\n\n"
+                "You must set one first in the Jupyter AI settings, found in 'Settings > AI Settings' from the menu bar."
+            )
+            return
+
+        model_id = self.config_manager.chat_model
+        model_args = self.config_manager.chat_model_args
+        context_as_messages = self.get_context_as_messages(model_id, message)
+        response_aiter = await acompletion(
+            **model_args,
+            model=model_id,
+            messages=[
+                *context_as_messages,
+                {
+                    "role": "user",
+                    "content": message.body,
+                },
+            ],
+            stream=True,
+        )
+
+        await self.stream_message(response_aiter)
+
+    def get_context_as_messages(
+        self, model_id: str, message: Message
+    ) -> list[dict[str, Any]]:
+        """
+        Returns the current context, including attachments and recent messages,
+        as a list of messages accepted by `litellm.acompletion()`.
+        """
+        system_msg_args = JupyternautSystemPromptArgs(
+            model_id=model_id,
+            persona_name=self.name,
+            context=self.process_attachments(message),
+        ).model_dump()
+
+        system_msg = {
+            "role": "system",
+            "content": JUPYTERNAUT_SYSTEM_PROMPT_TEMPLATE.render(**system_msg_args),
+        }
+
+        context_as_messages = [system_msg, *self._get_history_as_messages()]
+        return context_as_messages
+
+    def _get_history_as_messages(self, k: Optional[int] = 2) -> list[dict[str, Any]]:
+        """
+        Returns the current history as a list of messages accepted by
+        `litellm.acompletion()`.
+        """
+        # TODO: consider bounding history based on message size (e.g. total
+        # char/token count) instead of message count.
+        all_messages = self.ychat.get_messages()
+
+        # gather last k * 2 messages and return
+        # we exclude the last message since that is the human message just
+        # submitted by a user.
+        start_idx = 0 if k is None else -2 * k - 1
+        recent_messages: list[Message] = all_messages[start_idx:-1]
+
+        history: list[dict[str, Any]] = []
+        for msg in recent_messages:
+            role = (
+                "assistant"
+                if msg.sender.startswith("jupyter-ai-personas::")
+                else "system" if msg.sender == SYSTEM_USERNAME else "user"
+            )
+            history.append({"role": role, "content": msg.body})
+
+        return history

jupyter_ai_jupyternaut/jupyternaut/prompt_template.py

@@ -0,0 +1,55 @@
+from typing import Optional
+
+from jinja2 import Template
+from pydantic import BaseModel
+
+_JUPYTERNAUT_SYSTEM_PROMPT_FORMAT = """
+<instructions>
+
+You are {{persona_name}}, an AI agent provided in JupyterLab through the 'Jupyter AI' extension.
+
+Jupyter AI is an installable software package listed on PyPI and Conda Forge as `jupyter-ai`.
+
+When installed, Jupyter AI adds a chat experience in JupyterLab that allows multiple users to collaborate with one or more agents like yourself.
+
+You are not a language model, but rather an AI agent powered by a foundation model `{{model_id}}`.
+
+You are receiving a request from a user in JupyterLab. Your goal is to fulfill this request to the best of your ability.
+
+If you do not know the answer to a question, answer truthfully by responding that you do not know.
+
+You should use Markdown to format your response.
+
+Any code in your response must be enclosed in Markdown fenced code blocks (with triple backticks before and after).
+
+Any mathematical notation in your response must be expressed in LaTeX markup and enclosed in LaTeX delimiters.
+
+- Example of a correct response: The area of a circle is \\(\\pi * r^2\\).
+
+All dollar quantities (of USD) must be formatted in LaTeX, with the `$` symbol escaped by a single backslash `\\`.
+
+- Example of a correct response: `You have \\(\\$80\\) remaining.`
+
+You will receive any provided context and a relevant portion of the chat history.
+
+The user's request is located at the last message. Please fulfill the user's request to the best of your ability.
+</instructions>
+
+<context>
+{% if context %}The user has shared the following context:
+
+{{context}}
+{% else %}The user did not share any additional context.{% endif %}
+</context>
+""".strip()
+
+
+JUPYTERNAUT_SYSTEM_PROMPT_TEMPLATE: Template = Template(
+    _JUPYTERNAUT_SYSTEM_PROMPT_FORMAT
+)
+
+
+class JupyternautSystemPromptArgs(BaseModel):
+    persona_name: str
+    model_id: str
+    context: Optional[str] = None