Missing/incorrect formatting in the Jupyternaut response

[sundaraa-deshaw]

Description

I am using the jupyter ai extension with a custom model provider as per steps in https://jupyter-ai.readthedocs.io/en/latest/users/index.html#custom-model-providers

I have not modified the output format in the custom model provider, the provider looks like this:

class DESCOLLMProvider(BaseProvider, Llama2Chat):
    id = "desco_llama_provider"
    name = "desco llama provider"
    model_id_key = "model_id"
    models = [
        "codellama-7B-Instruct",
    ]

    def __init__(self, **kwargs):
        llm = LlamaCpp(
            model_path="/path/to/gguf/mode/file",
            temperature=0.2,
            max_tokens=640,
            top_p=0.0,
            verbose=True,  # Verbose is required to pass to the callback manager
            n_ctx=1024,
            top_k=1,
            n_gpu_layers=100,
            streaming=False
        )
        super().__init__(model_id="desco_llama_provider", llm=llm)

Reproduce

1. Select the registered custom model
2. Provide it with different inputs
3. After few Q&A iterations, the formatting of the response and subsequent response is garbled.

Actual response:

Once I clear the bot history, change model and revert the model back to the custom model, the formatting looks correct (which should be the ideal expected output). Now:

Expected behavior

Output format to look like:

Context

• Operating System and version: Fedora Linux 8
• Browser and version: Chrome build 119.0.6045.160
• JupyterLab version: 4.0.9

Generated prompt for the correct rendering

Prompt after formatting:
System: You are Jupyternaut, a conversational assistant living in JupyterLab to help users.
You are not a language model, but rather an application built on a foundation model from desco llama provider called desco_llama_provider.
You are talkative and you provide lots of specific details from the foundation model's context.
You may use Markdown to format your response.
Code blocks must be formatted in Markdown.
Math should be rendered with inline TeX markup, surrounded by $.
If you do not know the answer to a question, answer truthfully by responding that you do not know.
The following is a friendly conversation between you and a human.
Human: given a pandas dataframe with columns id, name, salary, write code to find the name with Nth highest salary

AI:   To find the name with the Nth highest salary in a pandas dataframe, you can use the `nlargest` method.
Here is an example:

import pandas as pd
# create sample data
data = {'id': [1, 2, 3], 'name': ['Alice', 'Bob', 'Charlie'], 'salary': [5000, 60000, 40000]}}

# find the name with the Nth highest salary

N = 2 # find the name with the 2nd highest salary

result = data.nlargest(N, 'salary'), on='id')

print(result))
Human: given a pandas dataframe with columns id, name, salary, write code to find the name with Nth highest salary

AI:   Here is some sample code that demonstrates how to find the name with the Nth highest salary in a pandas DataFrame:
import pandas as pd
# create sample data
data = {'id': [1, 2, 3], 'name': ['Alice', 'Bob', 'Charlie'], 'salary': [5000, 60000, 400000]}}

# find the name with the Nth highest salary

N = 2 # find the name with the 2nd highest salary

result = data.nlargest(N, 'salary'), on='id')

print(result))
Human: write code to make a http POST call to example101.com
AI:   Here is an example of how you can make a HTTP POST call using the `requests` library in Python:
```
import requests
url = "https://example101.com"
payload = {
    "key": "value",
}

headers = {
    'Content-Type': 'application/json',
}

response = requests.post(url, data=payload, headers=headers))
print(response.text))
```

Generated prompt for the incorrect rendering

> Entering new ConversationChain chain...
Prompt after formatting:
System: You are Jupyternaut, a conversational assistant living in JupyterLab to help users.
You are not a language model, but rather an application built on a foundation model from desco llama provider called desco_llama_provider.
You are talkative and you provide lots of specific details from the foundation model's context.
You may use Markdown to format your response.
Code blocks must be formatted in Markdown.
Math should be rendered with inline TeX markup, surrounded by $.
If you do not know the answer to a question, answer truthfully by responding that you do not know.
The following is a friendly conversation between you and a human.
Human: convert numpy array to a dataframe 
AI:   You can convert a NumPy array to a pandas DataFrame using the `to_df()` method of the NumPy array.
Here is an example:
import numpy as np
# create a NumPy array
arr = np.array([1,2],[3,4])
# convert the NumPy array to a pandas DataFrame
df = arr.to_df()
print(df))
Human: write code to make a http POST call to example101.com
AI:   Here is an example of how you can make a HTTP POST call to example101.com using Python:
import requests
url = "https://example101.com"
payload = {
    "key": "value",
}

response = requests.post(url, data=payload))
print(response.text))

As we can see, "```" is missing around the code snippet in the response rendered in the second prompt.

Is this expected?

[sundaraa-deshaw]

This happens with %ai cell magic command output as well

[sundaraa-deshaw]

I guess the Python comments "#" are treated as markdown headers and rendered. Is there a way to fix this by configuring the model provider? I am using the BaseChatModel here.

[JasonWeill]

@sundaraa-deshaw Sorry for the delay in responding to this! Typically, we expect language models to respond with code enclosed in triple backticks (``) on lines by themselves, perhaps with a language indicator on the first line (for example,python` after the triple backtick).

Is your custom model providing code blocks surrounded with triple backticks? If not, Jupyter AI won't be able to differentiate between Python code, where # foo indicates a comment, and Markdown, where # foo indicates a heading.