A more performant streaming json parser

[CNSeniorious000]

Description

I personally have been an LLM developer for a few month, structural output tools like instructor helped me a lot. These days I found the core implementation for streaming json parsing in instructor in instructor/dsl/partialjson.py, and think there is still some space to improve it:

• \ is not handled correctly:

  >>> from instructor.dsl.partialjson import JSONParser
  >>> parse = JSONParser().parse
  >>> parse('"\\t"')  # this is right because it in fact calls json.loads
  '\t'
  >>> parse('"\\t')  # this shows that `\\` not escaped correctly
  '\\t'

• type may change during iteration:

  >>> parse('{"a')
  {'a': None}
  >>> parse('{"ab')
  {'ab': None}
  >>> parse('{"ab": [')
  {'ab': []}

  In the example above, a is not a field in the pydantic model, but still it will be iterated from the iterator (let's think a is placeholder for a token).
  And then, the type of ab changed from None to list.

• the current implementation doesn't support non-standard json types like NaN and Infinity

• it is hard-linked with json.loads, but it maybe better to make use of pydantic's faster json parser, or use a 3-party json parser like ujson/orjson

Solution

I made a python package (as well as an almost the same js package) to parse partial json. It is designed to be flexible and performant. It also deal with NaN well (demo).

Flexible means, its implementation based on completing the suffix JSON tokens. So basically you just get a complete JSON and then send it to any json parse function you like.

Performance:

(both use json.dumps under the hood)

>>> from faker import Faker
>>> example_json = Faker().json(num_rows=200)

>>> from instructor.dsl.partialjson import JSONParser
>>> parse = JSONParser().parse
>>> @timeit parse(example_json[:-1])
2.05 ms ± 60.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

>>> from partial_json_parser import loads as parse
>>> %timeit parse(example_json[:-1])
1.46 ms ± 49.3 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

partial-json-parser is well-tested by the hypothesis test framework, which generates strange JSONs, then we make sure that from the first char to the total string is parsable. This package runs on py3.6 - py3.12, and runs even faster on pypy.

It is my personal project since half a year ago. I hope my work can benefit the LLM ecosystem.

Suggestion

If you agree with the Can I draft a PR, which will power instructor.Partial by partial-json-parser, with pydantic's json parser fix the \\ and type issues mentioned above?

[jxnl]

we're actually waiting on pydantic's very own jiter to support a faster partial

[shanktt]

Here's the PR they're working on: pydantic/jiter#66

[CNSeniorious000]

Thanks. Good to hear that the partial stuffs is going to be solved from the rust side