Fix JSON serialization error in Ollama models

[JanusChoi]

A slightly changes in order to fix: #1128

issue: When chat with Jupyternaut, getting error:

[D 2024-11-30 07:52:05.120 ServerApp] Broadcasting message: type='agent-stream-chunk' id='0a3f8b24522a46c9a694a8e2493e561b' content='' stream_complete=True metadata={'model': 'deepseek-v2:16b', 'created_at': '2024-11-29T23:52:04.8986414Z', 'done': True, 'done_reason': 'stop', 'total_duration': 4248175500, 'load_duration': 3461159000, 'prompt_eval_count': 265, 'prompt_eval_duration': 345000000, 'eval_count': 20, 'eval_duration': 322000000, 'message': Message(role='assistant', content='', images=None, tool_calls=None)} to all clients...
[E 2024-11-30 07:52:05.120 AiExtension] Object of type Message is not JSON serializable
[D 2024-11-30 07:52:05.204 ServerApp] Broadcasting message: id='0b38a635d3a34f0da5b5bb23e0a789d6' time=1732924325.204742 body='Sorry, an error occurred. Details below:\n\n```\nTraceback (most recent call last):\n  File "D:\\miniconda3\\envs\\jupyter-ai\\Lib\\site-packages\\jupyter_ai\\chat_handlers\\base.py", line 226, in on_message\n    await self.process_message(message)\n  File "D:\\miniconda3\\envs\\jupyter-ai\\Lib\\site-packages\\jupyter_ai\\chat_handlers\\default.py", line 71, in process_message\n    await self.stream_reply(inputs, message)\n  File "D:\\miniconda3\\envs\\jupyter-ai\\Lib\\site-packages\\jupyter_ai\\chat_handlers\\base.py", line 603, in stream_reply\n    self._send_stream_chunk(\n  File "D:\\miniconda3\\envs\\jupyter-ai\\Lib\\site-packages\\jupyter_ai\\chat_handlers\\base.py", line 518, in _send_stream_chunk\n    self.broadcast_message(stream_chunk_msg)\n  File "D:\\miniconda3\\envs\\jupyter-ai\\Lib\\site-packages\\jupyter_ai\\chat_handlers\\base.py", line 285, in broadcast_message\n    websocket.broadcast_message(message)\n  File "D:\\miniconda3\\envs\\jupyter-ai\\Lib\\site-packages\\jupyter_ai\\handlers.py", line 241, in broadcast_message\n    client.write_message(message.dict())\n  File "D:\\miniconda3\\envs\\jupyter-ai\\Lib\\site-packages\\tornado\\websocket.py", line 334, in write_message\n    message = tornado.escape.json_encode(message)\n              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n  File "D:\\miniconda3\\envs\\jupyter-ai\\Lib\\site-packages\\tornado\\escape.py", line 96, in json_encode\n    return json.dumps(value).replace("</", "<\\\\/")\n           ^^^^^^^^^^^^^^^^^\n  File "D:\\miniconda3\\envs\\jupyter-ai\\Lib\\json\\__init__.py", line 231, in dumps\n    return _default_encoder.encode(obj)\n           
^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n  File "D:\\miniconda3\\envs\\jupyter-ai\\Lib\\json\\encoder.py", line 200, in encode\n    chunks = self.iterencode(o, _one_shot=True)\n           
  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n  File "D:\\miniconda3\\envs\\jupyter-ai\\Lib\\json\\encoder.py", line 258, in iterencode\n    return _iterencode(o, 0)\n           ^^^^^^^^^^^^^^^^^\n  File "D:\\miniconda3\\envs\\jupyter-ai\\Lib\\json\\encoder.py", line 180, in default\n    raise TypeError(f\'Object of type {o.__class__.__name__} \'\nTypeError: Object of type Message is not JSON serializable\n\n```' reply_to='8bc8ec45-6453-40e3-8dc1-f3e1f1f40227' persona=Persona(name='Jupyternaut', avatar_route='api/ai/static/jupyternaut.svg') metadata={} type='agent' to all clients...
[I 2024-11-30 07:52:05.206 ServerApp] Default chat handler resolved in 4754 ms.

The error occurs when trying to serialize a Message object to JSON. Looking at the code, the Message class is defined as a Union type of several message classes:

Message = Union[
   ChatMessage,
   ConnectionMessage,
   ClearMessage,
   PendingMessage,
   ClosePendingMessage,
]

The error occurs because the metadata field in the message contains a Message object, which is not directly JSON serializable.

The issue is that the metadata field is defined as Dict[str, Any], which means it can contain any type of value. In this case, it appears to contain a Message object which can't be automatically serialized to JSON.

Edit the AgentStreamChunkMessage class to ensure metadata values are JSON serializable:

site-packages\jupyter_ai\models.py

import json

class AgentStreamChunkMessage(BaseModel):
    ## {...}

    @validator("metadata")
    def validate_metadata(cls, v):
        """Ensure metadata values are JSON serializable"""
        try:
            json.dumps(v)
            return v
        except TypeError as e:
            raise ValueError(f"Metadata must be JSON serializable: {str(e)}")

The metadata is being populated from the LLMResult's generation_info dictionary, which may contain objects that are not JSON serializable. We need to ensure that any non-serializable objects in the metadata are converted to a serializable format before being assigned.

So modify the MetadataCallbackHandler to handle this:

site-packages\jupyter_ai\callback_handlers\metadata.py

from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.outputs import LLMResult
import json


def convert_to_serializable(obj):
    """Convert an object to a JSON serializable format"""
    if hasattr(obj, 'dict') and callable(obj.dict):
        return obj.dict()
    if hasattr(obj, '__dict__'):
        return obj.__dict__
    return str(obj)


class MetadataCallbackHandler(BaseCallbackHandler):
    """
    When passed as a callback handler, this stores the LLMResult's
    `generation_info` dictionary in the `self.jai_metadata` instance attribute
    after the provider fully processes an input.

    If used in a streaming chat handler: the `metadata` field of the final
    `AgentStreamChunkMessage` should be set to `self.jai_metadata`.

    If used in a non-streaming chat handler: the `metadata` field of the
    returned `AgentChatMessage` should be set to `self.jai_metadata`.
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.jai_metadata = {}

    def on_llm_end(self, response: LLMResult, **kwargs) -> None:
        if not (len(response.generations) and len(response.generations[0])):
            return

        metadata = response.generations[0][0].generation_info or {}
        
        # Convert any non-serializable objects in metadata
        serializable_metadata = {}
        for key, value in metadata.items():
            try:
                json.dumps(value)
                serializable_metadata[key] = value
            except (TypeError, ValueError):
                serializable_metadata[key] = convert_to_serializable(value)
        
        self.jai_metadata = serializable_metadata

[krassowski]

Indeed, it looks like any arbitrary value may be present in generation_data (so it is not an issue with Ollama!) as per definition of this field in langchain:

    generation_info: Optional[dict[str, Any]] = None
    """Raw response from the provider.

    May include things like the reason for finishing or token log probabilities.
    """

[dlqqq]

@JanusChoi Thank you for contributing this fix so quickly! Looks like the fix works, per Sanjiv's testing. Awesome work. 🎉

Can you help modify this PR according to my review below? I've included a summary of @krassowski's review in mine. One of our contributors can help if you lack the time. 👍

[dlqqq]

I can help with these changes proposed above so we can include this in the next patch release. Working on this now.

[srdas]

Tested this with chat and RAG to make sure there is no error.

Also the function to convert to serializable works, as can be seen

All else looks good. Thanks so much @dlqqq (and @JanusChoi @krassowski) for this.

[dlqqq]

@meeseeksdev please backport to v3-dev