Repository: jupyterlab/jupyter-ai

Issue:
Prefix is not removed from inline completions when streaming
## Description

Using GPT 4 from Open AI:

No streaming:

![Screenshot from 2024-06-19 16-05-04](https://github.com/jupyterlab/jupyter-ai/assets/5832902/2bdca6c1-fa8b-4743-9896-f286264ef225)

Streaming:

![Screenshot from 2024-06-19 16-05-23](https://github.com/jupyterlab/jupyter-ai/assets/5832902/3b0814c9-6ee1-4143-a72f-43293cff6719)


## Reproduce

1. In inline completer settings change "streaming" to "Always"
2. Attempt to type something
3. See the prefix repeated

## Expected behavior

No prefix.

## Context
 
2.17.0


Discussion:
For bulk completion the post-processing method is always applied:

https://github.com/jupyterlab/jupyter-ai/blob/02d196685514b809d565e14e0e2dca7784706226/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py#L452-L458

For streaming it is applied conditionally:

https://github.com/jupyterlab/jupyter-ai/blob/02d196685514b809d565e14e0e2dca7784706226/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py#L488-L497

Some more refined logic is needed so that the prefix which may be added after opening ``` is also trimmed. It needs to keep track if trimming was performed to avoid removing the prefix-like string second time.



Patch:
diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py
index 8c8a1f917..39b1b8adf 100644
--- a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py
+++ b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py
@@ -485,7 +485,7 @@ async def stream_inline_completions(
         chain = self._create_completion_chain()
         token = completion.token_from_request(request, 0)
         model_arguments = completion.template_inputs_from_request(request)
-        suggestion = ""
+        suggestion = processed_suggestion = ""
 
         # send an incomplete `InlineCompletionReply`, indicating to the
         # client that LLM output is about to streamed across this connection.
@@ -505,17 +505,12 @@ async def stream_inline_completions(
 
         async for fragment in chain.astream(input=model_arguments):
             suggestion += fragment
-            if suggestion.startswith("```"):
-                if "\n" not in suggestion:
-                    # we are not ready to apply post-processing
-                    continue
-                else:
-                    suggestion = completion.post_process_suggestion(suggestion, request)
-            elif suggestion.rstrip().endswith("```"):
-                suggestion = completion.post_process_suggestion(suggestion, request)
+            processed_suggestion = completion.post_process_suggestion(
+                suggestion, request
+            )
             yield InlineCompletionStreamChunk(
                 type="stream",
-                response={"insertText": suggestion, "token": token},
+                response={"insertText": processed_suggestion, "token": token},
                 reply_to=request.number,
                 done=False,
             )
@@ -523,7 +518,7 @@ async def stream_inline_completions(
         # finally, send a message confirming that we are done
         yield InlineCompletionStreamChunk(
             type="stream",
-            response={"insertText": suggestion, "token": token},
+            response={"insertText": processed_suggestion, "token": token},
             reply_to=request.number,
             done=True,
         )


Test patch:
diff --git a/packages/jupyter-ai/jupyter_ai/tests/completions/test_handlers.py b/packages/jupyter-ai/jupyter_ai/tests/completions/test_handlers.py
index c5b5d1eea..d87b75b9d 100644
--- a/packages/jupyter-ai/jupyter_ai/tests/completions/test_handlers.py
+++ b/packages/jupyter-ai/jupyter_ai/tests/completions/test_handlers.py
@@ -98,13 +98,16 @@ async def test_handle_request(inline_handler):
     assert suggestions[0].insertText == "Test response"
 
 
+expected_suggestions_cases = [
+    ("```python\nTest python code\n```", "Test python code"),
+    ("```\ntest\n```\n   \n", "test"),
+    ("```hello```world```", "hello```world"),
+]
+
+
 @pytest.mark.parametrize(
     "response,expected_suggestion",
-    [
-        ("```python\nTest python code\n```", "Test python code"),
-        ("```\ntest\n```\n   \n", "test"),
-        ("```hello```world```", "hello```world"),
-    ],
+    expected_suggestions_cases,
 )
 async def test_handle_request_with_spurious_fragments(response, expected_suggestion):
     inline_handler = MockCompletionHandler(
@@ -128,6 +131,32 @@ async def test_handle_request_with_spurious_fragments(response, expected_suggest
     assert suggestions[0].insertText == expected_suggestion
 
 
+@pytest.mark.parametrize(
+    "response,expected_suggestion",
+    expected_suggestions_cases,
+)
+async def test_handle_request_with_spurious_fragments_stream(
+    response, expected_suggestion
+):
+    inline_handler = MockCompletionHandler(
+        lm_provider=MockProvider,
+        lm_provider_params={
+            "model_id": "model",
+            "responses": [response],
+        },
+    )
+    dummy_request = InlineCompletionRequest(
+        number=1, prefix="", suffix="", mime="", stream=True
+    )
+
+    await inline_handler.handle_stream_request(dummy_request)
+    assert len(inline_handler.messages) == 3
+    # the streamed fragment should not include spurious fragments
+    assert inline_handler.messages[1].response.insertText == expected_suggestion
+    # the final state should not include spurious fragments either
+    assert inline_handler.messages[2].response.insertText == expected_suggestion
+
+
 async def test_handle_stream_request():
     inline_handler = MockCompletionHandler(
         lm_provider=MockProvider,