[Bug]: Unable to use `json` on chat response

[JuHyung-Son]

Bug Description

chat response has Choice and Usage can be json serialize

Version

any version

Steps to Reproduce

response = llm.chat(messages=[
  ChatMessage(role="system", content="You are a helpful assistant."),
  ChatMessage(role="user", content="Hi, how are you?")
])

print(response)

{'id': 'chatcmpl-9hTPLBEbQZ2FCGSBMkrrBxR3VvESE',
 'choices': [Choice(finish_reason='stop', index=0, logprobs=None, message=ChatCompletionMessage(content="Hello! I'm just a computer program, so I don't have feelings, but I'm here and ready to help you. How can I assist you today?", role='assistant', function_call=None, tool_calls=None))],
 'created': 1720148175,
 'model': 'gpt-3.5-turbo-0125',
 'object': 'chat.completion',
 'system_fingerprint': None,
 'usage': CompletionUsage(completion_tokens=33, prompt_tokens=23, total_tokens=56)}

response.json()

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[53], line 1
----> 1 response.json()

File [~/upstage/projects/upstage-cookbook/.venv/lib/python3.10/site-packages/pydantic/v1/main.py:504](http://localhost:8888/lab/workspaces/auto-T/tree/~/upstage/projects/upstage-cookbook/.venv/lib/python3.10/site-packages/pydantic/v1/main.py#line=503), in BaseModel.json(self, include, exclude, by_alias, skip_defaults, exclude_unset, exclude_defaults, exclude_none, encoder, models_as_dict, **dumps_kwargs)
    502 if self.__custom_root_type__:
    503     data = data[ROOT_KEY]
--> 504 return self.__config__.json_dumps(data, default=encoder, **dumps_kwargs)

File [/usr/local/Cellar/python](http://localhost:8888/usr/local/Cellar/python)@3.10[/3.10.13_2/Frameworks/Python.framework/Versions/3.10/lib/python3.10/json/__init__.py:238](http://localhost:8888/3.10.13_2/Frameworks/Python.framework/Versions/3.10/lib/python3.10/json/__init__.py#line=237), in dumps(obj, skipkeys, ensure_ascii, check_circular, allow_nan, cls, indent, separators, default, sort_keys, **kw)
    232 if cls is None:
    233     cls = JSONEncoder
    234 return cls(
    235     skipkeys=skipkeys, ensure_ascii=ensure_ascii,
    236     check_circular=check_circular, allow_nan=allow_nan, indent=indent,
    237     separators=separators, default=default, sort_keys=sort_keys,
--> 238     **kw).encode(obj)

File [/usr/local/Cellar/python](http://localhost:8888/usr/local/Cellar/python)@3.10[/3.10.13_2/Frameworks/Python.framework/Versions/3.10/lib/python3.10/json/encoder.py:199](http://localhost:8888/3.10.13_2/Frameworks/Python.framework/Versions/3.10/lib/python3.10/json/encoder.py#line=198), in JSONEncoder.encode(self, o)
    195         return encode_basestring(o)
    196 # This doesn't pass the iterator directly to ''.join() because the
    197 # exceptions aren't as detailed.  The list call should be roughly
    198 # equivalent to the PySequence_Fast that ''.join() would do.
--> 199 chunks = self.iterencode(o, _one_shot=True)
    200 if not isinstance(chunks, (list, tuple)):
    201     chunks = list(chunks)

File [/usr/local/Cellar/python](http://localhost:8888/usr/local/Cellar/python)@3.10[/3.10.13_2/Frameworks/Python.framework/Versions/3.10/lib/python3.10/json/encoder.py:257](http://localhost:8888/3.10.13_2/Frameworks/Python.framework/Versions/3.10/lib/python3.10/json/encoder.py#line=256), in JSONEncoder.iterencode(self, o, _one_shot)
    252 else:
    253     _iterencode = _make_iterencode(
    254         markers, self.default, _encoder, self.indent, floatstr,
    255         self.key_separator, self.item_separator, self.sort_keys,
    256         self.skipkeys, _one_shot)
--> 257 return _iterencode(o, 0)

File [~/upstage/projects/upstage-cookbook/.venv/lib/python3.10/site-packages/pydantic/v1/json.py:90](http://localhost:8888/lab/workspaces/auto-T/tree/~/upstage/projects/upstage-cookbook/.venv/lib/python3.10/site-packages/pydantic/v1/json.py#line=89), in pydantic_encoder(obj)
     88     return encoder(obj)
     89 else:  # We have exited the for loop without finding a suitable encoder
---> 90     raise TypeError(f"Object of type '{obj.__class__.__name__}' is not JSON serializable")

TypeError: Object of type 'Choice' is not JSON serializable

Relevant Logs/Tracbacks

No response

[dosubot[bot]]

To serialize a chat response containing Choice and Usage objects to JSON without encountering a TypeError, you need to implement a custom JSON encoder. Here is an example of how you can do this:

import json
from typing import Any

class CustomJSONEncoder(json.JSONEncoder):
    def default(self, obj: Any) -> Any:
        if isinstance(obj, Choice):
            return obj.__dict__
        if isinstance(obj, Usage):
            return obj.__dict__
        return super().default(obj)

# Assuming `response` is the chat response object containing `Choice` and `Usage` objects
response_json = json.dumps(response, cls=CustomJSONEncoder)
print(response_json)

In this example, CustomJSONEncoder converts Choice and Usage objects to dictionaries using their __dict__ attribute, making them JSON serializable ^[1].

To continue talking to Dosu, mention @dosu.

[logan-markewich]

There's a similar fix applied for ChatMessage class that we could borrow to fix this

[logan-markewich]

(But also, an easy workaround is just taking the fields you want)

[houtaroy]

First, I don't recommend directly calling response.json(). Its return value is LlamaIndex's ChatResponse, which isn't the best choice if you want to log something like an interface log.

Secondly, if it's necessary, you can implement it yourself by encoding, like this:

import os
import json
from typing import Any

from llama_index.core.base.llms.types import ChatMessage
from llama_index.llms.openai import OpenAI
from openai.types.chat.chat_completion import Choice
from openai.types.completion_usage import CompletionUsage

def openai_json_encoder(obj: Any) -> Any:
    if isinstance(obj, (Choice, CompletionUsage)):
        return obj.model_dump()
    return json.JSONEncoder().default(obj)

llm = OpenAI()

response = llm.chat(messages=[
  ChatMessage(role="system", content="You are a helpful assistant."),
  ChatMessage(role="user", content="Hi, how are you?")
])

response.json(encoder=openai_json_encoder)

Lastly, if you have any other questions, we can continue to discuss.

@JuHyung-Son @logan-markewich

[houtaroy]

Of course, we can extend from ChatResponse, implement OpenAIChatResponse, and then use pydantic config like this:

class OpenAIChatResponse(ChatResponse):

    class Config:
        json_encoders = {
            Choice: lambda x: x.dict(),
            CompletionUsage: lambda x: x.dict()
        }

However, I don't think this is necessary.

[JuHyung-Son]

Okay I understand. Thanks. But if the json method actually does not work and it is not recommended to use, Do we even need a json method?

[houtaroy]

The method json is BaseModel's method in pydantic, My first reaction was to add a global converter like Spring Framework in Java