diff --git a/instructor/v2/core/json.py b/instructor/v2/core/json.py index e9238119c..ad005896c 100644 --- a/instructor/v2/core/json.py +++ b/instructor/v2/core/json.py @@ -142,9 +142,21 @@ def extract_json_from_stream(chunks: Iterable[str]) -> Generator[str, None, None delimiter_stack.pop() if not delimiter_stack: buffer.append(char) - yield from buffer + candidate = "".join(buffer) buffer = [] json_started = False + try: + json.loads(candidate) + except ValueError: + # The bracket-balanced span we just closed + # isn't actually valid JSON - e.g. a + # brace-delimited aside in the model's prose + # that happened to appear before the real + # payload. Discard it and keep scanning + # instead of emitting it as if it were the + # extracted JSON. + continue + yield from candidate continue buffer.append(char) @@ -231,10 +243,19 @@ async def extract_json_from_stream_async( delimiter_stack.pop() if not delimiter_stack: buffer.append(char) - for buffered_char in buffer: - yield buffered_char + candidate = "".join(buffer) buffer = [] json_started = False + try: + json.loads(candidate) + except ValueError: + # See the matching comment in + # extract_json_from_stream - discard + # bracket-balanced spans that aren't + # actually valid JSON and keep scanning. + continue + for buffered_char in candidate: + yield buffered_char continue buffer.append(char) diff --git a/tests/v2/test_json_helpers.py b/tests/v2/test_json_helpers.py index 8683d8254..20b29ee58 100644 --- a/tests/v2/test_json_helpers.py +++ b/tests/v2/test_json_helpers.py @@ -121,6 +121,24 @@ def test_extract_json_from_stream_preserves_backticks_in_fenced_string() -> None assert "".join(extract_json_from_stream(chunks)) == '{"code":"`inline`"}' +def test_extract_json_from_stream_discards_non_json_brace_span_before_payload() -> None: + # MD_JSON prompts often produce a prose preamble before the actual payload, + # and that preamble can itself contain a balanced pair of braces that is not + # JSON (e.g. a parenthetical aside). A naive bracket-matching scan treats that + # span as "the JSON" once it balances and emits it verbatim, so the real + # payload that follows gets appended onto invalid JSON instead of replacing + # it. The extractor should recognize the span isn't valid JSON, drop it, and + # keep scanning for the real payload. + chunks = [ + "I'll pull out the fields now. ", + "{Note: keeping the original casing} ", + "Here is the result: ", + '{"name":"Ada","age":30}', + ] + + assert "".join(extract_json_from_stream(chunks)) == '{"name":"Ada","age":30}' + + @pytest.mark.asyncio async def test_extract_json_from_stream_async_preserves_backticks_in_string() -> None: async def chunks(): @@ -176,3 +194,20 @@ async def chunks(): "".join([chunk async for chunk in extract_json_from_stream_async(chunks())]) == '{"path":"C:\\\\","name":"Ada"}' ) + + +@pytest.mark.asyncio +async def test_extract_json_from_stream_async_discards_non_json_brace_span_before_payload() -> None: + async def chunks(): + for chunk in [ + "I'll pull out the fields now. ", + "{Note: keeping the original casing} ", + "Here is the result: ", + '{"name":"Ada","age":30}', + ]: + yield chunk + + assert ( + "".join([chunk async for chunk in extract_json_from_stream_async(chunks())]) + == '{"name":"Ada","age":30}' + )