pvanand commited on
Commit
a9b0058
·
verified ·
1 Parent(s): 6b95868

update yaml parsing

Browse files
Files changed (1) hide show
  1. main.py +6 -3
main.py CHANGED
@@ -416,8 +416,11 @@ def parse_followup_response(response):
416
  if len(clarification_parts) > 1:
417
  clarification_yaml = clarification_parts[1].split("</clarification>")[0].strip()
418
  try:
419
- clarification = yaml.safe_load(clarification_yaml)
420
- except yaml.YAMLError:
 
 
 
421
  clarification = None
422
  else:
423
  clarification = None
@@ -453,7 +456,7 @@ async def followup_agent(query: FollowupQueryModel, background_tasks: Background
453
 
454
  result = {
455
  "response": response_content,
456
- "clarification": clarification
457
  }
458
 
459
  yield "\n\n" + json.dumps(result)
 
416
  if len(clarification_parts) > 1:
417
  clarification_yaml = clarification_parts[1].split("</clarification>")[0].strip()
418
  try:
419
+ # Add indentation to make it valid YAML
420
+ indented_yaml = "\n".join(" " + line for line in clarification_yaml.split("\n"))
421
+ clarification = yaml.load("questions:\n" + indented_yaml, Loader=SafeLoader)
422
+ except yaml.YAMLError as e:
423
+ logger.error(f"YAML parsing error: {e}")
424
  clarification = None
425
  else:
426
  clarification = None
 
456
 
457
  result = {
458
  "response": response_content,
459
+ "clarification": clarification['questions'] if clarification else None
460
  }
461
 
462
  yield "\n\n" + json.dumps(result)