Merge branch 'main' into notebook-7

README.md

@@ -10,7 +10,7 @@ in JupyterLab and the Jupyter Notebook. More specifically, Jupyter AI offers:
   This works anywhere the IPython kernel runs (JupyterLab, Jupyter Notebook, Google Colab, Kaggle, VSCode, etc.).
 * A native chat UI in JupyterLab that enables you to work with generative AI as a conversational assistant.
 * Support for a wide range of generative model providers, including AI21, Anthropic, AWS, Cohere,
-  Gemini, Hugging Face, NVIDIA, and OpenAI.
+  Gemini, Hugging Face, MistralAI, NVIDIA, and OpenAI.
 * Local model support through GPT4All, enabling use of generative AI models on consumer grade machines
   with ease and privacy.
 

docs/source/index.md

@@ -8,7 +8,7 @@ in JupyterLab and the Jupyter Notebook. More specifically, Jupyter AI offers:
   This works anywhere the IPython kernel runs (JupyterLab, Jupyter Notebook, Google Colab, VSCode, etc.).
 * A native chat UI in JupyterLab that enables you to work with generative AI as a conversational assistant.
 * Support for a wide range of generative model providers and models
-  (AI21, Anthropic, Cohere, Gemini, Hugging Face, OpenAI, SageMaker, NVIDIA, etc.).
+  (AI21, Anthropic, Cohere, Gemini, Hugging Face, MistralAI, OpenAI, SageMaker, NVIDIA, etc.).
 
 <img src="_static/jupyter-ai-screenshot.png"
     alt='A screenshot of Jupyter AI showing the chat interface and the magic commands'

docs/source/users/index.md

@@ -165,6 +165,7 @@ Jupyter AI supports the following model providers:
 | Gemini              | `gemini`             | `GOOGLE_API_KEY`           | `langchain-google-genai`        |
 | GPT4All             | `gpt4all`            | N/A                        | `gpt4all`                       |
 | Hugging Face Hub    | `huggingface_hub`    | `HUGGINGFACEHUB_API_TOKEN` | `huggingface_hub`, `ipywidgets`, `pillow` |
+| MistralAI           | `mistralai`          | `MISTRAL_API_KEY`          | `langchain-mistralai`           |
 | NVIDIA              | `nvidia-chat`        | `NVIDIA_API_KEY`           | `langchain_nvidia_ai_endpoints` |
 | OpenAI              | `openai`             | `OPENAI_API_KEY`           | `langchain-openai`              |
 | OpenAI (chat)       | `openai-chat`        | `OPENAI_API_KEY`           | `langchain-openai`              |
@@ -513,6 +514,36 @@ The `/learn` command also provides downloading and processing papers from the [a
 Use the `/export` command to export the chat history from the current session to a markdown file named `chat_history-YYYY-MM-DD-HH-mm.md`. Using `/export <file_name>` will export the chat history to `<file_name>-YYYY-MM-DD-HH-mm.md` instead. You can export chat history as many times as you like in a single session. Each successive export will include the entire chat history up to that point in the session.
 
 
+### Fixing a code cell with an error
+
+The `/fix` command can be used to fix any code cell with an error output in a
+Jupyter notebook file. To start, type `/fix` into the chat input. Jupyter AI
+will then prompt you to select a cell with error output before sending the
+request.
+
+<img src="../_static/fix-no-error-cell-selected.png"
+    alt='Screenshot of the chat input containing `/fix` without a code cell with error output selected.'
+    class="screenshot" />
+
+Then click on a code cell with error output. A blue bar should appear
+immediately to the left of the code cell.
+
+<img src="../_static/fix-error-cell-selected.png"
+    alt='Screenshot of a code cell with error output selected.'
+    class="screenshot" />
+
+After this, the Send button to the right of the chat input will be enabled, and
+you can use your mouse or keyboard to send `/fix` to Jupyternaut. The code cell
+and its associated error output are included in the message automatically. When
+complete, Jupyternaut will reply with suggested code that should fix the error.
+You can use the action toolbar under each code block to quickly replace the
+contents of the failing cell.
+
+<img src="../_static/fix-response.png"
+    alt='Screenshot of a response from `/fix`, with the "Replace active cell" action hovered.'
+    class="screenshot" style="max-width:65%" />
+
+
 ### Additional chat commands
 
 To clear the chat panel, use the `/clear` command. This does not reset the AI model; the model may still remember previous messages that you sent it, and it may use them to inform its responses.

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

@@ -0,0 +1,32 @@
+from jupyter_ai_magics.providers import BaseProvider, EnvAuthStrategy
+from langchain_mistralai import ChatMistralAI, MistralAIEmbeddings
+
+from ..embedding_providers import BaseEmbeddingsProvider
+
+
+class MistralAIProvider(BaseProvider, ChatMistralAI):
+    id = "mistralai"
+    name = "MistralAI"
+    models = [
+        "open-mistral-7b",
+        "open-mixtral-8x7b",
+        "open-mixtral-8x22b",
+        "mistral-small-latest",
+        "mistral-medium-latest",
+        "mistral-large-latest",
+        "codestral-latest",
+    ]
+    model_id_key = "model"
+    auth_strategy = EnvAuthStrategy(name="MISTRAL_API_KEY")
+    pypi_package_deps = ["langchain-mistralai"]
+
+
+class MistralAIEmbeddingsProvider(BaseEmbeddingsProvider, MistralAIEmbeddings):
+    id = "mistralai"
+    name = "MistralAI"
+    models = [
+        "mistral-embed",
+    ]
+    model_id_key = "model"
+    pypi_package_deps = ["langchain-mistralai"]
+    auth_strategy = EnvAuthStrategy(name="MISTRAL_API_KEY")

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

@@ -20,8 +20,6 @@
 
 from jsonpath_ng import parse
 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 (
     ChatPromptTemplate,
     HumanMessagePromptTemplate,
@@ -44,6 +42,7 @@
     SagemakerEndpoint,
     Together,
 )
+from langchain_community.llms.sagemaker_endpoint import LLMContentHandler
 
 # this is necessary because `langchain.pydantic_v1.main` does not include
 # `ModelMetaclass`, as it is not listed in `__all__` by the `pydantic.main`

packages/jupyter-ai-magics/pyproject.toml

@@ -41,13 +41,14 @@ all = [
     "huggingface_hub",
     "ipywidgets",
     "langchain_anthropic",
+    "langchain-mistralai",
     "langchain_nvidia_ai_endpoints",
+    "langchain-google-genai",
     "langchain-openai",
     "pillow",
     "boto3",
     "qianfan",
     "together",
-    "langchain-google-genai",
 ]
 
 [project.entry-points."jupyter_ai.model_providers"]
@@ -67,10 +68,12 @@ qianfan = "jupyter_ai_magics:QianfanProvider"
 nvidia-chat = "jupyter_ai_magics.partner_providers.nvidia:ChatNVIDIAProvider"
 together-ai = "jupyter_ai_magics:TogetherAIProvider"
 gemini = "jupyter_ai_magics.partner_providers.gemini:GeminiProvider"
+mistralai = "jupyter_ai_magics.partner_providers.mistralai:MistralAIProvider"
 
 [project.entry-points."jupyter_ai.embeddings_model_providers"]
 bedrock = "jupyter_ai_magics:BedrockEmbeddingsProvider"
 cohere = "jupyter_ai_magics:CohereEmbeddingsProvider"
+mistralai = "jupyter_ai_magics.partner_providers.mistralai:MistralAIEmbeddingsProvider"
 gpt4all = "jupyter_ai_magics:GPT4AllEmbeddingsProvider"
 huggingface_hub = "jupyter_ai_magics:HfHubEmbeddingsProvider"
 openai = "jupyter_ai_magics.partner_providers.openai:OpenAIEmbeddingsProvider"

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

@@ -3,6 +3,7 @@
 from .clear import ClearChatHandler
 from .default import DefaultChatHandler
 from .export import ExportChatHandler
+from .fix import FixChatHandler
 from .generate import GenerateChatHandler
 from .help import HelpChatHandler
 from .learn import LearnChatHandler

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

@@ -43,19 +43,6 @@ def create_llm_chain(
                 llm=llm, prompt=prompt_template, verbose=True, memory=self.memory
             )
 
-    def clear_memory(self):
-        # clear chain memory
-        if self.memory:
-            self.memory.clear()
-
-        # clear transcript for existing chat clients
-        reply_message = ClearMessage()
-        self.reply(reply_message)
-
-        # clear transcript for new chat clients
-        if self._chat_history:
-            self._chat_history.clear()
-
     async def process_message(self, message: HumanChatMessage):
         self.get_llm_chain()
         response = await self.llm_chain.apredict(input=message.body, stop=["\nHuman:"])

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

@@ -0,0 +1,103 @@
+from typing import Dict, Type
+
+from jupyter_ai.models import CellWithErrorSelection, HumanChatMessage
+from jupyter_ai_magics.providers import BaseProvider
+from langchain.chains import LLMChain
+from langchain.prompts import PromptTemplate
+
+from .base import BaseChatHandler, SlashCommandRoutingType
+
+FIX_STRING_TEMPLATE = """
+You are Jupyternaut, a conversational assistant living in JupyterLab. Please fix
+the notebook cell described below.
+
+Additional instructions:
+
+{extra_instructions}
+
+Input cell:
+
+```
+{cell_content}
+```
+
+Output error:
+
+```
+{traceback}
+
+{error_name}: {error_value}
+```
+""".strip()
+
+FIX_PROMPT_TEMPLATE = PromptTemplate(
+    input_variables=[
+        "extra_instructions",
+        "cell_content",
+        "traceback",
+        "error_name",
+        "error_value",
+    ],
+    template=FIX_STRING_TEMPLATE,
+)
+
+
+class FixChatHandler(BaseChatHandler):
+    """
+    Accepts a `HumanChatMessage` that includes a cell with error output and
+    recommends a fix as a reply. If a cell with error output is not included,
+    this chat handler does nothing.
+
+    `/fix` also accepts additional instructions in natural language as an
+    arbitrary number of arguments, e.g.
+
+    ```
+    /fix use the numpy library to implement this function instead.
+    ```
+    """
+
+    id = "fix"
+    name = "Fix error cell"
+    help = "Fix an error cell selected in your notebook"
+    routing_type = SlashCommandRoutingType(slash_id="fix")
+    uses_llm = True
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def create_llm_chain(
+        self, provider: Type[BaseProvider], provider_params: Dict[str, str]
+    ):
+        unified_parameters = {
+            **provider_params,
+            **(self.get_model_parameters(provider, provider_params)),
+        }
+        llm = provider(**unified_parameters)
+
+        self.llm = llm
+        self.llm_chain = LLMChain(llm=llm, prompt=FIX_PROMPT_TEMPLATE, verbose=True)
+
+    async def process_message(self, message: HumanChatMessage):
+        if not (message.selection and message.selection.type == "cell-with-error"):
+            self.reply(
+                "`/fix` requires an active code cell with error output. Please click on a cell with error output and retry.",
+                message,
+            )
+            return
+
+        # hint type of selection
+        selection: CellWithErrorSelection = message.selection
+
+        # parse additional instructions specified after `/fix`
+        extra_instructions = message.body[4:].strip() or "None."
+
+        self.get_llm_chain()
+        response = await self.llm_chain.apredict(
+            extra_instructions=extra_instructions,
+            stop=["\nHuman:"],
+            cell_content=selection.source,
+            error_name=selection.error.name,
+            error_value=selection.error.value,
+            traceback="\n".join(selection.error.traceback),
+        )
+        self.reply(response, message)

packages/jupyter-ai/jupyter_ai/extension.py

@@ -17,6 +17,7 @@
     ClearChatHandler,
     DefaultChatHandler,
     ExportChatHandler,
+    FixChatHandler,
     GenerateChatHandler,
     HelpChatHandler,
     LearnChatHandler,
@@ -264,13 +265,17 @@ def initialize_settings(self):
         ask_chat_handler = AskChatHandler(**chat_handler_kwargs, retriever=retriever)
 
         export_chat_handler = ExportChatHandler(**chat_handler_kwargs)
+
+        fix_chat_handler = FixChatHandler(**chat_handler_kwargs)
+
         jai_chat_handlers = {
             "default": default_chat_handler,
             "/ask": ask_chat_handler,
             "/clear": clear_chat_handler,
             "/generate": generate_chat_handler,
             "/learn": learn_chat_handler,
             "/export": export_chat_handler,
+            "/fix": fix_chat_handler,
         }
 
         help_chat_handler = HelpChatHandler(

packages/jupyter-ai/jupyter_ai/handlers.py

@@ -207,6 +207,7 @@ async def on_message(self, message):
             id=chat_message_id,
             time=time.time(),
             body=chat_request.prompt,
+            selection=chat_request.selection,
             client=self.chat_client,
         )
 

packages/jupyter-ai/jupyter_ai/models.py

@@ -8,9 +8,29 @@
 DEFAULT_CHUNK_OVERLAP = 100
 
 
+class CellError(BaseModel):
+    name: str
+    value: str
+    traceback: List[str]
+
+
+class CellWithErrorSelection(BaseModel):
+    type: Literal["cell-with-error"] = "cell-with-error"
+    source: str
+    error: CellError
+
+
+Selection = Union[CellWithErrorSelection]
+
+
 # the type of message used to chat with the agent
 class ChatRequest(BaseModel):
     prompt: str
+    # TODO: This currently is only used when a user runs the /fix slash command.
+    # In the future, the frontend should set the text selection on this field in
+    # the `HumanChatMessage` it sends to JAI, instead of appending the text
+    # selection to `body` in the frontend.
+    selection: Optional[Selection]
 
 
 class ChatUser(BaseModel):
@@ -55,6 +75,7 @@ class HumanChatMessage(BaseModel):
     time: float
     body: str
     client: ChatClient
+    selection: Optional[Selection]
 
 
 class ConnectionMessage(BaseModel):