Trim leading whitespace when processing (jupyterlab#900)

jtpio · Jul 17, 2024 · 12d069e · 12d069e
1 parent 83c0cf8
commit 12d069e
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/completion_utils.py b/packages/jupyter-ai-magics/jupyter_ai_magics/completion_utils.py
@@ -38,8 +38,9 @@ def post_process_suggestion(suggestion: str, request: InlineCompletionRequest) -
         for identifier in markdown_identifiers.get(language, [language])
     ] + ["```"]
     for opening in bad_openings:
-        if suggestion.startswith(opening):
-            suggestion = suggestion[len(opening) :].lstrip()
+        # ollama models tend to add spurious whitespace
+        if suggestion.lstrip().startswith(opening):
+            suggestion = suggestion.lstrip()[len(opening) :].lstrip()
             # check for the prefix inclusion (only if there was a bad opening)
             if suggestion.startswith(request.prefix):
                 suggestion = suggestion[len(request.prefix) :]

diff --git a/packages/jupyter-ai/jupyter_ai/tests/completions/test_handlers.py b/packages/jupyter-ai/jupyter_ai/tests/completions/test_handlers.py
@@ -102,6 +102,8 @@ async def test_handle_request(inline_handler):
     ("```python\nTest python code\n```", "Test python code"),
     ("```\ntest\n```\n   \n", "test"),
     ("```hello```world```", "hello```world"),
+    (" ```\nprint(test)\n```", "print(test)"),
+    ("``` \nprint(test)\n```", "print(test)"),
 ]