allow extract result to be non dict (#4069)

This commit is contained in:
LawyZheng 2025-11-22 10:36:43 +08:00 committed by GitHub
parent 8fb46ef1ca
commit 7c189818d9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 29 additions and 11 deletions

View file

@ -165,7 +165,9 @@ def _coerce_response_to_dict(response: Any) -> dict[str, Any]:
raise InvalidLLMResponseType(type(response).__name__)
def parse_api_response(response: litellm.ModelResponse, add_assistant_prefix: bool = False) -> dict[str, Any]:
def parse_api_response(
response: litellm.ModelResponse, add_assistant_prefix: bool = False, force_dict: bool = True
) -> dict[str, Any] | Any:
content = None
try:
content = response.choices[0].message.content
@ -174,6 +176,8 @@ def parse_api_response(response: litellm.ModelResponse, add_assistant_prefix: bo
content = "{" + content
parsed = json_repair.loads(content)
if not force_dict:
return parsed
return _coerce_response_to_dict(parsed)
except Exception:
@ -186,6 +190,8 @@ def parse_api_response(response: litellm.ModelResponse, add_assistant_prefix: bo
raise EmptyLLMResponseError(str(response))
content = _try_to_extract_json_from_markdown_format(content)
parsed = commentjson.loads(content)
if not force_dict:
return parsed
return _coerce_response_to_dict(parsed)
except Exception as e:
if content:
@ -196,6 +202,8 @@ def parse_api_response(response: litellm.ModelResponse, add_assistant_prefix: bo
)
try:
parsed = _fix_and_parse_json_string(content)
if not force_dict:
return parsed
return _coerce_response_to_dict(parsed)
except Exception as e2:
LOG.exception("Failed to auto-fix LLM response.", error=str(e2))