Merge branch 'jupyterlab:main' into use-rendermime-markdown

docs/source/users/index.md

@@ -765,6 +765,28 @@ The `--response-path` option is a [JSONPath](https://goessner.net/articles/JsonP
 You can specify an allowlist, to only allow only a certain list of providers, or
 a blocklist, to block some providers.
 
+### Configuring default models and API keys
+
+This configuration allows for setting a default language and embedding models, and their corresponding API keys.
+These values are offered as a starting point for users, so they don't have to select the models and API keys, however,
+the selections they make in the settings panel will take precedence over these values.
+
+Specify default language model
+```bash
+jupyter lab --AiExtension.default_language_model=bedrock-chat:anthropic.claude-v2
+```
+
+Specify default embedding model
+```bash
+jupyter lab --AiExtension.default_embeddings_model=bedrock:amazon.titan-embed-text-v1
+```
+
+Specify default API keys
+```bash
+jupyter lab --AiExtension.default_api_keys={'OPENAI_API_KEY': 'sk-abcd'}
+```
+
+
 ### Blocklisting providers
 
 This configuration allows for blocking specific providers in the settings panel.

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

@@ -481,8 +481,13 @@ def run_ai_cell(self, args: CellArgs, prompt: str):
         if args.model_id in self.custom_model_registry and isinstance(
             self.custom_model_registry[args.model_id], LLMChain
         ):
+            # Get the output, either as raw text or as the contents of the 'text' key of a dict
+            invoke_output = self.custom_model_registry[args.model_id].invoke(prompt)
+            if isinstance(invoke_output, dict):
+                invoke_output = invoke_output.get("text")
+
             return self.display_output(
-                self.custom_model_registry[args.model_id].run(prompt),
+                invoke_output,
                 args.format,
                 {"jupyter_ai": {"custom_chain_id": args.model_id}},
             )

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

@@ -11,7 +11,13 @@
 from langchain.chat_models.base import BaseChatModel
 from langchain.llms.sagemaker_endpoint import LLMContentHandler
 from langchain.llms.utils import enforce_stop_tokens
-from langchain.prompts import PromptTemplate
+from langchain.prompts import (
+    ChatPromptTemplate,
+    HumanMessagePromptTemplate,
+    MessagesPlaceholder,
+    PromptTemplate,
+    SystemMessagePromptTemplate,
+)
 from langchain.pydantic_v1 import BaseModel, Extra, root_validator
 from langchain.schema import LLMResult
 from langchain.utils import get_from_dict_or_env
@@ -42,6 +48,49 @@
     from pydantic.main import ModelMetaclass
 
 
+CHAT_SYSTEM_PROMPT = """
+You are Jupyternaut, a conversational assistant living in JupyterLab to help users.
+You are not a language model, but rather an application built on a foundation model from {provider_name} called {local_model_id}.
+You are talkative and you provide lots of specific details from the foundation model's context.
+You may use Markdown to format your response.
+Code blocks must be formatted in Markdown.
+Math should be rendered with inline TeX markup, surrounded by $.
+If you do not know the answer to a question, answer truthfully by responding that you do not know.
+The following is a friendly conversation between you and a human.
+""".strip()
+
+CHAT_DEFAULT_TEMPLATE = """Current conversation:
+{history}
+Human: {input}
+AI:"""
+
+
+COMPLETION_SYSTEM_PROMPT = """
+You are an application built to provide helpful code completion suggestions.
+You should only produce code. Keep comments to minimum, use the
+programming language comment syntax. Produce clean code.
+The code is written in JupyterLab, a data analysis and code development
+environment which can execute code extended with additional syntax for
+interactive features, such as magics.
+""".strip()
+
+# only add the suffix bit if present to save input tokens/computation time
+COMPLETION_DEFAULT_TEMPLATE = """
+The document is called `{{filename}}` and written in {{language}}.
+{% if suffix %}
+The code after the completion request is:
+
+```
+{{suffix}}
+```
+{% endif %}
+
+Complete the following code:
+
+```
+{{prefix}}"""
+
+
 class EnvAuthStrategy(BaseModel):
     """Require one auth token via an environment variable."""
 
@@ -265,6 +314,55 @@ def get_prompt_template(self, format) -> PromptTemplate:
         else:
             return self.prompt_templates["text"]  # Default to plain format
 
+    def get_chat_prompt_template(self) -> PromptTemplate:
+        """
+        Produce a prompt template optimised for chat conversation.
+        The template should take two variables: history and input.
+        """
+        name = self.__class__.name
+        if self.is_chat_provider:
+            return ChatPromptTemplate.from_messages(
+                [
+                    SystemMessagePromptTemplate.from_template(
+                        CHAT_SYSTEM_PROMPT
+                    ).format(provider_name=name, local_model_id=self.model_id),
+                    MessagesPlaceholder(variable_name="history"),
+                    HumanMessagePromptTemplate.from_template("{input}"),
+                ]
+            )
+        else:
+            return PromptTemplate(
+                input_variables=["history", "input"],
+                template=CHAT_SYSTEM_PROMPT.format(
+                    provider_name=name, local_model_id=self.model_id
+                )
+                + "\n\n"
+                + CHAT_DEFAULT_TEMPLATE,
+            )
+
+    def get_completion_prompt_template(self) -> PromptTemplate:
+        """
+        Produce a prompt template optimised for inline code or text completion.
+        The template should take variables: prefix, suffix, language, filename.
+        """
+        if self.is_chat_provider:
+            return ChatPromptTemplate.from_messages(
+                [
+                    SystemMessagePromptTemplate.from_template(COMPLETION_SYSTEM_PROMPT),
+                    HumanMessagePromptTemplate.from_template(
+                        COMPLETION_DEFAULT_TEMPLATE, template_format="jinja2"
+                    ),
+                ]
+            )
+        else:
+            return PromptTemplate(
+                input_variables=["prefix", "suffix", "language", "filename"],
+                template=COMPLETION_SYSTEM_PROMPT
+                + "\n\n"
+                + COMPLETION_DEFAULT_TEMPLATE,
+                template_format="jinja2",
+            )
+
     @property
     def is_chat_provider(self):
         return isinstance(self, BaseChatModel)

packages/jupyter-ai/jupyter_ai/chat_handlers/default.py

@@ -4,32 +4,9 @@
 from jupyter_ai_magics.providers import BaseProvider
 from langchain.chains import ConversationChain
 from langchain.memory import ConversationBufferWindowMemory
-from langchain.prompts import (
-    ChatPromptTemplate,
-    HumanMessagePromptTemplate,
-    MessagesPlaceholder,
-    PromptTemplate,
-    SystemMessagePromptTemplate,
-)
 
 from .base import BaseChatHandler, SlashCommandRoutingType
 
-SYSTEM_PROMPT = """
-You are Jupyternaut, a conversational assistant living in JupyterLab to help users.
-You are not a language model, but rather an application built on a foundation model from {provider_name} called {local_model_id}.
-You are talkative and you provide lots of specific details from the foundation model's context.
-You may use Markdown to format your response.
-Code blocks must be formatted in Markdown.
-Math should be rendered with inline TeX markup, surrounded by $.
-If you do not know the answer to a question, answer truthfully by responding that you do not know.
-The following is a friendly conversation between you and a human.
-""".strip()
-
-DEFAULT_TEMPLATE = """Current conversation:
-{history}
-Human: {input}
-AI:"""
-
 
 class DefaultChatHandler(BaseChatHandler):
     id = "default"
@@ -49,27 +26,10 @@ def create_llm_chain(
         model_parameters = self.get_model_parameters(provider, provider_params)
         llm = provider(**provider_params, **model_parameters)
 
-        if llm.is_chat_provider:
-            prompt_template = ChatPromptTemplate.from_messages(
-                [
-                    SystemMessagePromptTemplate.from_template(SYSTEM_PROMPT).format(
-                        provider_name=provider.name, local_model_id=llm.model_id
-                    ),
-                    MessagesPlaceholder(variable_name="history"),
-                    HumanMessagePromptTemplate.from_template("{input}"),
-                ]
-            )
-            self.memory = ConversationBufferWindowMemory(return_messages=True, k=2)
-        else:
-            prompt_template = PromptTemplate(
-                input_variables=["history", "input"],
-                template=SYSTEM_PROMPT.format(
-                    provider_name=provider.name, local_model_id=llm.model_id
-                )
-                + "\n\n"
-                + DEFAULT_TEMPLATE,
-            )
-            self.memory = ConversationBufferWindowMemory(k=2)
+        prompt_template = llm.get_chat_prompt_template()
+        self.memory = ConversationBufferWindowMemory(
+            return_messages=llm.is_chat_provider, k=2
+        )
 
         self.llm = llm
         self.llm_chain = ConversationChain(

packages/jupyter-ai/jupyter_ai/completions/handlers/default.py

@@ -18,32 +18,6 @@
 )
 from .base import BaseInlineCompletionHandler
 
-SYSTEM_PROMPT = """
-You are an application built to provide helpful code completion suggestions.
-You should only produce code. Keep comments to minimum, use the
-programming language comment syntax. Produce clean code.
-The code is written in JupyterLab, a data analysis and code development
-environment which can execute code extended with additional syntax for
-interactive features, such as magics.
-""".strip()
-
-AFTER_TEMPLATE = """
-The code after the completion request is:
-
-```
-{suffix}
-```
-""".strip()
-
-DEFAULT_TEMPLATE = """
-The document is called `{filename}` and written in {language}.
-{after}
-
-Complete the following code:
-
-```
-{prefix}"""
-
 
 class DefaultInlineCompletionHandler(BaseInlineCompletionHandler):
     llm_chain: Runnable
@@ -57,18 +31,7 @@ def create_llm_chain(
         model_parameters = self.get_model_parameters(provider, provider_params)
         llm = provider(**provider_params, **model_parameters)
 
-        if llm.is_chat_provider:
-            prompt_template = ChatPromptTemplate.from_messages(
-                [
-                    SystemMessagePromptTemplate.from_template(SYSTEM_PROMPT),
-                    HumanMessagePromptTemplate.from_template(DEFAULT_TEMPLATE),
-                ]
-            )
-        else:
-            prompt_template = PromptTemplate(
-                input_variables=["prefix", "suffix", "language", "filename"],
-                template=SYSTEM_PROMPT + "\n\n" + DEFAULT_TEMPLATE,
-            )
+        prompt_template = llm.get_completion_prompt_template()
 
         self.llm = llm
         self.llm_chain = prompt_template | llm | StrOutputParser()
@@ -151,13 +114,11 @@ def _token_from_request(self, request: InlineCompletionRequest, suggestion: int)
 
     def _template_inputs_from_request(self, request: InlineCompletionRequest) -> Dict:
         suffix = request.suffix.strip()
-        # only add the suffix template if the suffix is there to save input tokens/computation time
-        after = AFTER_TEMPLATE.format(suffix=suffix) if suffix else ""
         filename = request.path.split("/")[-1] if request.path else "untitled"
 
         return {
             "prefix": request.prefix,
-            "after": after,
+            "suffix": suffix,
             "language": request.language,
             "filename": filename,
             "stop": ["\n```"],

packages/jupyter-ai/jupyter_ai/config_manager.py

@@ -105,6 +105,7 @@ def __init__(
         blocked_providers: Optional[List[str]],
         allowed_models: Optional[List[str]],
         blocked_models: Optional[List[str]],
+        defaults: dict,
         *args,
         **kwargs,
     ):
@@ -120,6 +121,8 @@ def __init__(
         self._blocked_providers = blocked_providers
         self._allowed_models = allowed_models
         self._blocked_models = blocked_models
+        self._defaults = defaults
+        """Provider defaults."""
 
         self._last_read: Optional[int] = None
         """When the server last read the config file. If the file was not
@@ -146,14 +149,20 @@ def _init_validator(self) -> Validator:
             self.validator = Validator(schema)
 
     def _init_config(self):
+        default_config = self._init_defaults()
         if os.path.exists(self.config_path):
-            self._process_existing_config()
+            self._process_existing_config(default_config)
         else:
-            self._create_default_config()
+            self._create_default_config(default_config)
 
-    def _process_existing_config(self):
+    def _process_existing_config(self, default_config):
         with open(self.config_path, encoding="utf-8") as f:
-            config = GlobalConfig(**json.loads(f.read()))
+            existing_config = json.loads(f.read())
+            merged_config = Merger.merge(
+                default_config,
+                {k: v for k, v in existing_config.items() if v is not None},
+            )
+            config = GlobalConfig(**merged_config)
             validated_config = self._validate_lm_em_id(config)
 
             # re-write to the file to validate the config and apply any
@@ -192,14 +201,23 @@ def _validate_lm_em_id(self, config):
 
         return config
 
-    def _create_default_config(self):
-        properties = self.validator.schema.get("properties", {})
+    def _create_default_config(self, default_config):
+        self._write_config(GlobalConfig(**default_config))
+
+    def _init_defaults(self):
         field_list = GlobalConfig.__fields__.keys()
+        properties = self.validator.schema.get("properties", {})
         field_dict = {
             field: properties.get(field).get("default") for field in field_list
         }
-        default_config = GlobalConfig(**field_dict)
-        self._write_config(default_config)
+        if self._defaults is None:
+            return field_dict
+
+        for field in field_list:
+            default_value = self._defaults.get(field)
+            if default_value is not None:
+                field_dict[field] = default_value
+        return field_dict
 
     def _read_config(self) -> GlobalConfig:
         """Returns the user's current configuration as a GlobalConfig object.