开发者#

开发者文档适用于希望增强 Jupyter AI 功能的作者。

如果您有兴趣为 Jupyter AI 做出贡献，请参阅我们的 贡献者指南。

Pydantic 兼容性#

Jupyter AI 完全兼容使用 Pydantic v1 或 Pydantic v2 的 Python 环境。Jupyter AI 从 langchain.pydantic_v1 模块导入 Pydantic 类。开发者在扩展 Jupyter AI 类时应执行相同的操作。

有关在安装了 Pydantic v2 的环境中使用 langchain.pydantic_v1 的更多详细信息，请参阅 LangChain 关于 Pydantic 兼容性的文档。

Jupyter AI 模块 cookiecutter#

我们提供了一个 cookiecutter 模板，可用于生成预配置的 Jupyter AI 模块。这是一个 Python 包，公开了一个模板模型提供者和斜杠命令，用于与 Jupyter AI 集成。开发者可以根据需要扩展生成的 AI 模块。

要使用 cookiecutter 生成新的 AI 模块，请从仓库根目录运行以下命令：

pip install cookiecutter
cd packages/
cookiecutter jupyter-ai-module-cookiecutter

最后一个命令将打开一个向导，允许您设置包名称和其他一些元数据字段。默认情况下，包的名称为 jupyter-ai-test。

要在本地安装新的 AI 模块并使用生成的模板提供者和斜杠命令：

cd jupyter-ai-test/
pip install -e .

然后，在重新启动 JupyterLab 后，您将能够使用测试提供者和斜杠命令。

本页文档的其余部分详细说明了如何定义自定义模型提供者和斜杠命令。

自定义模型提供者#

您可以使用 LangChain 框架 API 定义新的提供者。自定义提供者继承自 jupyter-ai 的 BaseProvider 和 langchain 的 LLM。您可以导入 LangChain LLM 列表 中的预定义模型，或定义 自定义 LLM。在下面的示例中，我们使用一个虚拟的 FakeListLLM 模型定义了一个提供者，该模型从 responses 关键字参数返回响应。

# my_package/my_provider.py
from jupyter_ai_magics import BaseProvider
from langchain_community.llms import FakeListLLM


class MyProvider(BaseProvider, FakeListLLM):
    id = "my_provider"
    name = "My Provider"
    model_id_key = "model"
    models = [
        "model_a",
        "model_b"
    ]
    def __init__(self, **kwargs):
        model = kwargs.get("model_id")
        kwargs["responses"] = (
            ["This is a response from model 'a'"]
            if model == "model_a" else
            ["This is a response from model 'b'"]
        )
        super().__init__(**kwargs)

如果新提供者继承自 BaseChatModel，它将在聊天 UI 和魔法命令中都可用。否则，用户只能通过魔法命令使用新提供者。

要使新提供者可用，您需要将其声明为 入口点：

# my_package/pyproject.toml
[project]
name = "my_package"
version = "0.0.1"

[project.entry-points."jupyter_ai.model_providers"]
my-provider = "my_provider:MyProvider"

要测试上述最小提供者包是否有效，请使用以下命令安装它：

# 从 `my_package` 目录
pip install -e .

然后，重新启动 JupyterLab。您现在应该会在日志中看到一条信息消息，提到您的新提供者的 id：

[I 2023-10-29 13:56:16.915 AiExtension] 已注册模型提供者 `my_provider`。

自定义嵌入提供者#

要提供自定义嵌入模型，应定义一个嵌入提供者，实现 jupyter-ai 的 BaseEmbeddingsProvider 和 langchain 的 [Embeddings][Embeddings] 抽象类的 API。

from jupyter_ai_magics import BaseEmbeddingsProvider
from langchain.embeddings import FakeEmbeddings

class MyEmbeddingsProvider(BaseEmbeddingsProvider, FakeEmbeddings):
    id = "my_embeddings_provider"
    name = "My Embeddings Provider"
    model_id_key = "model"
    models = ["my_model"]

    def __init__(self, **kwargs):
        super().__init__(size=300, **kwargs)

Jupyter AI 使用入口点来发现嵌入提供者。在 pyproject.toml 文件中，将您的自定义嵌入提供者添加到 [project.entry-points."jupyter_ai.embeddings_model_providers"] 部分：

[project.entry-points."jupyter_ai.embeddings_model_providers"]
my-provider = "my_provider:MyEmbeddingsProvider"

自定义补全提供者#

任何从 BaseProvider 派生的模型提供者都可以用作补全提供者。 然而，某些提供者可能会从自定义处理补全请求中受益。

在 BaseProvider 的子类中，可以重写两个异步方法：

  • generate_inline_completions：接受一个请求（InlineCompletionRequest）并返回 InlineCompletionReply

  • stream_inline_completions：接受一个请求并生成一个初始回复（InlineCompletionReply），其中 isIncomplete 设置为 True，随后是后续的块（InlineCompletionStreamChunk）

在流式传输时，对于 stream_inline_completions() 方法的每次调用，所有回复和块应包含一个常量且唯一的字符串令牌来标识流。除给定项的最后一个块外，所有块的 done 值应设置为 False。

以下示例展示了一个自定义补全提供者的实现，该实现包含一次性发送多个补全的方法，以及并发流式传输多个补全的方法。 此示例中使用的 merge_iterators 函数的实现和解释可以在这里找到。

class MyCompletionProvider(BaseProvider, FakeListLLM):
    id = "my_provider"
    name = "My Provider"
    model_id_key = "model"
    models = ["model_a"]

    def __init__(self, **kwargs):
        kwargs["responses"] = ["This fake response will not be used for completion"]
        super().__init__(**kwargs)

    async def generate_inline_completions(self, request: InlineCompletionRequest):
        return InlineCompletionReply(
            list=InlineCompletionList(items=[
                {"insertText": "一只蚂蚁在忙自己的事"},
                {"insertText": "一只虫子在寻找零食"}
            ]),
            reply_to=request.number,
        )

    async def stream_inline_completions(self, request: InlineCompletionRequest):
        token_1 = f"t{request.number}s0"
        token_2 = f"t{request.number}s1"

        yield InlineCompletionReply(
            list=InlineCompletionList(
                items=[
                    {"insertText": "一只", "isIncomplete": True, "token": token_1},
                    {"insertText": "", "isIncomplete": True, "token": token_2}
                ]
            ),
            reply_to=request.number,
        )

        # 使用 merge_iterators
        async for reply in merge_iterators([
            self._stream("大象在雨中跳舞", request.number, token_1, start_with="一只"),
            self._stream("一群鸟在山周围飞翔", request.number, token_2)
        ]):
            yield reply

    async def _stream(self, sentence, request_number, token, start_with = ""):
        suggestion = start_with

        for fragment in sentence.split():
            await asyncio.sleep(0.75)
            suggestion += " " + fragment
            yield InlineCompletionStreamChunk(
                type="stream",
                response={"insertText": suggestion, "token": token},
                reply_to=request_number,
                done=False
            )

        # 最后，发送一条确认完成的消息
        yield InlineCompletionStreamChunk(
            type="stream",
            response={"insertText": suggestion, "token": token},
            reply_to=request_number,
            done=True,
        )

使用完整的笔记本内容进行补全#

InlineCompletionRequest 包含当前文档（文件或笔记本）的 path。 内联补全提供者可以使用此路径从磁盘中提取笔记本的内容， 但如果用户最近未保存笔记本，则该内容可能已过时。

通过将可能过时的前/后单元格内容与描述当前单元格（由 cell_id 标识）最新状态的 prefix 和 suffix 结合，可以稍微提高建议的准确性。

然而，从磁盘中读取完整的笔记本对于较大的笔记本来说可能很慢，这与内联补全的低延迟要求相冲突。

更好的方法是使用在启用 协作 文档模型时持久化在 jupyter-server 上的笔记本文档的实时副本。 需要安装两个包来访问协作模型：

  • jupyter-server-ydoc (>= 1.0) 在 jupyter-server 上运行时存储协作模型

  • jupyter-docprovider (>= 1.0) 重新配置 JupyterLab/Notebook 以使用协作模型 这两个包会随着 jupyter-collaboration（在 v3.0 或更新版本中）自动安装，然而，安装 jupyter-collaboration 并不是利用 协作 模型的必要条件。

下面的代码片段展示了如何从协作模型的内存副本中获取特定类型的所有单元格内容（无需额外的磁盘读取）。

from jupyter_ydoc import YNotebook


class MyCompletionProvider(BaseProvider, FakeListLLM):
    id = "my_provider"
    name = "My Provider"
    model_id_key = "model"
    models = ["model_a"]

    def __init__(self, **kwargs):
        kwargs["responses"] = ["This fake response will not be used for completion"]
        super().__init__(**kwargs)

    async def _get_prefix_and_suffix(self, request: InlineCompletionRequest):
        prefix = request.prefix
        suffix = request.suffix.strip()

        server_ydoc = self.server_settings.get("jupyter_server_ydoc", None)
        if not server_ydoc:
            # fallback to prefix/suffix from single cell
            return prefix, suffix

        is_notebook = request.path.endswith("ipynb")
        document = await server_ydoc.get_document(
            path=request.path,
            content_type="notebook" if is_notebook else "file",
            file_format="json" if is_notebook else "text"
        )
        if not document or not isinstance(document, YNotebook):
            return prefix, suffix

        cell_type = "markdown" if request.language == "markdown" else "code"

        is_before_request_cell = True
        before = []
        after = [suffix]

        for cell in document.ycells:
            if is_before_request_cell and cell["id"] == request.cell_id:
                is_before_request_cell = False
                continue
            if cell["cell_type"] != cell_type:
                continue
            source = cell["source"].to_py()
            if is_before_request_cell:
                before.append(source)
            else:
                after.append(source)

        before.append(prefix)
        prefix = "\n\n".join(before)
        suffix = "\n\n".join(after)
        return prefix, suffix

    async def generate_inline_completions(self, request: InlineCompletionRequest):
        prefix, suffix = await self._get_prefix_and_suffix(request)

        return InlineCompletionReply(
            list=InlineCompletionList(items=[
                {"insertText": your_llm_function(prefix, suffix)}
            ]),
            reply_to=request.number,
        )

提示模板#

每个提供者可以为每种支持的格式定义 提示模板。提示模板指导语言模型以特定格式生成输出。默认的提示模板是一个将格式映射到模板的 Python 字典。编写 BaseProvider 子类的开发者可以通过实现自己的get_prompt_template 函数，按输出格式、模型以及提交的提示来覆盖模板。每个提示模板都包含字符串 {prompt}，当用户运行魔法命令时，该字符串会被用户提供的提示替换。

自定义提示模板#

要修改给定格式的提示模板，请覆盖 get_prompt_template 方法：

from langchain.prompts import PromptTemplate


class MyProvider(BaseProvider, FakeListLLM):
    # (... 如上所述的属性 ...)
    def get_prompt_template(self, format) -> PromptTemplate:
        if format === "code":
            return PromptTemplate.from_template(
                "{prompt}\n\nProduce output as source code only, "
                "with no text or explanation before or after it."
            )
        return super().get_prompt_template(format)

请注意，这仅适用于 Jupyter AI 魔法（%ai 和 %%ai 魔法命令）。自定义提示模板尚未在聊天界面中使用。

聊天界面中的自定义斜杠命令#

您可以通过创建一个继承自 BaseChatHandler 的新类，向聊天界面添加自定义斜杠命令。设置其 id、name、用于在用户界面中显示的 help 消息以及 routing_type。每个自定义斜杠命令必须具有唯一的斜杠命令。斜杠命令只能包含 ASCII 字母、数字和下划线。每个斜杠命令必须是唯一的；自定义斜杠命令不能替换内置的斜杠命令。

在 Python 代码中添加您的自定义处理程序：

from jupyter_ai.chat_handlers.base import BaseChatHandler, SlashCommandRoutingType
from jupyter_ai.models import HumanChatMessage

class CustomChatHandler(BaseChatHandler):
    id = "custom"
    name = "自定义"
    help = "一个执行自定义操作的聊天处理器"
    routing_type = SlashCommandRoutingType(slash_id="custom")

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    async def process_message(self, message: HumanChatMessage):
        # 在这里放置你的自定义逻辑
        self.reply("<你的响应>", message)

Jupyter AI 使用入口点来支持自定义斜杠命令。 在 pyproject.toml 文件中，将您的自定义处理程序添加到 [project.entry-points."jupyter_ai.chat_handlers"] 部分：

[project.entry-points."jupyter_ai.chat_handlers"]
custom = "custom_package:CustomChatHandler"

然后，安装您的包，以便 Jupyter AI 将自定义聊天处理程序添加到现有的聊天处理程序中。

自定义消息页脚#

您可以提供一个自定义消息页脚，该页脚将在 UI 中每个消息下方渲染。为此，您需要编写或安装一个包含插件的 labextension，该插件提供 IJaiMessageFooter 令牌。此插件应返回一个 IJaiMessageFooter 对象，该对象定义要渲染的自定义页脚。

IJaiMessageFooter 对象包含一个属性 component，它应引用定义自定义消息页脚的 React 组件。Jupyter AI 将在每个聊天消息下方渲染此组件，并将每个聊天消息的定义作为对象传递给组件的 message 属性。message 属性采用 AiService.ChatMessage 类型，其中 AiService 从 @jupyter-ai/core/handler 导入。

以下是一个参考插件，它在每个代理消息下方显示一些自定义文本：

import React from 'react';
import {
  JupyterFrontEnd,
  JupyterFrontEndPlugin
} from '@jupyterlab/application';
import { IJaiMessageFooter, IJaiMessageFooterProps } from '@jupyter-ai/core/tokens';

export const footerPlugin: JupyterFrontEndPlugin<IJaiMessageFooter> = {
  id: '@your-org/your-package:custom-footer',
  autoStart: true,
  requires: [],
  provides: IJaiMessageFooter,
  activate: (app: JupyterFrontEnd): IJaiMessageFooter => {
    return {
      component: MessageFooter
    };
  }
};

function MessageFooter(props: IJaiMessageFooterProps) {
  if (props.message.type !== 'agent' && props.message.type !== 'agent-stream') {
    return null;
  }

  return (
    <div>这是一个测试页脚，渲染在每个代理消息下方。</div>
  );
}
On this page
Show Source