mirror of
https://github.com/open-webui/open-webui.git
synced 2026-09-17 23:52:29 +00:00
62 lines
2.5 KiB
Python
62 lines
2.5 KiB
Python
"""The app-wide JSON codec, selected by the ``ENABLE_ORJSON`` env var.
|
|
|
|
Every module that would otherwise reach for stdlib ``json`` imports ``JSONCodec``
|
|
from here, so the whole app switches implementation from a single flag. With the
|
|
flag off these are stdlib ``json`` and engineio's codec verbatim, so the default
|
|
behaviour is exactly what it was before orjson entered the picture.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json as stdlib_json
|
|
|
|
from engineio import json as engineio_json
|
|
from open_webui.env import ENABLE_ORJSON
|
|
|
|
if ENABLE_ORJSON:
|
|
import orjson
|
|
|
|
# stdlib escapes these, orjson emits them raw, and Python treats all three as
|
|
# line boundaries: one raw separator splits an SSE frame that a reader
|
|
# reassembles with ``splitlines()``.
|
|
LINE_SEPARATOR_ESCAPES = str.maketrans({'\u2028': '\\u2028', '\u2029': '\\u2029', '\x85': '\\u0085'})
|
|
|
|
class ORJSONCodec:
|
|
"""stdlib-``json``-compatible codec backed by orjson.
|
|
|
|
Anything orjson rejects (non-str dict keys, ints beyond 64 bits, ``NaN``
|
|
literals) falls back to engineio's stdlib-based codec, which keeps its
|
|
oversized-integer guard for untrusted client payloads.
|
|
"""
|
|
|
|
JSONDecodeError = engineio_json.JSONDecodeError
|
|
|
|
@staticmethod
|
|
def dumps(obj, *args, **kwargs):
|
|
# orjson can't honor stdlib json options; compact separators are its default.
|
|
if args or (kwargs and kwargs != {'separators': (',', ':')}):
|
|
return engineio_json.dumps(obj, *args, **kwargs)
|
|
try:
|
|
serialized = orjson.dumps(obj).decode('utf-8')
|
|
except (TypeError, ValueError):
|
|
return engineio_json.dumps(obj, *args, **kwargs)
|
|
if '\u2028' in serialized or '\u2029' in serialized or '\x85' in serialized:
|
|
return serialized.translate(LINE_SEPARATOR_ESCAPES)
|
|
return serialized
|
|
|
|
@staticmethod
|
|
def loads(s, *args, **kwargs):
|
|
if args or kwargs:
|
|
return engineio_json.loads(s, *args, **kwargs)
|
|
try:
|
|
return orjson.loads(s)
|
|
except (TypeError, ValueError):
|
|
return engineio_json.loads(s, *args, **kwargs)
|
|
|
|
# Drop-in for stdlib ``json``: ``JSONCodec.dumps`` / ``JSONCodec.loads``.
|
|
JSONCodec = ORJSONCodec
|
|
# Codec handed to the socket.io/engineio managers, which default to their own.
|
|
SOCKETIO_JSON = ORJSONCodec
|
|
else:
|
|
JSONCodec = stdlib_json
|
|
SOCKETIO_JSON = engineio_json
|