Add request number and logging (do not use print), fix typos

krassowski · krassowski · commit 00bb6f3d32da · 2025-03-25T14:11:27.000Z
Add test for LLM completion

Add `@pytest.mark.asyncio` decorators
diff --git a/IPython/terminal/shortcuts/auto_suggest.py b/IPython/terminal/shortcuts/auto_suggest.py
@@ -185,6 +185,7 @@ def __init__(self):
         self.skip_lines = 0
         self._connected_apps = []
         self._llm_provider = None
+        self._request_number = 0
 
     def reset_history_position(self, _: Buffer) -> None:
         self.skip_lines = 0
@@ -350,7 +351,7 @@ async def error_catcher(buffer):
             try:
                 await self._trigger_llm_core(buffer)
             except Exception as e:
-                get_ipython().log.error("error")
+                get_ipython().log.error(f"error {e}")
                 raise
 
         # here we need a cancellable task so we can't just await the error catched
@@ -365,9 +366,8 @@ async def _trigger_llm_core(self, buffer: Buffer):
         provider to stream it's response back to us iteratively setting it as
         the suggestion on the current buffer.
 
-        Unlike with JupyterAi, as we do not have multiple cell, the cell number
-        is always set to `0`, note that we _could_ set it to a new number each
-        time and ignore threply from past numbers.
+        Unlike with JupyterAi, as we do not have multiple cell, the cell id
+        is always set to `None`.
 
         We set the prefix to the current cell content, but could also inset the
         rest of the history or even just the non-fail history.
@@ -389,10 +389,12 @@ async def _trigger_llm_core(self, buffer: Buffer):
 
         hm = buffer.history.shell.history_manager
         prefix = self._llm_prefixer(hm)
-        print(prefix)
+        get_ipython().log.debug(f"prefix: {prefix}")
 
+        self._request_number += 1
+        request_number = self._request_number
         request = jai_models.InlineCompletionRequest(
-            number=0,
+            number=request_number,
             prefix=prefix + buffer.document.text,
             suffix="",
             mime="text/x-python",
@@ -405,6 +407,9 @@ async def _trigger_llm_core(self, buffer: Buffer):
         async for reply_and_chunks in self._llm_provider.stream_inline_completions(
             request
         ):
+            if self._request_number != request_number:
+                # If a new suggestion was requested, skip processing this one.
+                return
             if isinstance(reply_and_chunks, jai_models.InlineCompletionReply):
                 if len(reply_and_chunks.list.items) > 1:
                     raise ValueError(
diff --git a/pyproject.toml b/pyproject.toml
@@ -77,6 +77,7 @@ test = [
 test_extra = [
     "ipython[test]",
     "curio",
+    "jupyter_ai",
     "matplotlib!=3.2.0",
     "nbformat",
     "numpy>=1.23",
diff --git a/tests/fake_llm.py b/tests/fake_llm.py
@@ -0,0 +1,88 @@
+import asyncio
+
+try:
+    from jupyter_ai_magics import BaseProvider
+    from langchain_community.llms import FakeListLLM
+except ImportError:
+
+    class BaseProvider:
+        pass
+
+    class FakeListLLM:
+        pass
+
+
+FIBONACCI = """\
+def fib(n):
+    if n < 2: return n
+    return fib(n - 1) + fib(n - 2)
+"""
+
+
+class FibonacciCompletionProvider(BaseProvider, FakeListLLM):  # type: ignore[misc, valid-type]
+
+    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"]
+        kwargs["model_id"] = "model_a"
+        super().__init__(**kwargs)
+
+    async def generate_inline_completions(self, request):
+        raise ValueError("IPython only supports streaming models.")
+
+    async def stream_inline_completions(self, request):
+        from jupyter_ai.completions.models import (
+            InlineCompletionList,
+            InlineCompletionReply,
+            InlineCompletionStreamChunk,
+        )
+
+        assert request.number > 0
+        token = f"t{request.number}s0"
+        last_line = request.prefix.rstrip("\n").splitlines()[-1]
+
+        if not FIBONACCI.startswith(last_line):
+            return
+
+        yield InlineCompletionReply(
+            list=InlineCompletionList(
+                items=[
+                    {"insertText": "", "isIncomplete": True, "token": token},
+                ]
+            ),
+            reply_to=request.number,
+        )
+
+        async for reply in self._stream(
+            FIBONACCI[len(last_line) :],
+            request.number,
+            token,
+        ):
+            yield reply
+
+    async def _stream(self, sentence, request_number, token, start_with=""):
+        suggestion = start_with
+
+        for fragment in sentence.split(" "):
+            await asyncio.sleep(0.05)
+            if suggestion:
+                suggestion += " "
+            suggestion += fragment
+            yield InlineCompletionStreamChunk(
+                type="stream",
+                response={"insertText": suggestion, "token": token},
+                reply_to=request_number,
+                done=False,
+            )
+
+        # finally, send a message confirming that we are done
+        yield InlineCompletionStreamChunk(
+            type="stream",
+            response={"insertText": suggestion, "token": token},
+            reply_to=request_number,
+            done=True,
+        )
diff --git a/tests/test_shortcuts.py b/tests/test_shortcuts.py
@@ -7,12 +7,14 @@
     accept_word,
     accept_and_keep_cursor,
     discard,
+    llm_autosuggestion,
     NavigableAutoSuggestFromHistory,
     swap_autosuggestion_up,
     swap_autosuggestion_down,
 )
 from IPython.terminal.shortcuts.auto_match import skip_over
 from IPython.terminal.shortcuts import create_ipython_shortcuts, reset_search_buffer
+from IPython.testing import decorators as dec
 
 from prompt_toolkit.history import InMemoryHistory
 from prompt_toolkit.buffer import Buffer
@@ -34,6 +36,26 @@ def make_event(text, cursor, suggestion):
     return event
 
 
+try:
+    from .fake_llm import FIBONACCI
+except ImportError:
+    FIBONACCI = None
+
+
+@dec.skip_without("jupyter_ai")
+@pytest.mark.asyncio
+async def test_llm_autosuggestion():
+    provider = NavigableAutoSuggestFromHistory()
+    ip = get_ipython()
+    ip.auto_suggest = provider
+    ip.llm_provider_class = "tests.fake_llm.FibonacciCompletionProvider"
+    text = "def fib"
+    event = make_event(text, len(text), "")
+    event.current_buffer.history.shell.history_manager.get_range = Mock(return_value=[])
+    await llm_autosuggestion(event)
+    assert event.current_buffer.suggestion.text == FIBONACCI[len(text) :]
+
+
 @pytest.mark.parametrize(
     "text, suggestion, expected",
     [
@@ -219,6 +241,7 @@ def test_other_providers():
     assert swap_autosuggestion_down(event) is None
 
 
+@pytest.mark.asyncio
 async def test_navigable_provider():
     provider = NavigableAutoSuggestFromHistory()
     history = InMemoryHistory(history_strings=["very_a", "very", "very_b", "very_c"])
@@ -271,6 +294,7 @@ def get_suggestion():
     assert get_suggestion().text == "_a"
 
 
+@pytest.mark.asyncio
 async def test_navigable_provider_multiline_entries():
     provider = NavigableAutoSuggestFromHistory()
     history = InMemoryHistory(history_strings=["very_a\nvery_b", "very_c"])
