improve response cleaning

The prompt for the model starts with a code block (```). When testing watsonx models (llama and granite), they would generate the closing block in the response.
This commit is contained in:
Robby
2024-07-29 10:26:58 -04:00
committed by GitHub
parent 41588efe9a
commit 172d0c0358

View File

@ -557,7 +557,7 @@ def _fix_key_value(key: str, value: str):
def load_yaml(response_text: str, keys_fix_yaml: List[str] = [], first_key="", last_key="") -> dict:
response_text = response_text.strip('\n').removeprefix('```yaml').rstrip('`')
response_text = response_text.strip().removeprefix("```yaml").removesuffix("```").strip()
try:
data = yaml.safe_load(response_text)
except Exception as e: