[HF Maintenance] Fix the json errors (#2)

- [HF Maintenance] Fix the json errors (4eb73651d1388d0e2b98c20359203a134f2a4adc)

prompt_augment.py

@@ -212,7 +212,26 @@ Please strictly follow the rewriting rules below:
                 time.sleep(1)
 
         # polished_prompt = json.loads(completion.choices[0].message.content)['Rewritten']
-        polished_prompt = json.loads(response.output.choices[0].message.content[0]['text'])['Rewritten']
+        try:
+            raw_text = response.output.choices[0].message.content[0]['text']
+            print(f"Raw API response: {repr(raw_text[:500])}")
+            
+            # The thinking model may wrap JSON in markdown or include <think> tags
+            # Try to extract JSON from the response
+            json_match = re.search(r'\{[^{}]*"Rewritten"\s*:\s*"[^"]*"[^{}]*\}', raw_text, re.DOTALL)
+            if json_match:
+                polished_prompt = json.loads(json_match.group())['Rewritten']
+            else:
+                # More flexible: find any JSON block (possibly in ```json ... ```)
+                json_block = re.search(r'```json\s*(\{.*?\})\s*```', raw_text, re.DOTALL)
+                if json_block:
+                    polished_prompt = json.loads(json_block.group(1))['Rewritten']
+                else:
+                    print("Could not extract JSON from response, using original prompt")
+                    polished_prompt = original_prompt
+        except Exception as e:
+            print(f"Prompt rewrite failed: {e}, using original prompt")
+            polished_prompt = original_prompt
         
         return polished_prompt  # + magic_prompt