diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_agents.py b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_agents.py index 87d80ab32c..4ec71b6f19 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_agents.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_agents.py @@ -16,6 +16,7 @@ import json import logging import uuid +from collections.abc import Iterator from dataclasses import dataclass, field from typing import Any, cast @@ -36,6 +37,200 @@ logger = logging.getLogger(__name__) +_CODE_FENCE = "```" +_JSON_CODE_FENCE_QUALIFIER = "json" +_MAX_JSON_DECODE_BUDGET_MULTIPLIER = 4 +_NO_JSON = object() + + +def _iter_fenced_blocks(text: str, *, require_json_qualifier: bool) -> Iterator[str]: + """Yield non-overlapping fenced blocks in source order.""" + search_start = 0 + while True: + opening_index = text.find(_CODE_FENCE, search_start) + if opening_index < 0: + return + + content_start = opening_index + len(_CODE_FENCE) + if require_json_qualifier: + if not text.startswith(_JSON_CODE_FENCE_QUALIFIER, content_start): + search_start = content_start + continue + + qualifier_end = content_start + len(_JSON_CODE_FENCE_QUALIFIER) + if ( + qualifier_end < len(text) + and not text[qualifier_end].isspace() + and text[qualifier_end] not in "{[" + and not text.startswith(_CODE_FENCE, qualifier_end) + ): + search_start = content_start + continue + content_start = qualifier_end + + while content_start < len(text) and text[content_start].isspace(): + content_start += 1 + + closing_index = text.find(_CODE_FENCE, content_start) + if closing_index < 0: + return + + yield text[content_start:closing_index].strip() + search_start = closing_index + len(_CODE_FENCE) + + +def _index_escaped_quotes(text: str) -> bytearray: + """Index quote characters preceded by an odd-length backslash run.""" + escaped_quotes = bytearray(len(text)) + backslash_count = 0 + + for index, char in enumerate(text): + if char == "\\": + backslash_count += 1 + continue + + if char == '"' and backslash_count % 2 == 1: + escaped_quotes[index] = 1 + backslash_count = 0 + + return escaped_quotes + + +def _index_json_candidates_forward(text: str, escaped_quotes: bytearray) -> set[tuple[int, int]]: + """Index JSON candidate ranges from left to right.""" + candidates: set[tuple[int, int]] = set() + object_openings: list[int] = [] + array_openings: list[int] = [] + in_string = False + + for index, char in enumerate(text): + if not object_openings and not array_openings: + if char in "{[": + (object_openings if char == "{" else array_openings).append(index) + continue + + if char == '"' and not escaped_quotes[index]: + in_string = not in_string + continue + + if in_string: + continue + + if char in "{[": + (object_openings if char == "{" else array_openings).append(index) + elif char == "}" and object_openings: + candidates.add((object_openings.pop(), index)) + elif char == "]" and array_openings: + candidates.add((array_openings.pop(), index)) + + return candidates + + +def _index_json_candidates_reverse(text: str, escaped_quotes: bytearray) -> set[tuple[int, int]]: + """Index JSON candidate ranges from right to left.""" + candidates: set[tuple[int, int]] = set() + object_closings: list[int] = [] + array_closings: list[int] = [] + in_string = False + + for index in range(len(text) - 1, -1, -1): + char = text[index] + if not object_closings and not array_closings: + if char in "}]": + (object_closings if char == "}" else array_closings).append(index) + continue + + if char == '"' and not escaped_quotes[index]: + in_string = not in_string + continue + + if in_string: + continue + + if char in "}]": + (object_closings if char == "}" else array_closings).append(index) + elif char == "{" and object_closings: + candidates.add((index, object_closings.pop())) + elif char == "[" and array_closings: + candidates.add((index, array_closings.pop())) + + return candidates + + +def _find_last_decodable_json(text: str) -> Any: + """Find the last decodable JSON object or array within text.""" + escaped_quotes = _index_escaped_quotes(text) + candidates = _index_json_candidates_forward(text, escaped_quotes) + candidates.update(_index_json_candidates_reverse(text, escaped_quotes)) + + candidate_groups: list[tuple[int, int, list[tuple[int, int]]]] = [] + for candidate in sorted(candidates): + json_start, json_end = candidate + if not candidate_groups or json_start > candidate_groups[-1][1]: + candidate_groups.append((json_start, json_end, [candidate])) + continue + + group_start, group_end, group_candidates = candidate_groups[-1] + group_candidates.append(candidate) + candidate_groups[-1] = (group_start, max(group_end, json_end), group_candidates) + + for group_start, group_end, group_candidates in reversed(candidate_groups): + group_span = group_end - group_start + 1 + primary_decode_budget = group_span * (_MAX_JSON_DECODE_BUDGET_MULTIPLIER // 2) + recovery_decode_budget = group_span * ( + _MAX_JSON_DECODE_BUDGET_MULTIPLIER - (_MAX_JSON_DECODE_BUDGET_MULTIPLIER // 2) + ) + attempted_candidates: set[tuple[int, int]] = set() + last_json: Any = _NO_JSON + consumed_end = -1 + candidate_index = 0 + + while candidate_index < len(group_candidates) and primary_decode_budget > 0: + json_start, json_end = group_candidates[candidate_index] + candidate_index += 1 + if json_start <= consumed_end: + continue + + candidate_length = json_end - json_start + 1 + if candidate_length > primary_decode_budget: + continue + + primary_decode_budget -= candidate_length + attempted_candidates.add((json_start, json_end)) + try: + last_json = json.loads(text[json_start : json_end + 1]) + except json.JSONDecodeError: + continue + + consumed_end = json_end + while candidate_index < len(group_candidates) and group_candidates[candidate_index][0] <= consumed_end: + candidate_index += 1 + + recovery_candidates = sorted( + group_candidates, + key=lambda candidate: (candidate[1] - candidate[0], -candidate[0]), + ) + for json_start, json_end in recovery_candidates: + if recovery_decode_budget == 0: + break + if (json_start, json_end) in attempted_candidates or json_start <= consumed_end: + continue + + candidate_length = json_end - json_start + 1 + if candidate_length > recovery_decode_budget: + continue + + recovery_decode_budget -= candidate_length + try: + return json.loads(text[json_start : json_end + 1]) + except json.JSONDecodeError: + continue + + if last_json is not _NO_JSON: + return last_json + + raise json.JSONDecodeError("No valid JSON found in response", text, 0) + def _extract_json_from_response(text: str) -> Any: r"""Extract and parse JSON from an agent response. @@ -58,13 +253,11 @@ def _extract_json_from_response(text: str) -> Any: text: The raw text response from an agent Returns: - Parsed JSON as a Python dict/list, or None if parsing fails + Parsed JSON, or None if the response is empty. Raises: json.JSONDecodeError: If no valid JSON can be extracted """ - import re - if not text: return None @@ -79,96 +272,18 @@ def _extract_json_from_response(text: str) -> Any: except json.JSONDecodeError: pass - # Try extracting from markdown code blocks: ```json ... ``` or ``` ... ``` - # Use the last code block if there are multiple - code_block_patterns = [ - r"```json\s*\n?(.*?)\n?```", # ```json ... ``` - r"```\s*\n?(.*?)\n?```", # ``` ... ``` - ] - for pattern in code_block_patterns: - matches = list(re.finditer(pattern, text, re.DOTALL)) - if matches: - # Try the last match first (most likely to be the final result) - for match in reversed(matches): - try: - return json.loads(match.group(1).strip()) - except json.JSONDecodeError: - continue - - # Find ALL JSON objects {...} or arrays [...] in the text and return the last valid one - # This handles cases where agents stream multiple JSON objects (partial, then final) - all_json_objects: list[Any] = [] - - pos = 0 - while pos < len(text): - # Find next { or [ - json_start = -1 - bracket_char = None - for i in range(pos, len(text)): - if text[i] == "{": - json_start = i - bracket_char = "{" - break - if text[i] == "[": - json_start = i - bracket_char = "[" - break - - if json_start < 0: - break # No more JSON objects - - # Find matching closing bracket - open_bracket = bracket_char - close_bracket = "}" if open_bracket == "{" else "]" - depth = 0 - in_string = False - escape_next = False - found_end = False - - for i in range(json_start, len(text)): - char = text[i] - - if escape_next: - escape_next = False - continue - - if char == "\\": - escape_next = True - continue - - if char == '"' and not escape_next: - in_string = not in_string - continue - - if in_string: + # Exactly-qualified JSON fences take precedence over plain fences. + for require_json_qualifier in (True, False): + last_fenced_json: Any = _NO_JSON + for block in _iter_fenced_blocks(text, require_json_qualifier=require_json_qualifier): + try: + last_fenced_json = json.loads(block) + except json.JSONDecodeError: continue + if last_fenced_json is not _NO_JSON: + return last_fenced_json - if char == open_bracket: - depth += 1 - elif char == close_bracket: - depth -= 1 - if depth == 0: - # Found the end - potential_json = text[json_start : i + 1] - try: - parsed = json.loads(potential_json) - all_json_objects.append(parsed) - except json.JSONDecodeError: - pass - pos = i + 1 - found_end = True - break - - if not found_end: - # Malformed JSON, move past the start character - pos = json_start + 1 - - # Return the last valid JSON object (most likely to be the final/complete result) - if all_json_objects: - return all_json_objects[-1] - - # Unable to extract JSON - raise json.JSONDecodeError("No valid JSON found in response", text, 0) + return _find_last_decodable_json(text) def _validate_conversation_history(messages: list[Message], agent_name: str) -> None: diff --git a/python/packages/declarative/tests/test_graph_executors.py b/python/packages/declarative/tests/test_graph_executors.py index 5dce56d4e7..96315bec06 100644 --- a/python/packages/declarative/tests/test_graph_executors.py +++ b/python/packages/declarative/tests/test_graph_executors.py @@ -1322,6 +1322,340 @@ def test_multiple_json_objects_with_text_between(self): result = _extract_json_from_response(text) assert result == {"status": "complete", "id": 42} + def test_multiple_qualified_code_blocks_returns_last_valid(self): + """Test that the last valid qualified code block is returned.""" + from agent_framework_declarative._workflows._executors_agents import ( + _extract_json_from_response, + ) + + text = """```json +{"status": "pending"} +``` +```json +{"status": "complete"} +```""" + result = _extract_json_from_response(text) + assert result == {"status": "complete"} + + def test_invalid_later_qualified_code_block_uses_previous_valid(self): + """Test that an invalid later block does not replace an earlier valid block.""" + from agent_framework_declarative._workflows._executors_agents import ( + _extract_json_from_response, + ) + + text = """```json +{"status": "complete"} +``` +```json +not valid JSON +```""" + result = _extract_json_from_response(text) + assert result == {"status": "complete"} + + def test_qualified_code_block_takes_precedence_over_plain_block(self): + """Test that a qualified block is preferred over a later plain block.""" + from agent_framework_declarative._workflows._executors_agents import ( + _extract_json_from_response, + ) + + text = """```json +{"source": "qualified"} +``` +``` +{"source": "plain"} +```""" + result = _extract_json_from_response(text) + assert result == {"source": "qualified"} + + def test_invalid_qualified_code_block_falls_through_to_plain_block(self): + """Test that plain blocks are considered when qualified blocks are invalid.""" + from agent_framework_declarative._workflows._executors_agents import ( + _extract_json_from_response, + ) + + text = """```json +not valid JSON +``` +``` +{"source": "plain"} +```""" + result = _extract_json_from_response(text) + assert result == {"source": "plain"} + + @pytest.mark.parametrize(("json_text", "expected"), [("null", None), ("false", False), ("0", 0)]) + def test_json_scalar_in_qualified_code_block(self, json_text, expected): + """Test that valid JSON scalars are not confused with a missing result.""" + from agent_framework_declarative._workflows._executors_agents import ( + _extract_json_from_response, + ) + + result = _extract_json_from_response(f"```json\n{json_text}\n```") + assert result == expected + + def test_unrecognized_code_block_qualifier_is_not_removed(self): + """Test that the plain-block pass does not consume language qualifiers.""" + import json + + from agent_framework_declarative._workflows._executors_agents import ( + _extract_json_from_response, + ) + + with pytest.raises(json.JSONDecodeError): + _extract_json_from_response("```yaml\nfalse\n```") + + def test_inline_json_code_block(self): + """Test extracting JSON from an inline qualified code block.""" + from agent_framework_declarative._workflows._executors_agents import ( + _extract_json_from_response, + ) + + result = _extract_json_from_response('Result: ```json{"status": "complete"}```.') + assert result == {"status": "complete"} + + def test_inline_json_array_code_block(self): + """Test extracting a JSON array from an inline qualified code block.""" + from agent_framework_declarative._workflows._executors_agents import ( + _extract_json_from_response, + ) + + result = _extract_json_from_response("Result: ```json[1, 2]```.") + assert result == [1, 2] + + def test_json5_scalar_is_not_treated_as_json_qualified(self): + """Test that a JSON5 qualifier prefix is not interpreted as JSON.""" + import json + + from agent_framework_declarative._workflows._executors_agents import ( + _extract_json_from_response, + ) + + with pytest.raises(json.JSONDecodeError): + _extract_json_from_response("```json5```") + + @pytest.mark.parametrize("qualifier", ["json5", "jsonc"]) + def test_nonstandard_json_qualified_object_uses_general_fallback(self, qualifier): + """Test that objects in nonstandard JSON blocks are recovered by fallback.""" + from agent_framework_declarative._workflows._executors_agents import ( + _extract_json_from_response, + ) + + result = _extract_json_from_response(f'```{qualifier}\n{{"status": "complete"}}\n```') + assert result == {"status": "complete"} + + def test_nonstandard_json_block_does_not_take_qualified_precedence(self): + """Test that JSON5 blocks do not take precedence over plain blocks.""" + from agent_framework_declarative._workflows._executors_agents import ( + _extract_json_from_response, + ) + + text = """```json5 +{"source": "json5"} +``` +``` +{"source": "plain"} +```""" + result = _extract_json_from_response(text) + assert result == {"source": "plain"} + + def test_json_code_block_with_crlf(self): + """Test extracting JSON from a code block with CRLF line endings.""" + from agent_framework_declarative._workflows._executors_agents import ( + _extract_json_from_response, + ) + + result = _extract_json_from_response('Result:\r\n```json\r\n{"status": "complete"}\r\n```\r\n') + assert result == {"status": "complete"} + + def test_array_with_brackets_and_escapes_in_string(self): + """Test nested delimiters and escapes inside JSON strings.""" + from agent_framework_declarative._workflows._executors_agents import ( + _extract_json_from_response, + ) + + text = r'Info: [{"message": "Use [x] and {y}", "path": "C:\\temp"}]' + result = _extract_json_from_response(text) + assert result == [{"message": "Use [x] and {y}", "path": r"C:\temp"}] + + def test_unterminated_code_block_raises_error(self): + """Test that an unterminated whitespace-heavy code block fails safely.""" + import json + + from agent_framework_declarative._workflows._executors_agents import ( + _extract_json_from_response, + ) + + text = f"```json\n{' ' * 64}X" + with pytest.raises(json.JSONDecodeError): + _extract_json_from_response(text) + + @pytest.mark.parametrize("text", ["{" * 64 + "X", "[" * 64 + "X"]) + def test_repeated_unmatched_brackets_raise_error(self, text): + """Test that repeated unmatched opening brackets fail safely.""" + import json + + from agent_framework_declarative._workflows._executors_agents import ( + _extract_json_from_response, + ) + + with pytest.raises(json.JSONDecodeError): + _extract_json_from_response(text) + + def test_valid_json_after_unmatched_outer_bracket(self): + """Test recovering valid JSON nested after an unmatched outer bracket.""" + from agent_framework_declarative._workflows._executors_agents import ( + _extract_json_from_response, + ) + + result = _extract_json_from_response('{{"status": "complete"}') + assert result == {"status": "complete"} + + def test_valid_json_after_mismatched_bracket_candidate(self): + """Test recovering valid JSON after a malformed mixed-bracket candidate.""" + from agent_framework_declarative._workflows._executors_agents import ( + _extract_json_from_response, + ) + + text = '{"broken": [} then {"status": "complete"} ]}' + result = _extract_json_from_response(text) + assert result == {"status": "complete"} + + def test_valid_outer_json_is_preferred_over_crossing_reverse_candidate(self): + """Test a reverse-indexed crossing candidate does not override valid JSON.""" + from agent_framework_declarative._workflows._executors_agents import ( + _extract_json_from_response, + ) + + result = _extract_json_from_response('{"s": "["}"]') + assert result == {"s": "["} + + def test_valid_outer_json_is_preferred_over_nested_candidate(self): + """Test a malformed wrapper does not cause a nested value to win.""" + from agent_framework_declarative._workflows._executors_agents import ( + _extract_json_from_response, + ) + + result = _extract_json_from_response('{"mixed": [} {"final": {"nested": 1}} ]}') + assert result == {"final": {"nested": 1}} + + def test_valid_json_after_double_escaped_fragment(self): + """Test recovering valid JSON after a double-escaped malformed fragment.""" + from agent_framework_declarative._workflows._executors_agents import ( + _extract_json_from_response, + ) + + text = r'{\"partial\": true} then {"status": "complete"}' + result = _extract_json_from_response(text) + assert result == {"status": "complete"} + + def test_valid_json_after_double_escaped_fenced_fragment(self): + """Test recovering valid JSON after an invalid fenced fragment.""" + from agent_framework_declarative._workflows._executors_agents import ( + _extract_json_from_response, + ) + + text = """```json +{\\"partial\\": true} +``` +{"status": "complete"}""" + result = _extract_json_from_response(text) + assert result == {"status": "complete"} + + def test_valid_json_after_brace_in_quoted_explanation(self): + """Test recovering valid JSON after a brace in quoted explanatory text.""" + from agent_framework_declarative._workflows._executors_agents import ( + _extract_json_from_response, + ) + + text = 'The model said "use { as the opener" before {"status": "complete"}' + result = _extract_json_from_response(text) + assert result == {"status": "complete"} + + def test_valid_json_after_unterminated_quoted_candidate(self): + """Test recovering valid JSON after an unterminated quoted candidate.""" + from agent_framework_declarative._workflows._executors_agents import ( + _extract_json_from_response, + ) + + text = '{"partial: true} then {"status": "complete"}' + result = _extract_json_from_response(text) + assert result == {"status": "complete"} + + def test_latest_json_after_valid_and_unterminated_candidates(self): + """Test returning the latest JSON after an earlier valid and poisoned candidate.""" + from agent_framework_declarative._workflows._executors_agents import ( + _extract_json_from_response, + ) + + text = '{"status": "pending"} {"partial: true} then {"status": "complete"}' + result = _extract_json_from_response(text) + assert result == {"status": "complete"} + + def test_valid_json_after_two_poisoned_candidates(self): + """Test recovering valid JSON after two malformed candidates.""" + from agent_framework_declarative._workflows._executors_agents import ( + _extract_json_from_response, + ) + + text = '{"mixed": [} {"partial: true} then {"status": "complete"}' + result = _extract_json_from_response(text) + assert result == {"status": "complete"} + + def test_valid_json_after_many_malformed_candidates(self): + """Test recovery is not limited by the number of malformed candidates.""" + from agent_framework_declarative._workflows._executors_agents import ( + _extract_json_from_response, + ) + + text = "[{}[[" * 50 + '{"final": "value"}' + result = _extract_json_from_response(text) + assert result == {"final": "value"} + + def test_review_reported_malformed_candidate_sequence(self): + """Test the review-reported malformed prefix before final JSON.""" + from agent_framework_declarative._workflows._executors_agents import ( + _extract_json_from_response, + ) + + result = _extract_json_from_response('[{}[["{"final": "value"}') + assert result == {"final": "value"} + + def test_valid_json_before_nested_malformed_suffix(self): + """Test a malformed suffix cannot consume the earlier candidate's decode budget.""" + from agent_framework_declarative._workflows._executors_agents import ( + _extract_json_from_response, + ) + + result = _extract_json_from_response('{"good": 1} [[[[[[[[[[x]]]]]]]]]]') + assert result == {"good": 1} + + def test_last_sibling_json_inside_malformed_wrapper(self): + """Test that the last valid sibling wins inside a malformed wrapper.""" + from agent_framework_declarative._workflows._executors_agents import ( + _extract_json_from_response, + ) + + result = _extract_json_from_response('[x {"first": 1} {"last": 2}]') + assert result == {"last": 2} + + def test_valid_json_inside_deeply_nested_malformed_wrapper(self): + """Test recovery budget is reserved for a deeply nested valid value.""" + from agent_framework_declarative._workflows._executors_agents import ( + _extract_json_from_response, + ) + + result = _extract_json_from_response('[[[[x {"final": "value"}]]]]') + assert result == {"final": "value"} + + def test_valid_json_between_malformed_prefix_and_suffix(self): + """Test recovery prioritizes compact JSON over malformed wrappers.""" + from agent_framework_declarative._workflows._executors_agents import ( + _extract_json_from_response, + ) + + text = ("[" * 4) + 'x {"good": 1} ' + ("[" * 9) + "x" + ("]" * 13) + result = _extract_json_from_response(text) + assert result == {"good": 1} + class TestPowerFxConditionalImport: """The _declarative_base module should be importable without dotnet/powerfx."""