From 97fa2f5280a741d71fecc0deda4e670bd49a75bc Mon Sep 17 00:00:00 2001 From: krassowski <5832902+krassowski@users.noreply.github.com> Date: Wed, 8 May 2024 11:23:19 +0100 Subject: [PATCH 1/5] Document how to create completions using full notebook --- docs/source/developers/index.md | 88 +++++++++++++++++++++++++++++++++ 1 file changed, 88 insertions(+) diff --git a/docs/source/developers/index.md b/docs/source/developers/index.md index aac92328..0f85a423 100644 --- a/docs/source/developers/index.md +++ b/docs/source/developers/index.md @@ -228,6 +228,94 @@ class MyCompletionProvider(BaseProvider, FakeListLLM): ) ``` + +#### Using the full notebook content for completions + +The `InlineCompletionRequest` contains the `path` of the current document (file or notebook). +Inline completion providers can use this path to extract the content of the notebook from the disk, +however such content may be outdated if the user has not saved the notebook recently. + +The accuracy of the suggestions can be slightly improved by combining the potentially outdated content of previous/following cells +with the `prefix` and `suffix` which describe the up-to-date state of the current cell (identified by `cell_id`). + +Still, reading the full notebook from the disk may be slow for larger notebooks, which conflicts with the low latency requirement of inline completion. + +A better approach is to use the live copy of the notebook document that is persisted on the jupyter-server when *collaborative* document models are enabled. +Two packages need to be installed to access the collaborative models: +- `jupyter-server-ydoc` (>= 1.0) stores the collaborative models in the jupyter-server on runtime +- `jupyter-docprovider` (>= 1.0) reconfigures JupyterLab/Notebook to use the collaborative models + +Both packages are automatically installed with `jupyter-collaboration` (in v3.0 or newer), however installing `jupyter-collaboration` is not required to take advantage of *collaborative* models. + +The snippet below demonstrates how to retrieve the content of all cells of a given type from the in-memory copy of the collaborative model (without additional disk reads). + +```python +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) + + def _get_prefix_and_suffix(self, request: InlineCompletionRequest): + prefix = request.prefix + suffix = request.suffix.strip() + + server_ydoc = self.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 = 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 = self._get_prefix_and_suffix(request) + + return InlineCompletionReply( + list=InlineCompletionList(items=[ + {"insertText": your_llm_function(prefix, suffix)} + ]), + reply_to=request.number, + ) +``` + + ## Prompt templates Each provider can define **prompt templates** for each supported format. A prompt From 55184c6a54a87c25d0b4c193bc3ac3048c6021e1 Mon Sep 17 00:00:00 2001 From: krassowski <5832902+krassowski@users.noreply.github.com> Date: Sun, 26 May 2024 17:10:03 +0100 Subject: [PATCH 2/5] Pass a subset of server-settings down to providers --- docs/source/developers/index.md | 2 +- .../jupyter_ai_magics/providers.py | 8 ++++++++ packages/jupyter-ai/jupyter_ai/extension.py | 13 ++++++++++++- 3 files changed, 21 insertions(+), 2 deletions(-) diff --git a/docs/source/developers/index.md b/docs/source/developers/index.md index 0f85a423..6b0bd4c5 100644 --- a/docs/source/developers/index.md +++ b/docs/source/developers/index.md @@ -267,7 +267,7 @@ class MyCompletionProvider(BaseProvider, FakeListLLM): prefix = request.prefix suffix = request.suffix.strip() - server_ydoc = self.settings.get("jupyter_server_ydoc", None) + server_ydoc = self.server_settings.get("jupyter_server_ydoc", None) if not server_ydoc: # fallback to prefix/suffix from single cell return prefix, suffix diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py index 1bbc4ce5..bef7efac 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py @@ -5,6 +5,7 @@ import io import json from concurrent.futures import ThreadPoolExecutor +from types import MappingProxyType from typing import ( Any, AsyncIterator, @@ -265,6 +266,13 @@ class Config: provider is selected. """ + server_settings: ClassVar[Optional[MappingProxyType[str, Any]]] = None + """Settings passed on from jupyter-ai package. + + The same server settings are shared between all providers. + Providers are not allowed to mutate this dictionary. + """ + @classmethod def chat_models(self): """Models which are suitable for chat.""" diff --git a/packages/jupyter-ai/jupyter_ai/extension.py b/packages/jupyter-ai/jupyter_ai/extension.py index 420b72ad..768ba7ae 100644 --- a/packages/jupyter-ai/jupyter_ai/extension.py +++ b/packages/jupyter-ai/jupyter_ai/extension.py @@ -1,11 +1,12 @@ import os import re import time +import types from dask.distributed import Client as DaskClient from importlib_metadata import entry_points from jupyter_ai.chat_handlers.learn import Retriever -from jupyter_ai_magics import JupyternautPersona +from jupyter_ai_magics import BaseProvider, JupyternautPersona from jupyter_ai_magics.utils import get_em_providers, get_lm_providers from jupyter_server.extension.application import ExtensionApp from tornado.web import StaticFileHandler @@ -202,6 +203,16 @@ def initialize_settings(self): defaults=defaults, ) + # Expose a subset of settings as read-only to the providers + exposed_server_settings = ["jupyter_server_ydoc"] + BaseProvider.server_settings = types.MappingProxyType( + { + key: value + for key, value in self.settings.items() + if key in exposed_server_settings + } + ) + self.log.info("Registered providers.") self.log.info(f"Registered {self.name} server extension") From 02d5ca2974d64a57e397ce3c4a5b8cb96dad1611 Mon Sep 17 00:00:00 2001 From: krassowski <5832902+krassowski@users.noreply.github.com> Date: Sun, 26 May 2024 17:48:08 +0100 Subject: [PATCH 3/5] Do not filter the settings object just yet Depending on the order in which extensions are initialized `jupyter_server_ydoc` may or may not be there; instead we may want to define a custom proxy which would restrict access on access to `__getitem__()` --- packages/jupyter-ai/jupyter_ai/extension.py | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/packages/jupyter-ai/jupyter_ai/extension.py b/packages/jupyter-ai/jupyter_ai/extension.py index 768ba7ae..c3661c11 100644 --- a/packages/jupyter-ai/jupyter_ai/extension.py +++ b/packages/jupyter-ai/jupyter_ai/extension.py @@ -204,14 +204,7 @@ def initialize_settings(self): ) # Expose a subset of settings as read-only to the providers - exposed_server_settings = ["jupyter_server_ydoc"] - BaseProvider.server_settings = types.MappingProxyType( - { - key: value - for key, value in self.settings.items() - if key in exposed_server_settings - } - ) + BaseProvider.server_settings = types.MappingProxyType(self.settings) self.log.info("Registered providers.") From 34d71679b5d067f7c563f5278d2c8b2476915bbb Mon Sep 17 00:00:00 2001 From: krassowski <5832902+krassowski@users.noreply.github.com> Date: Sun, 26 May 2024 18:01:53 +0100 Subject: [PATCH 4/5] `jupyter_server_ydoc` is accessible in web_app.settings --- packages/jupyter-ai/jupyter_ai/extension.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/packages/jupyter-ai/jupyter_ai/extension.py b/packages/jupyter-ai/jupyter_ai/extension.py index c3661c11..efb83b3d 100644 --- a/packages/jupyter-ai/jupyter_ai/extension.py +++ b/packages/jupyter-ai/jupyter_ai/extension.py @@ -204,7 +204,9 @@ def initialize_settings(self): ) # Expose a subset of settings as read-only to the providers - BaseProvider.server_settings = types.MappingProxyType(self.settings) + BaseProvider.server_settings = types.MappingProxyType( + self.serverapp.web_app.settings + ) self.log.info("Registered providers.") From aa3436f1ba369c301485c75e7120eac93d07e806 Mon Sep 17 00:00:00 2001 From: krassowski <5832902+krassowski@users.noreply.github.com> Date: Sun, 26 May 2024 18:02:23 +0100 Subject: [PATCH 5/5] Add missing async/await --- docs/source/developers/index.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/source/developers/index.md b/docs/source/developers/index.md index 6b0bd4c5..6763486d 100644 --- a/docs/source/developers/index.md +++ b/docs/source/developers/index.md @@ -263,7 +263,7 @@ class MyCompletionProvider(BaseProvider, FakeListLLM): kwargs["responses"] = ["This fake response will not be used for completion"] super().__init__(**kwargs) - def _get_prefix_and_suffix(self, request: InlineCompletionRequest): + async def _get_prefix_and_suffix(self, request: InlineCompletionRequest): prefix = request.prefix suffix = request.suffix.strip() @@ -273,7 +273,7 @@ class MyCompletionProvider(BaseProvider, FakeListLLM): return prefix, suffix is_notebook = request.path.endswith("ipynb") - document = server_ydoc.get_document( + 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" @@ -305,7 +305,7 @@ class MyCompletionProvider(BaseProvider, FakeListLLM): return prefix, suffix async def generate_inline_completions(self, request: InlineCompletionRequest): - prefix, suffix = self._get_prefix_and_suffix(request) + prefix, suffix = await self._get_prefix_and_suffix(request) return InlineCompletionReply( list=InlineCompletionList(items=[