Improve performance of streaming JSON output by using an event-based JSON parser

[doerwalter]

Checked

• I searched existing ideas and did not find a similar one
• I added a very descriptive title
• I've clearly described the feature request and motivation for it

Feature request

Using langchain_core.output_parsers.json.JsonOutputParser in streaming mode is really slow for larger JSON output. I'd like to propose an improved JSON parser that uses the event-based JSON parserijson.

Motivation

langchain_core.output_parsers.json.JsonOutputParser is slow in streaming mode since it reparses the JSON input multiple times during the run. The Python module ijson (available at https://pypi.org/project/ijson/) provides an event-based JSON parser that can be used to speed up streaming JSON output.

The following test script compares the output of the current langchain_core.output_parsers.json.JsonOutputParser with the new incrementaljson.IncrementalJsonOutputParser (attached below in the "proposal" section):

import incrementaljson

import datetime, asyncio

from langchain_core.output_parsers import json
from langchain_openai import ChatOpenAI

t0 = None

prompt = (
	"output a list of the 100 most populous countries, their populations, "
	"their land area in square kilometers, their highest elevation in meters in JSON format. "
	"Use a dict with an outer key of `countries` which contains a list of countries. "
	"Each country should have the keys `name`, `population`, `area` and `elevation`."
)

def test(parser):
	print(f"{'='*40} Start {parser!r}")
	model = ChatOpenAI(
		model_name="gpt-4o-mini",
		temperature=0.6,
		api_key="..."
	)

	chain = (model | parser)

	global t0
	t1 = datetime.datetime.now()
	t0 = t1
	t2 = None
	text = None
	count = 0
	async def run():
		nonlocal t2
		nonlocal count
		nonlocal text
		async for text in chain.astream(prompt):
			if t2 is None:
				t2 = datetime.datetime.now()
			count += 1
			print(f"#{count:,} - [{datetime.datetime.now() - t0}]: {text}", flush=True)
	asyncio.run(run())
	t3 = datetime.datetime.now()

	print()
	print(f"{'='*40} Finished {parser!r} with {count:,} events, time until first output {t2-t1}, total time {t3-t1}")


test(incrementaljson.IncrementalJsonOutputParser())
test(json.JsonOutputParser())

This test produces JSON that is around 6K in length. Running the script outputs the following timing information:

Finished IncrementalJsonOutputParser() with 422 events, time until first output 0:00:01.012896, total time 0:00:45.439018
Finished JsonOutputParser() with 855 events, time until first output 0:00:00.989100, total time 0:04:56.244498

i.e. the current JsonOutputParser takes nearly five minutes, IncrementalJsonOutputParser finishes in 45 seconds.

Proposal (If applicable)

The following module implements the proposed new incremental JSON parser:

"""
:mod:`incrementaljson` provides an alternative JSON parser.

The :class:`IncrementalJsonOutputParser` is a subclass of
:class:`langchain_core.output_parsers.json.JsonOutputParser`, so it should
behave identically but offers improved performance for streaming output.

Implementation notes
--------------------

The improved performance is achieved by using :mod:`ijson`, an incremental or
event-based parser (similar to a SAX parser in the XML world). Unlike
traditional parsers that construct an object directly, this parser reports
"events" from the JSON input to the calling application step-by-step.
This allows the application to build the object incrementally as it consumes
parsing events, giving it more control and avoiding the "all-or-nothing" parsing
approach used by :mod:`json`. By managing this process, the parser can operate
within a generator or similar structure, yielding control back to the
application as needed. This approach improves performance over the classic
"all-or-nothing" parse for three main reasons:

1.	In the "all-or-nothing" approach, the input must be a complete JSON string.
	If it’s incomplete, :func:`langchain_core.utils.json.parse_partial_json`
	tries to fix it by appending missing delimiters, ensuring a complete JSON
	string. Decoding such partial JSON input requires multiple calls to
	:func:`json.loads` (and a preliminary pass over the JSON input to detect
	missing delimiters). Note that this approach will only ever do only *one*
	:func:`json.loads` call when the JSON is complete, specifically when
	the final closing delimiter (usually a `}`) has been received.

2.	In the "all-or-nothing" approach, as new input arrives, it is concatenated
	with the previous input and re-parsed from the beginning (as described in
	step 1). This means that if the input is divided into 100 parts, the first
	part will have been parsed 199 times.

3.	Each parse call in step 2 creates a new object from scratch (twice),
	which is then discarded after output. This puts unnecessary strain on the
	garbage collector.

The :class:`IncrementalJsonOutputParser` avoids these issues as follows:

1.	The incremental parser only has to react to incomplete input *once*
	(when no further input is available).

2.	Each character in the input is parsed *once*.

3.	The parser builds a single object over the entire parsing process
	(this does slightly alter semantics: any modifications made to the object
	by the application during parsing will persist for the remainder of the
	process).

In "diff" mode, :class:`IncrementalJsonOutputParser` is even faster:
In the "all-or-nothing" approach, each output step requires a full comparison
to determine the differences between the old and the new JSON object.

In contrast, :class:`IncrementalJsonOutputParser` modifies the target object
directly, so it *knows* which modification has been applied, allowing it to
output this change directly.
"""

from typing import *

from ijson import common
import ijson
from langchain_core.output_parsers import json
from langchain_core.messages import base


def extract_text(input: Iterator[str | base.BaseMessage]) -> Iterator[str]:
	"""
	Iterates over ``input`` and extract the message text of all parts.
	"""
	for chunk in input:
		if isinstance(chunk, base.BaseMessage):
			chunk = chunk.content
		if chunk:
			yield chunk


async def aextract_text(input: AsyncIterator[str | base.BaseMessage]) -> AsyncIterator[str]:
	"""
	Iterates over ``input`` and extract the message text of all parts.

	This is the asynchronous version of :func:`extract_text`.
	"""
	async for chunk in input:
		if isinstance(chunk, base.BaseMessage):
			chunk = chunk.content
		if chunk:
			yield chunk


def extract_json_from_markdown(input: Iterator[str]) -> Iterator[str]:
	"""
	Iterates over ``input`` and extract a Markdown block with JSON.

	Everything before and after that block will be dropped.

	This assumes that the parts before and after the JSON block is small enough
	to be collected as a string, and to be searched multiple times.
	"""
	buffer = ""

	start_start_strings = ("```json\n", "```\n")
	middle_start_strings = ("\n```json\n", "\n```\n")
	end_string = "\n```"

	input = iter(input)

	# Find the start of the JSON part
	while True:
		# Search for a JSON block at the start of the output
		if buffer.startswith(start_start_strings):
			for s in start_start_strings:
				if buffer.startswith(s):
					buffer = buffer[len(s):]
					break
			break
		# Search for a JSON block at the middle of the output
		# (i.e. with a line feed before that)
		found = False
		for s in middle_start_strings:
			pos = buffer.find(s)
			if pos >= 0:
				buffer = buffer[pos+len(s):]
				found = True
				break
		if found:
			break
		try:
			buffer += next(input)
		except StopIteration:
			break

	# We're inside the JSON
	while True:
		pos = buffer.find(end_string)
		if pos >= 0:
			yield buffer[:pos]
			break
		# `buffer` might end with the first n-1 character of `end_string`
		# and the next character coming in might be the last character of
		# `end_string`, so we can't throw away those n-1 characters,
		# but we can output those before that.
		if len(buffer) > len(end_string):
			pos = len(buffer)-len(end_string)+1
			(result, buffer) = (buffer[:pos], buffer[pos:])
			yield result
		try:
			buffer += next(input)
		except StopIteration:
			break

	# We're outside of the JSON again, simply consume the rest of the iterator
	# without yielding anything
	for _ in input:
		pass


async def aextract_json_from_markdown(input: AsyncIterator[str]) -> Iterator[str]:
	"""
	Iterates over ``input`` and extract a Markdown block with JSON.

	Everything before and after that block will be dropped.

	This assumes that the parts before and after the block is small enough
	to be collected as a string, and to be searched multiple times.

	This is the asynchronous version of :func:`extract_json_from_markdown`.
	"""
	buffer = ""

	start_start_strings = ("```json\n", "```\n")
	middle_start_strings = ("\n```json\n", "\n```\n")
	end_string = "\n```"

	input = aiter(input)

	# Find the start of the JSON part
	while True:
		# Search for a JSON block at the start of the output
		if buffer.startswith(start_start_strings):
			for s in start_start_strings:
				if buffer.startswith(s):
					buffer = buffer[len(s):]
					break
			break
		# Search for a JSON block at the middle of the output
		# (i.e. with a line feed before that)
		found = False
		for s in middle_start_strings:
			pos = buffer.find(s)
			if pos >= 0:
				buffer = buffer[pos+len(s):]
				found = True
				break
		if found:
			break
		try:
			buffer += await anext(input)
		except StopAsyncIteration:
			break

	# We're inside the JSON
	while True:
		pos = buffer.find(end_string)
		if pos >= 0:
			yield buffer[:pos]
			break
		# `buffer` might end with the first n-1 character of `end_string`
		# and the next character coming in might be the last character of
		# `end_string` so we can't throw away those n-1 characters,
		# but we can output those before that.
		if len(buffer) > len(end_string):
			pos = len(buffer)-len(end_string)+1
			(result, buffer) = (buffer[:pos], buffer[pos:])
			yield result
		try:
			buffer += await anext(input)
		except StopAsyncIteration:
			break

	# We're outside of the JSON again, simply consume the rest of the iterator
	# without yielding anything
	async for _ in input:
		pass


class IteratorAsFile:
	r"""
	:class:`IteratorAsFile` wraps an iterator over strings or
	:class:`base.BaseMessage`\s into a file-like interface.

	The only method that :class:`IteratorAsFile provides is :meth:`read`.
	"""

	def __init__(self, iterator: Iterator[str | base.BaseMessage]):
		self.iterator = extract_json_from_markdown(extract_text(iterator))
		self.buffer = ""

	def read(self, size : int = -1) -> str:
		while size == -1 or len(self.buffer) < size:
			try:
				self.buffer += next(self.iterator)
			except StopIteration:
				break

		if size == -1:
			(result, self.buffer) = (self.buffer, "")
		else:
			(result, self.buffer) = (self.buffer[:size], self.buffer[size:])

		return result


class AsyncIteratorAsFile:
	r"""
	:class:`IteratorAsFile` wraps an iterator over strings or
	:class:`base.BaseMessage`\s into a file-like interface.

	The only method that :class:`IteratorAsFile provides is :meth:`read`.

	This is the asynchronous version of :class:`IteratorAsFile`.
	"""

	def __init__(self, iterator: AsyncIterator[str | base.BaseMessage]):
		self.iterator = aextract_json_from_markdown(aextract_text(iterator))
		self.buffer = ""

	async def read(self, size : int = -1) -> str:
		while size == -1 or len(self.buffer) < size:
			try:
				self.buffer += await anext(self.iterator)
			except StopAsyncIteration:
				break

		if size == -1:
			(result, self.buffer) = (self.buffer, "")
		else:
			(result, self.buffer) = (self.buffer[:size], self.buffer[size:])

		return result


class IncrementalJsonOutputParser(json.JsonOutputParser):
	"""
	:class:`IncrementalJsonOutputParser` is a subclass of :class:`json.JsonOutputParser`
	the should perform better for streaming output, since it incrementally parses
	the incoming JSON.
	"""

	def _transform(self, input: Iterator[Union[str, base.BaseMessage]]) -> Iterator[Any]:
		_stack = []
		_keys = []
		_result = None

		def _add(value):
			nonlocal _result

			if _stack:
				if isinstance(_keys[-1], int):
					_stack[-1].append(value)
					_keys[-1] += 1
				else:
					_stack[-1][_keys[-1]] = value
			else:
				_result = value

		def output(value):
			if self.diff:
				return [dict(op="add", path="".join(f"/{key}" for key in _keys), value=value)]
			else:
				return _result

		try:
			for (event, value) in ijson.basic_parse(extract_json_from_markdown(extract_text(input))):
				if event == "start_array":
					value = []
					_add(value)
					_stack.append(value)
					_keys.append(None)
					yield output(value)
				elif event == "end_array":
					_stack.pop()
					_keys.pop()
				elif event == "start_map":
					value = {}
					_add(value)
					_stack.append(value)
					_keys.append("")
					yield output(value)
				elif event == "end_map":
					_stack.pop()
					_keys.pop()
				elif event == "map_key":
					_keys[-1] = value
				else:
					_add(value)
					yield output(value)
			yield _result
		except common.IncompleteJSONError:
			pass

	async def _atransform(self, input: AsyncIterator[Union[str, base.BaseMessage]]) -> AsyncIterator[Any]:
		_stack = []
		_keys = []
		_result = None

		def _add(value):
			nonlocal _result

			if _stack:
				if isinstance(_keys[-1], int):
					_stack[-1].append(value)
					_keys[-1] += 1
				else:
					_stack[-1][_keys[-1]] = value
			else:
				_result = value

		def output(value):
			if self.diff:
				return [dict(op="add", path="".join(f"/{key}" for key in _keys), value=value)]
			else:
				return _result

		try:
			async for (event, value) in ijson.basic_parse_async(AsyncIteratorAsFile(input), buf_size=10):
				if event == "start_array":
					value = []
					_add(value)
					yield output(value)
					_stack.append(value)
					_keys.append(0)
				elif event == "end_array":
					_stack.pop()
					_keys.pop()
				elif event == "start_map":
					value = {}
					_add(value)
					yield output(value)
					_stack.append(value)
					_keys.append("")
				elif event == "end_map":
					_stack.pop()
					_keys.pop()
				elif event == "map_key":
					_keys[-1] = value
				else:
					_add(value)
					yield output(value)
		except common.IncompleteJSONError:
			pass