Ollama models return JSON serializing error

[srdas]

Description

Ollama models generate error after returning a valid response:

Reproduce

Choose any Ollama model (the error is not specific to a particular Ollama model):

Then try any prompt. It will return a correct response in one exchange and will follow it up with an unnecessary error as shown above. Error is:

Traceback (most recent call last):
  File "/Users/sanjivda/GitHub/jupyter-ai/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py", line 226, in on_message
    await self.process_message(message)
  File "/Users/sanjivda/GitHub/jupyter-ai/packages/jupyter-ai/jupyter_ai/chat_handlers/default.py", line 71, in process_message
    await self.stream_reply(inputs, message)
  File "/Users/sanjivda/GitHub/jupyter-ai/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py", line 603, in stream_reply
    self._send_stream_chunk(
  File "/Users/sanjivda/GitHub/jupyter-ai/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py", line 518, in _send_stream_chunk
    self.broadcast_message(stream_chunk_msg)
  File "/Users/sanjivda/GitHub/jupyter-ai/packages/jupyter-ai/jupyter_ai/chat_handlers/base.py", line 285, in broadcast_message
    websocket.broadcast_message(message)
  File "/Users/sanjivda/GitHub/jupyter-ai/packages/jupyter-ai/jupyter_ai/handlers.py", line 241, in broadcast_message
    client.write_message(message.dict())
  File "/opt/anaconda3/envs/jupyter-ai/lib/python3.12/site-packages/tornado/websocket.py", line 334, in write_message
    message = tornado.escape.json_encode(message)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/jupyter-ai/lib/python3.12/site-packages/tornado/escape.py", line 96, in json_encode
    return json.dumps(value).replace("</", "<\\/")
           ^^^^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/jupyter-ai/lib/python3.12/json/__init__.py", line 231, in dumps
    return _default_encoder.encode(obj)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/jupyter-ai/lib/python3.12/json/encoder.py", line 200, in encode
    chunks = self.iterencode(o, _one_shot=True)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/jupyter-ai/lib/python3.12/json/encoder.py", line 258, in iterencode
    return _iterencode(o, 0)
           ^^^^^^^^^^^^^^^^^
  File "/opt/anaconda3/envs/jupyter-ai/lib/python3.12/json/encoder.py", line 180, in default
    raise TypeError(f'Object of type {o.__class__.__name__} '
TypeError: Object of type Message is not JSON serializable

Expected behavior

No JSON serializing error

Environment

Jupyter AI v2.28.2

[edwardochoaphd]

getting same error