perf: use the orjson codec to parse Jupyter kernel messages (#27812)

The code interpreter parses every message from the Jupyter kernel websocket with stdlib `json`, in a loop that runs for the duration of an execution. Messages carrying large stdout or a base64 image payload are the expensive ones.

It now goes through `JSONCodec`, which selects orjson when `ENABLE_ORJSON` is set. Every consumer of the parsed message reads strings only: `content.text`, `content.data['text/plain']` and `['image/png']`, `content.traceback`, and `content.execution_state`. Jupyter renders large integers into `text/plain` as strings rather than JSON numbers, so no numeric round trip is involved.

The one-shot `execute_request` message this module sends keeps stdlib `json`; it is a small fixed-shape dict sent once per execution.

With `ENABLE_ORJSON` unset, which is the default, `JSONCodec` is stdlib `json` and this call site behaves exactly as before.
This commit is contained in:
Classic298
2026-07-31 23:32:05 +02:00
committed by GitHub
parent d03e9af0b7
commit e2221fb662

View File

@@ -7,6 +7,7 @@ from typing import Optional
import aiohttp
import websockets
from open_webui.env import AIOHTTP_CLIENT_ALLOW_REDIRECTS
from open_webui.utils.json_codec import JSONCodec
from pydantic import BaseModel
logger = logging.getLogger(__name__)
@@ -157,7 +158,7 @@ class JupyterCodeExecuter:
try:
# wait for message
message = await asyncio.wait_for(ws.recv(), self.timeout)
message_data = json.loads(message)
message_data = JSONCodec.loads(message)
# msg id not match, skip
if message_data.get('parent_header', {}).get('msg_id') != msg_id:
continue