ignore header description in ai response

This commit is contained in:
koid
2023-12-20 11:13:14 +09:00
parent 68c26b362b
commit 16b61eb4e8
2 changed files with 18 additions and 0 deletions

View File

@ -317,6 +317,12 @@ def _fix_key_value(key: str, value: str):
def load_yaml(response_text: str) -> dict: def load_yaml(response_text: str) -> dict:
# remove everything before the first ```yaml
snipet_pattern = r'```(yaml)?[\s\S]*?```'
snipet = re.search(snipet_pattern, response_text)
if snipet:
response_text = snipet.group()
response_text = response_text.removeprefix('```yaml').rstrip('`') response_text = response_text.removeprefix('```yaml').rstrip('`')
try: try:
data = yaml.safe_load(response_text) data = yaml.safe_load(response_text)

View File

@ -15,6 +15,18 @@ class TestLoadYaml:
expected_output = {'name': 'John Smith', 'age': 35} expected_output = {'name': 'John Smith', 'age': 35}
assert load_yaml(yaml_str) == expected_output assert load_yaml(yaml_str) == expected_output
def test_load_valid_yaml_with_description(self):
yaml_str = '''\
Here is the answer in YAML format:
```yaml
name: John Smith
age: 35
```
'''
expected_output = {'name': 'John Smith', 'age': 35}
assert load_yaml(yaml_str) == expected_output
def test_load_invalid_yaml1(self): def test_load_invalid_yaml1(self):
yaml_str = \ yaml_str = \
'''\ '''\