refactor: improve type handling in structured output parsing

- Changed the type of `structured_output` from `dict` to `Mapping` for better type flexibility.
- Introduced a new variable `parsed` to hold the result of JSON validation, enhancing clarity.
- Updated error handling to use a temporary variable `temp_parsed` for JSON repair, ensuring correct type casting before assignment.
pull/21565/head
Yeuoly 11 months ago
parent 47f7a32aeb
commit 79290687bf

@ -1,7 +1,7 @@
import json import json
from collections.abc import Generator, Mapping, Sequence from collections.abc import Generator, Mapping, Sequence
from copy import deepcopy from copy import deepcopy
from typing import Any, Literal, Optional, overload from typing import Any, Literal, Optional, cast, overload
import json_repair import json_repair
from pydantic import TypeAdapter, ValidationError from pydantic import TypeAdapter, ValidationError
@ -233,7 +233,8 @@ def _handle_prompt_based_schema(
def _parse_structured_output(result_text: str) -> Mapping[str, Any]: def _parse_structured_output(result_text: str) -> Mapping[str, Any]:
structured_output: dict[str, Any] = {} structured_output: Mapping[str, Any] = {}
parsed: Mapping[str, Any] = {}
try: try:
parsed = TypeAdapter(Mapping).validate_json(result_text) parsed = TypeAdapter(Mapping).validate_json(result_text)
if not isinstance(parsed, dict): if not isinstance(parsed, dict):
@ -241,14 +242,14 @@ def _parse_structured_output(result_text: str) -> Mapping[str, Any]:
structured_output = parsed structured_output = parsed
except ValidationError as e: except ValidationError as e:
# if the result_text is not a valid json, try to repair it # if the result_text is not a valid json, try to repair it
parsed = json_repair.loads(result_text) temp_parsed = json_repair.loads(result_text)
if not isinstance(parsed, dict): if not isinstance(temp_parsed, dict):
# handle reasoning model like deepseek-r1 got '<think>\n\n</think>\n' prefix # handle reasoning model like deepseek-r1 got '<think>\n\n</think>\n' prefix
if isinstance(parsed, list): if isinstance(temp_parsed, list):
parsed = next((item for item in parsed if isinstance(item, dict)), {}) parsed = next((item for item in temp_parsed if isinstance(item, dict)), {})
else: else:
raise OutputParserError(f"Failed to parse structured output: {result_text}") raise OutputParserError(f"Failed to parse structured output: {result_text}")
structured_output = parsed structured_output = cast(dict, parsed)
return structured_output return structured_output

Loading…
Cancel
Save