640 lines
22 KiB
Python
640 lines
22 KiB
Python
|
from copy import deepcopy
|
||
|
|
import json
|
||
|
|
from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
|
||
|
|
|
||
|
|
from guardrails_ai.types import ReAsk
|
||
|
|
from guardrails.classes.execution.guard_execution_options import GuardExecutionOptions
|
||
|
|
from guardrails.classes.output_type import OutputTypes
|
||
|
|
from guardrails.prompt.instructions import Instructions
|
||
|
|
from guardrails.prompt.prompt import Prompt
|
||
|
|
from guardrails.prompt.messages import Messages
|
||
|
|
from guardrails.schema.generator import generate_example
|
||
|
|
from guardrails.schema.rail_schema import json_schema_to_rail_output
|
||
|
|
from guardrails.types.validator import ValidatorMap
|
||
|
|
from guardrails.utils.constants import constants
|
||
|
|
from guardrails.utils.prompt_utils import prompt_content_for_schema, prompt_uses_xml
|
||
|
|
|
||
|
|
|
||
|
|
### Classes/Types ###
|
||
|
|
class FieldReAsk(ReAsk):
|
||
|
|
"""An implementation of ReAsk that is used to reask for a specific field.
|
||
|
|
Inherits from ReAsk.
|
||
|
|
|
||
|
|
Attributes:
|
||
|
|
path (Optional[List[Any]]): a list of keys that
|
||
|
|
designated the path to the field that failed validation.
|
||
|
|
"""
|
||
|
|
|
||
|
|
# FIXME: This shouldn't be optional
|
||
|
|
# We should be able to assign it on init now
|
||
|
|
path: Optional[List[Any]] = None
|
||
|
|
|
||
|
|
|
||
|
|
class SkeletonReAsk(ReAsk):
|
||
|
|
"""An implementation of ReAsk that is used to reask for structured data
|
||
|
|
when the response does not match the expected schema.
|
||
|
|
|
||
|
|
Inherits from ReAsk.
|
||
|
|
"""
|
||
|
|
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
class NonParseableReAsk(ReAsk):
|
||
|
|
"""An implementation of ReAsk that is used to reask for structured data
|
||
|
|
when the response is not parseable as JSON.
|
||
|
|
|
||
|
|
Inherits from ReAsk.
|
||
|
|
"""
|
||
|
|
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
def to_reask(obj: Any) -> FieldReAsk | SkeletonReAsk | NonParseableReAsk | ReAsk:
|
||
|
|
reask = ReAsk.model_validate(obj)
|
||
|
|
if hasattr(reask, "path"):
|
||
|
|
return FieldReAsk.model_validate(reask.model_dump())
|
||
|
|
|
||
|
|
if reask.fail_results and len(reask.fail_results) == 1:
|
||
|
|
error_message = reask.fail_results[0].error_message
|
||
|
|
if error_message == "Output is not parseable as JSON":
|
||
|
|
return NonParseableReAsk.model_validate(reask.model_dump())
|
||
|
|
elif "JSON does not match schema" in error_message:
|
||
|
|
return SkeletonReAsk.model_validate(reask.model_dump())
|
||
|
|
|
||
|
|
return reask
|
||
|
|
|
||
|
|
|
||
|
|
### Internal Helper Methods ###
|
||
|
|
def get_reask_subschema(
|
||
|
|
json_schema: Dict[str, Any],
|
||
|
|
reasks: Optional[List[FieldReAsk]] = None,
|
||
|
|
) -> Dict[str, Any]:
|
||
|
|
"""Prune schema of any subschemas that are not in `reasks`.
|
||
|
|
|
||
|
|
Return the schema with only the subschemas that are being `reask`ed for and
|
||
|
|
their parents. If `reasks` is None, return the entire schema. If an
|
||
|
|
subschema is removed, remove all ancestors that have no children.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
root: A JSON Schema
|
||
|
|
reasks: The fields that are to be reasked.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
A JSON Schema.
|
||
|
|
"""
|
||
|
|
root = deepcopy(json_schema)
|
||
|
|
|
||
|
|
if reasks is None:
|
||
|
|
return root
|
||
|
|
|
||
|
|
# Find all elements that are to be retained
|
||
|
|
# NOTE: At this point, in the case of discriminated unions,
|
||
|
|
# the LLM has already decided which subschema of the union to use.
|
||
|
|
# This means that we can flatten complex schema compositions, e.g. anyOf's,
|
||
|
|
# and just build a subschema that represents the resolved schema
|
||
|
|
# of the LLM response.
|
||
|
|
# schema_paths_to_retain = []
|
||
|
|
# for reask in reasks:
|
||
|
|
# path = reask.path
|
||
|
|
# if path is None:
|
||
|
|
# raise RuntimeError("FieldReAsk path is None")
|
||
|
|
# schema_path = "$"
|
||
|
|
# for part in path:
|
||
|
|
# if isinstance(part, int):
|
||
|
|
# schema_path += ".items"
|
||
|
|
# else:
|
||
|
|
# schema_path += f".properties.{path}"
|
||
|
|
# schema_paths_to_retain.append(schema_path)
|
||
|
|
|
||
|
|
# # Remove all elements that are not to be retained
|
||
|
|
# def _prune_schema(schema: Dict[str, Any]) -> None:
|
||
|
|
# if schema.get("type") == SimpleTypes.ARRAY:
|
||
|
|
# if schema.children.item not in retain:
|
||
|
|
# del schema._children["item"]
|
||
|
|
# else:
|
||
|
|
# _prune_schema(schema.children.item)
|
||
|
|
# else: # if isinstance(schema, ObjectType):
|
||
|
|
# for child_name, child in vars(schema.children).items():
|
||
|
|
# if child not in retain:
|
||
|
|
# del schema._children[child_name]
|
||
|
|
# else:
|
||
|
|
# _prune_schema(child)
|
||
|
|
|
||
|
|
# _prune_schema(root)
|
||
|
|
|
||
|
|
# FIXME: PUNT
|
||
|
|
return root
|
||
|
|
|
||
|
|
|
||
|
|
def prune_obj_for_reasking(obj: Any) -> Union[None, Dict, List, ReAsk]:
|
||
|
|
"""After validation, we get a nested dictionary where some keys may be
|
||
|
|
ReAsk objects.
|
||
|
|
|
||
|
|
This function prunes the validated form of any object that is not a ReAsk object.
|
||
|
|
It also keeps all of the ancestors of the ReAsk objects.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
obj: The validated object.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
The pruned validated object.
|
||
|
|
"""
|
||
|
|
|
||
|
|
if isinstance(obj, ReAsk):
|
||
|
|
return obj
|
||
|
|
elif isinstance(obj, list):
|
||
|
|
pruned_list = []
|
||
|
|
for item in obj:
|
||
|
|
pruned_output = prune_obj_for_reasking(item)
|
||
|
|
if pruned_output is not None:
|
||
|
|
pruned_list.append(pruned_output)
|
||
|
|
if len(pruned_list):
|
||
|
|
return pruned_list
|
||
|
|
return None
|
||
|
|
elif isinstance(obj, dict):
|
||
|
|
pruned_json = {}
|
||
|
|
for key, value in obj.items():
|
||
|
|
if isinstance(value, FieldReAsk):
|
||
|
|
pruned_json[key] = value
|
||
|
|
elif isinstance(value, dict):
|
||
|
|
pruned_output = prune_obj_for_reasking(value)
|
||
|
|
if pruned_output is not None:
|
||
|
|
pruned_json[key] = pruned_output
|
||
|
|
elif isinstance(value, list):
|
||
|
|
pruned_list = []
|
||
|
|
for item in value:
|
||
|
|
pruned_output = prune_obj_for_reasking(item)
|
||
|
|
if pruned_output is not None:
|
||
|
|
pruned_list.append(pruned_output)
|
||
|
|
if len(pruned_list):
|
||
|
|
pruned_json[key] = pruned_list
|
||
|
|
|
||
|
|
if len(pruned_json):
|
||
|
|
return pruned_json
|
||
|
|
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def update_response_by_path(output: dict, path: List[Any], value: Any) -> None:
|
||
|
|
"""Update the output by path.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
output: The output.
|
||
|
|
path: The path to the element to be updated.
|
||
|
|
value: The value to be updated.
|
||
|
|
"""
|
||
|
|
for key in path[:-1]:
|
||
|
|
output = output[key]
|
||
|
|
output[path[-1]] = value
|
||
|
|
|
||
|
|
|
||
|
|
### Guard Execution Methods ###
|
||
|
|
def introspect(
|
||
|
|
data: Optional[Union[ReAsk, str, Dict, List]],
|
||
|
|
) -> Tuple[Sequence[ReAsk], Optional[Union[str, Dict, List]]]:
|
||
|
|
if isinstance(data, FieldReAsk):
|
||
|
|
return [data], None
|
||
|
|
elif isinstance(data, SkeletonReAsk):
|
||
|
|
return [data], None
|
||
|
|
elif isinstance(data, NonParseableReAsk):
|
||
|
|
return [data], None
|
||
|
|
return gather_reasks(data)
|
||
|
|
|
||
|
|
|
||
|
|
def get_reask_setup_for_string(
|
||
|
|
output_type: OutputTypes,
|
||
|
|
output_schema: Dict[str, Any],
|
||
|
|
validation_map: ValidatorMap,
|
||
|
|
reasks: Sequence[ReAsk],
|
||
|
|
*,
|
||
|
|
validation_response: Optional[Union[str, List, Dict, ReAsk]] = None,
|
||
|
|
prompt_params: Optional[Dict[str, Any]] = None,
|
||
|
|
exec_options: Optional[GuardExecutionOptions] = None,
|
||
|
|
) -> Tuple[Dict[str, Any], Messages]:
|
||
|
|
prompt_params = prompt_params or {}
|
||
|
|
exec_options = exec_options or GuardExecutionOptions()
|
||
|
|
|
||
|
|
schema_prompt_content = prompt_content_for_schema(
|
||
|
|
output_type, output_schema, validation_map
|
||
|
|
)
|
||
|
|
xml_output_schema = json_schema_to_rail_output(
|
||
|
|
json_schema=output_schema, validator_map=validation_map
|
||
|
|
)
|
||
|
|
|
||
|
|
reask_prompt_template = None
|
||
|
|
|
||
|
|
reask_prompt_template = Prompt(
|
||
|
|
constants["high_level_string_reask_prompt"]
|
||
|
|
+ constants["complete_string_suffix"]
|
||
|
|
)
|
||
|
|
|
||
|
|
error_messages = "\n".join(
|
||
|
|
[
|
||
|
|
f"- {fail_result.error_message}"
|
||
|
|
for reask in reasks
|
||
|
|
for fail_result in reask.fail_results or []
|
||
|
|
]
|
||
|
|
)
|
||
|
|
|
||
|
|
prompt = reask_prompt_template.format(
|
||
|
|
# FIXME: How do we properly type this?
|
||
|
|
# Solution will have to come from Runner all the way down to here
|
||
|
|
previous_response=validation_response.incorrect_value, # type: ignore
|
||
|
|
error_messages=error_messages,
|
||
|
|
output_schema=schema_prompt_content,
|
||
|
|
xml_output_schema=xml_output_schema,
|
||
|
|
**prompt_params,
|
||
|
|
)
|
||
|
|
|
||
|
|
instructions = None
|
||
|
|
instructions = Instructions("You are a helpful assistant.")
|
||
|
|
instructions = instructions.format(
|
||
|
|
output_schema=schema_prompt_content,
|
||
|
|
xml_output_schema=xml_output_schema,
|
||
|
|
**prompt_params,
|
||
|
|
)
|
||
|
|
messages = None
|
||
|
|
if exec_options.reask_messages:
|
||
|
|
messages = Messages(exec_options.reask_messages)
|
||
|
|
|
||
|
|
if messages is None:
|
||
|
|
messages = Messages(
|
||
|
|
[
|
||
|
|
{"role": "system", "content": instructions},
|
||
|
|
{"role": "user", "content": prompt},
|
||
|
|
]
|
||
|
|
)
|
||
|
|
|
||
|
|
messages = messages.format(
|
||
|
|
output_schema=schema_prompt_content,
|
||
|
|
xml_output_schema=xml_output_schema,
|
||
|
|
**prompt_params,
|
||
|
|
)
|
||
|
|
|
||
|
|
return output_schema, messages
|
||
|
|
|
||
|
|
|
||
|
|
def get_original_prompt(exec_options: Optional[GuardExecutionOptions] = None) -> str:
|
||
|
|
exec_options = exec_options or GuardExecutionOptions()
|
||
|
|
original_messages = exec_options.messages or []
|
||
|
|
messages_prompt = next(
|
||
|
|
(
|
||
|
|
h.get("content")
|
||
|
|
for h in original_messages
|
||
|
|
if isinstance(h, dict) and h.get("role") == "user"
|
||
|
|
),
|
||
|
|
"",
|
||
|
|
)
|
||
|
|
original_prompt = messages_prompt or ""
|
||
|
|
return original_prompt
|
||
|
|
|
||
|
|
|
||
|
|
def get_reask_setup_for_json(
|
||
|
|
output_type: OutputTypes,
|
||
|
|
output_schema: Dict[str, Any],
|
||
|
|
validation_map: ValidatorMap,
|
||
|
|
reasks: Sequence[ReAsk],
|
||
|
|
*,
|
||
|
|
parsing_response: Optional[Union[str, List, Dict, ReAsk]] = None,
|
||
|
|
validation_response: Optional[Union[str, List, Dict, ReAsk]] = None,
|
||
|
|
use_full_schema: Optional[bool] = False,
|
||
|
|
prompt_params: Optional[Dict[str, Any]] = None,
|
||
|
|
exec_options: Optional[GuardExecutionOptions] = None,
|
||
|
|
) -> Tuple[Dict[str, Any], Messages]:
|
||
|
|
reask_schema = output_schema
|
||
|
|
is_skeleton_reask = not any(isinstance(reask, FieldReAsk) for reask in reasks)
|
||
|
|
is_nonparseable_reask = any(
|
||
|
|
isinstance(reask, NonParseableReAsk) for reask in reasks
|
||
|
|
)
|
||
|
|
error_messages = {}
|
||
|
|
prompt_params = prompt_params or {}
|
||
|
|
exec_options = exec_options or GuardExecutionOptions()
|
||
|
|
original_prompt = get_original_prompt(exec_options)
|
||
|
|
use_xml = prompt_uses_xml(original_prompt)
|
||
|
|
|
||
|
|
reask_prompt_template = None
|
||
|
|
|
||
|
|
if is_nonparseable_reask:
|
||
|
|
if reask_prompt_template is None:
|
||
|
|
suffix = (
|
||
|
|
constants["xml_suffix_without_examples"]
|
||
|
|
if use_xml
|
||
|
|
else constants["json_suffix_without_examples"]
|
||
|
|
)
|
||
|
|
reask_prompt_template = Prompt(
|
||
|
|
constants["high_level_json_parsing_reask_prompt"] + suffix
|
||
|
|
)
|
||
|
|
np_reask: NonParseableReAsk = next(
|
||
|
|
r for r in reasks if isinstance(r, NonParseableReAsk)
|
||
|
|
)
|
||
|
|
# Give the LLM what it gave us that couldn't be parsed as JSON
|
||
|
|
reask_value = np_reask.incorrect_value
|
||
|
|
elif is_skeleton_reask:
|
||
|
|
if reask_prompt_template is None:
|
||
|
|
reask_prompt = constants["high_level_skeleton_reask_prompt"]
|
||
|
|
|
||
|
|
if use_xml:
|
||
|
|
reask_prompt = (
|
||
|
|
reask_prompt + constants["xml_suffix_with_structure_example"]
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
reask_prompt = (
|
||
|
|
reask_prompt
|
||
|
|
+ constants["error_messages"]
|
||
|
|
+ constants["json_suffix_with_structure_example"]
|
||
|
|
)
|
||
|
|
|
||
|
|
reask_prompt_template = Prompt(reask_prompt)
|
||
|
|
|
||
|
|
# Validation hasn't happend yet
|
||
|
|
# and the problem is with the json the LLM gave us.
|
||
|
|
# Give it this same json and tell it to fix it.
|
||
|
|
reask_value = validation_response if use_xml else parsing_response
|
||
|
|
skeleton_reask: SkeletonReAsk = next(
|
||
|
|
r for r in reasks if isinstance(r, SkeletonReAsk)
|
||
|
|
)
|
||
|
|
fail_results = skeleton_reask.fail_results or []
|
||
|
|
first_fail_result = fail_results[0]
|
||
|
|
error_messages = first_fail_result.error_message if first_fail_result else ""
|
||
|
|
else:
|
||
|
|
if use_full_schema:
|
||
|
|
# Give the LLM the full JSON that failed validation
|
||
|
|
reask_value = validation_response if use_xml else parsing_response
|
||
|
|
# Don't prune the tree if we're reasking with pydantic model
|
||
|
|
# (and openai function calling)
|
||
|
|
else:
|
||
|
|
# Prune out the individual fields that did not fail validation.
|
||
|
|
# Only reask for field that did fail.
|
||
|
|
reask_value = prune_obj_for_reasking(validation_response)
|
||
|
|
|
||
|
|
# Generate a subschema that matches the specific fields we're reasking for.
|
||
|
|
field_reasks = [r for r in reasks if isinstance(r, FieldReAsk)]
|
||
|
|
reask_schema = get_reask_subschema(output_schema, field_reasks)
|
||
|
|
|
||
|
|
if reask_prompt_template is None:
|
||
|
|
suffix = (
|
||
|
|
constants["xml_suffix_without_examples"]
|
||
|
|
if use_xml
|
||
|
|
else constants["json_suffix_without_examples"]
|
||
|
|
)
|
||
|
|
reask_prompt_template = Prompt(
|
||
|
|
constants["high_level_json_reask_prompt"] + suffix
|
||
|
|
)
|
||
|
|
|
||
|
|
error_messages = {
|
||
|
|
".".join(str(p) for p in r.path): "; ".join( # type: ignore
|
||
|
|
f.error_message for f in r.fail_results or []
|
||
|
|
)
|
||
|
|
for r in reasks
|
||
|
|
if isinstance(r, FieldReAsk)
|
||
|
|
}
|
||
|
|
|
||
|
|
stringified_schema = prompt_content_for_schema(
|
||
|
|
output_type, reask_schema, validation_map
|
||
|
|
)
|
||
|
|
xml_output_schema = json_schema_to_rail_output(
|
||
|
|
json_schema=output_schema, validator_map=validation_map
|
||
|
|
)
|
||
|
|
|
||
|
|
json_example = json.dumps(
|
||
|
|
generate_example(reask_schema),
|
||
|
|
indent=2,
|
||
|
|
)
|
||
|
|
|
||
|
|
def reask_decoder(obj: ReAsk):
|
||
|
|
decoded = {}
|
||
|
|
for k, v in obj.__dict__.items():
|
||
|
|
if k in ["path", "additional_properties"]:
|
||
|
|
continue
|
||
|
|
if k == "fail_results":
|
||
|
|
k = "error_messages"
|
||
|
|
v = [result.error_message for result in v]
|
||
|
|
decoded[k] = v
|
||
|
|
return decoded
|
||
|
|
|
||
|
|
prompt = reask_prompt_template.format(
|
||
|
|
previous_response=json.dumps(
|
||
|
|
reask_value, indent=2, default=reask_decoder, ensure_ascii=False
|
||
|
|
),
|
||
|
|
output_schema=stringified_schema,
|
||
|
|
xml_output_schema=xml_output_schema,
|
||
|
|
json_example=json_example,
|
||
|
|
error_messages=json.dumps(error_messages),
|
||
|
|
**prompt_params,
|
||
|
|
)
|
||
|
|
|
||
|
|
instructions = None
|
||
|
|
instructions_const = (
|
||
|
|
constants["high_level_xml_instructions"]
|
||
|
|
if use_xml
|
||
|
|
else constants["high_level_json_instructions"]
|
||
|
|
)
|
||
|
|
instructions = Instructions(instructions_const)
|
||
|
|
instructions = instructions.format(**prompt_params)
|
||
|
|
|
||
|
|
messages = None
|
||
|
|
if exec_options.reask_messages:
|
||
|
|
messages = Messages(exec_options.reask_messages)
|
||
|
|
else:
|
||
|
|
messages = Messages(
|
||
|
|
[
|
||
|
|
{"role": "system", "content": instructions},
|
||
|
|
{"role": "user", "content": prompt},
|
||
|
|
]
|
||
|
|
)
|
||
|
|
|
||
|
|
return reask_schema, messages
|
||
|
|
|
||
|
|
|
||
|
|
def get_reask_setup(
|
||
|
|
output_type: OutputTypes,
|
||
|
|
output_schema: Dict[str, Any],
|
||
|
|
validation_map: ValidatorMap,
|
||
|
|
reasks: Sequence[ReAsk],
|
||
|
|
*,
|
||
|
|
parsing_response: Optional[Union[str, List, Dict, ReAsk]] = None,
|
||
|
|
validation_response: Optional[Union[str, List, Dict, ReAsk]] = None,
|
||
|
|
use_full_schema: Optional[bool] = False,
|
||
|
|
prompt_params: Optional[Dict[str, Any]] = None,
|
||
|
|
exec_options: Optional[GuardExecutionOptions] = None,
|
||
|
|
) -> Tuple[Dict[str, Any], Messages]:
|
||
|
|
prompt_params = prompt_params or {}
|
||
|
|
exec_options = exec_options or GuardExecutionOptions()
|
||
|
|
|
||
|
|
if output_type == OutputTypes.STRING:
|
||
|
|
return get_reask_setup_for_string(
|
||
|
|
output_type=output_type,
|
||
|
|
output_schema=output_schema,
|
||
|
|
validation_map=validation_map,
|
||
|
|
reasks=reasks,
|
||
|
|
validation_response=validation_response,
|
||
|
|
prompt_params=prompt_params,
|
||
|
|
exec_options=exec_options,
|
||
|
|
)
|
||
|
|
return get_reask_setup_for_json(
|
||
|
|
output_type=output_type,
|
||
|
|
output_schema=output_schema,
|
||
|
|
validation_map=validation_map,
|
||
|
|
reasks=reasks,
|
||
|
|
parsing_response=parsing_response,
|
||
|
|
validation_response=validation_response,
|
||
|
|
use_full_schema=use_full_schema,
|
||
|
|
prompt_params=prompt_params,
|
||
|
|
exec_options=exec_options,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
### Post-Processing Methods ###
|
||
|
|
def gather_reasks(
|
||
|
|
validated_output: Optional[Union[ReAsk, str, Dict, List]],
|
||
|
|
) -> Tuple[List[ReAsk], Optional[Union[str, Dict, List]]]:
|
||
|
|
"""Traverse output and gather all ReAsk objects.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
validated_output (Union[str, Dict, ReAsk], optional): The output of a model.
|
||
|
|
Each value can be a ReAsk, a list, a dictionary, or a single value.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
A list of ReAsk objects found in the output.
|
||
|
|
"""
|
||
|
|
if validated_output is None:
|
||
|
|
return [], None
|
||
|
|
if isinstance(validated_output, ReAsk):
|
||
|
|
return [validated_output], None
|
||
|
|
if isinstance(validated_output, str):
|
||
|
|
return [], validated_output
|
||
|
|
|
||
|
|
reasks = []
|
||
|
|
|
||
|
|
def _gather_reasks_in_dict(
|
||
|
|
original: Dict, valid_output: Dict, path: Optional[List[Union[str, int]]] = None
|
||
|
|
) -> None:
|
||
|
|
if path is None:
|
||
|
|
path = []
|
||
|
|
for field, value in original.items():
|
||
|
|
if isinstance(value, FieldReAsk):
|
||
|
|
value.path = path + [field]
|
||
|
|
reasks.append(value)
|
||
|
|
del valid_output[field]
|
||
|
|
|
||
|
|
if isinstance(value, dict):
|
||
|
|
_gather_reasks_in_dict(value, valid_output[field], path + [field])
|
||
|
|
|
||
|
|
if isinstance(value, list):
|
||
|
|
_gather_reasks_in_list(value, valid_output[field], path + [field])
|
||
|
|
return
|
||
|
|
|
||
|
|
def _gather_reasks_in_list(
|
||
|
|
original: List, valid_output: List, path: Optional[List[Union[str, int]]] = None
|
||
|
|
) -> None:
|
||
|
|
if path is None:
|
||
|
|
path = []
|
||
|
|
for idx, item in enumerate(original):
|
||
|
|
if isinstance(item, FieldReAsk):
|
||
|
|
item.path = path + [idx]
|
||
|
|
reasks.append(item)
|
||
|
|
del valid_output[idx]
|
||
|
|
elif isinstance(item, dict):
|
||
|
|
_gather_reasks_in_dict(item, valid_output[idx], path + [idx])
|
||
|
|
elif isinstance(item, list):
|
||
|
|
_gather_reasks_in_list(item, valid_output[idx], path + [idx])
|
||
|
|
return
|
||
|
|
|
||
|
|
if isinstance(validated_output, Dict):
|
||
|
|
valid_output = deepcopy(validated_output)
|
||
|
|
_gather_reasks_in_dict(validated_output, valid_output)
|
||
|
|
return reasks, valid_output
|
||
|
|
elif isinstance(validated_output, List):
|
||
|
|
valid_output = deepcopy(validated_output)
|
||
|
|
_gather_reasks_in_list(validated_output, valid_output)
|
||
|
|
return reasks, valid_output
|
||
|
|
return reasks, None
|
||
|
|
|
||
|
|
|
||
|
|
def sub_reasks_with_fixed_values(value: Any) -> Any:
|
||
|
|
"""Substitute ReAsk objects with their fixed values recursively.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
value: Either a list, a dictionary, a ReAsk object or a scalar value.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
The value with ReAsk objects replaced with their fixed values.
|
||
|
|
"""
|
||
|
|
copy = deepcopy(value)
|
||
|
|
if isinstance(copy, list):
|
||
|
|
for index, item in enumerate(copy):
|
||
|
|
copy[index] = sub_reasks_with_fixed_values(item)
|
||
|
|
elif isinstance(copy, dict):
|
||
|
|
for dict_key, dict_value in value.items():
|
||
|
|
copy[dict_key] = sub_reasks_with_fixed_values(dict_value)
|
||
|
|
elif isinstance(copy, FieldReAsk):
|
||
|
|
fail_results = copy.fail_results or []
|
||
|
|
first_fail_result = fail_results[0]
|
||
|
|
|
||
|
|
fix_value = first_fail_result.fix_value if first_fail_result else None
|
||
|
|
# TODO handle multiple fail results
|
||
|
|
# Leave the ReAsk in place if there is no fix value
|
||
|
|
# This allows us to determine the proper status for the call
|
||
|
|
copy = fix_value if fix_value is not None else copy
|
||
|
|
|
||
|
|
return copy
|
||
|
|
|
||
|
|
|
||
|
|
def merge_reask_output(previous_response, reask_response) -> Dict:
|
||
|
|
"""Merge the reask output into the original output.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
prev_logs: validation output object from the previous iteration.
|
||
|
|
current_logs: validation output object from the current iteration.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
The merged output.
|
||
|
|
"""
|
||
|
|
if isinstance(previous_response, ReAsk):
|
||
|
|
return reask_response
|
||
|
|
|
||
|
|
# FIXME: Uncommenet when field level reask is fixed
|
||
|
|
# This used to be necessary for field level reask because
|
||
|
|
# the schema was pruned to only the properties that failed.
|
||
|
|
# This caused previous keys that were correct to be pruned during schemafication.
|
||
|
|
# pruned_reask_json = prune_obj_for_reasking(previous_response)
|
||
|
|
pruned_reask_json = previous_response
|
||
|
|
|
||
|
|
# Reask output and reask json have the same structure, except that values
|
||
|
|
# of the reask json are ReAsk objects. We want to replace the ReAsk objects
|
||
|
|
# with the values from the reask output.
|
||
|
|
merged_json = deepcopy(previous_response)
|
||
|
|
|
||
|
|
def update_reasked_elements(pruned_reask_json, reask_response_dict):
|
||
|
|
if isinstance(pruned_reask_json, dict):
|
||
|
|
for key, value in pruned_reask_json.items():
|
||
|
|
if isinstance(value, FieldReAsk):
|
||
|
|
if value.path is None:
|
||
|
|
raise RuntimeError(
|
||
|
|
"FieldReAsk object must have a path attribute."
|
||
|
|
)
|
||
|
|
corrected_value = reask_response_dict.get(key)
|
||
|
|
update_response_by_path(merged_json, value.path, corrected_value)
|
||
|
|
else:
|
||
|
|
update_reasked_elements(
|
||
|
|
pruned_reask_json[key], reask_response_dict[key]
|
||
|
|
)
|
||
|
|
elif isinstance(pruned_reask_json, list):
|
||
|
|
for i, item in enumerate(pruned_reask_json):
|
||
|
|
if isinstance(item, FieldReAsk):
|
||
|
|
if item.path is None:
|
||
|
|
raise RuntimeError(
|
||
|
|
"FieldReAsk object must have a path attribute."
|
||
|
|
)
|
||
|
|
corrected_value = reask_response_dict[i]
|
||
|
|
update_response_by_path(merged_json, item.path, corrected_value)
|
||
|
|
else:
|
||
|
|
update_reasked_elements(
|
||
|
|
pruned_reask_json[i], reask_response_dict[i]
|
||
|
|
)
|
||
|
|
|
||
|
|
update_reasked_elements(pruned_reask_json, reask_response)
|
||
|
|
|
||
|
|
return merged_json
|