Update app.py

app.py

@@ -57,55 +57,39 @@ class PromptRefiner:
             temperature=0.7
         )
         response_content = response.choices[0].message.content.strip()
-
-        # Parse the response
-        result = self._parse_response(response_content)
-            
-        return (
-            result.get('initial_prompt_evaluation', ''),
-            result.get('refined_prompt', ''),
-            result.get('explanation_of_refinements', ''),
-            result
-            )
-
-    def _parse_response(self, response_content: str) -> dict:
         try:
-            # Try to find JSON in response
+            # Extract JSON from between <json> tags
             json_match = re.search(r'<json>\s*(.*?)\s*</json>', response_content, re.DOTALL)
             if json_match:
                 json_str = json_match.group(1)
+                # Remove newlines and escape quotes within the JSON string
                 json_str = re.sub(r'\n\s*', ' ', json_str)
                 json_str = json_str.replace('"', '\\"')
+                # Wrap the entire string in quotes and parse it
                 json_output = json.loads(f'"{json_str}"')
-                
+                # Ensure json_output is a dictionary
                 if isinstance(json_output, str):
                     json_output = json.loads(json_output)
-                output={
-                    key: value.replace('\\"', '"') if isinstance(value, str) else value
-                    for key, value in json_output.items()
-                }
-                output['response_content']=json_output
-                # Clean up JSON values
-                return output
-            
-            # Fallback to regex parsing if no JSON found
+                # Unescape the parsed JSON
+                for key, value in json_output.items():
+                    if isinstance(value, str):
+                        json_output[key] = value.replace('\\"', '"')
+                return RefinementOutput(**json_output, raw_content=response_content)
+            else:
+                raise ValueError("No JSON found in the response")
+        except (json.JSONDecodeError, ValueError) as e:
+            print(f"Error parsing JSON: {e}")
+            print(f"Raw content: {response_content}")
+            # If JSON parsing fails, attempt to extract the content manually
             output = {}
             for key in ["initial_prompt_evaluation", "refined_prompt", "explanation_of_refinements"]:
                 pattern = rf'"{key}":\s*"(.*?)"(?:,|\}})'
                 match = re.search(pattern, response_content, re.DOTALL)
-                output[key] = match.group(1).replace('\\n', '\n').replace('\\"', '"') if match else ""
-            output['response_content']=response_content
-            return output
-
-        except (json.JSONDecodeError, ValueError) as e:
-            print(f"Error parsing response: {e}")
-            print(f"Raw content: {response_content}")
-            return {
-                "initial_prompt_evaluation": "Error parsing response",
-                "refined_prompt": "",
-                "explanation_of_refinements": str(e),
-                'response_content':str(e)
-            }
+                if match:
+                    output[key] = match.group(1).replace('\\n', '\n').replace('\\"', '"')
+                else:
+                    output[key] = "" # Set empty string if content not found
+            return RefinementOutput(**output, raw_content=response_content)
                 
     def apply_prompt(self, prompt: str, model: str) -> str:
         try: