keys fallback

This commit is contained in:
mrT23
2024-07-03 20:29:17 +03:00
parent 0bac03496a
commit e30c70d2ca
2 changed files with 3 additions and 2 deletions

View File

@ -565,7 +565,7 @@ def try_fix_yaml(response_text: str,
except: except:
pass pass
# forth fallback - try to extract yaml snippet by first_key and 'last_key' # forth fallback - try to extract yaml snippet by 'first_key' and 'last_key'
# note that 'last_key' can be in practice a key that is not the last key in the yaml snippet. # note that 'last_key' can be in practice a key that is not the last key in the yaml snippet.
# it just needs to be some inner key, so we can look for newlines after it # it just needs to be some inner key, so we can look for newlines after it
if first_key and last_key: if first_key and last_key:

View File

@ -231,7 +231,8 @@ class PRCodeSuggestions:
def _prepare_pr_code_suggestions(self, predictions: str) -> Dict: def _prepare_pr_code_suggestions(self, predictions: str) -> Dict:
data = load_yaml(predictions.strip(), data = load_yaml(predictions.strip(),
keys_fix_yaml=["relevant_file", "suggestion_content", "existing_code", "improved_code"]) keys_fix_yaml=["relevant_file", "suggestion_content", "existing_code", "improved_code"],
first_key="code_suggestions",last_key="label")
if isinstance(data, list): if isinstance(data, list):
data = {'code_suggestions': data} data = {'code_suggestions': data}