Inline completion strips out new line when it should be kept

[krassowski]

Description

New line is stripped out if model returns

```

prefix.

Reproduce

1. Write comment like # load json file
2. See suggestion gets applied like: # load json file import

Expected behavior

# load json file
import

Context

This is an issue in

jupyter-ai/packages/jupyter-ai-magics/jupyter_ai_magics/completion_utils.py

Lines 25 to 47 in fcb2d71

	def post_process_suggestion(suggestion: str, request: InlineCompletionRequest) -> str:
	"""Remove spurious fragments from the suggestion.
	While most models (especially instruct and infill models do not require
	any pre-processing, some models such as gpt-4 which only have chat APIs
	may require removing spurious fragments. This function uses heuristics
	and request data to remove such fragments.
	"""
	# gpt-4 tends to add "```python" or similar
	language = request.language or "python"
	markdown_identifiers = {"ipython": ["ipython", "python", "py"]}
	bad_openings = [
	f"```{identifier}"
	for identifier in markdown_identifiers.get(language, [language])
	] + ["```"]
	for opening in bad_openings:
	# 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) :]
	break

[krassowski]

To fix it we need to:

- suggestion = suggestion.lstrip()[len(opening) :].lstrip() 
+ suggestion = suggestion.lstrip()[len(opening) :]
  # check for the prefix inclusion (only if there was a bad opening) 
- if suggestion.startswith(request.prefix): 
-      suggestion = suggestion[len(request.prefix) :]
+ if suggestion.lstrip().startswith(request.prefix): 
+      suggestion = suggestion.lstrip()[len(request.prefix) :]

Plus tests cases need to be added in

jupyter-ai/packages/jupyter-ai/jupyter_ai/tests/completions/test_handlers.py

Lines 101 to 133 in fcb2d71

	expected_suggestions_cases = [
	("```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)"),
	]
	@pytest.mark.parametrize(
	"response,expected_suggestion",
	expected_suggestions_cases,
	)
	async def test_handle_request_with_spurious_fragments(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=False
	)
	await inline_handler.handle_request(dummy_request)
	# should write a single reply
	assert len(inline_handler.messages) == 1
	# reply should contain a single suggestion
	suggestions = inline_handler.messages[0].list.items
	assert len(suggestions) == 1
	# the suggestion should include insert text from LLM without spurious fragments
	assert suggestions[0].insertText == expected_suggestion