Update app.py

app.py

@@ -57,39 +57,62 @@ 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
+            )
+
+        except Exception as e:
+            return (
+                f"Error: {str(e)}",
+                "",
+                "An unexpected error occurred.",
+                {}
+            )        
+    def _parse_response(self, response_content: str) -> dict:
         try:
-            # Extract JSON from between <json> tags
+            # Try to find JSON in response
             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)
-                # 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={
+                    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
             output = {}
             for key in ["initial_prompt_evaluation", "refined_prompt", "explanation_of_refinements"]:
                 pattern = rf'"{key}":\s*"(.*?)"(?:,|\}})'
                 match = re.search(pattern, response_content, re.DOTALL)
-                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)
+                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)
+            }
                 
     def apply_prompt(self, prompt: str, model: str) -> str:
         try:
@@ -200,7 +223,6 @@ class PromptRefiner:
                     "content": prompt
                 }
             ]'''
-
             response = self.client.chat.completions.create(
                 model=model,
                 messages=messages,
@@ -210,8 +232,8 @@ class PromptRefiner:
 
             output = response.choices[0].message.content.strip()
             # Basic post-processing
-            output = output.replace('\n\n', '\n').strip()
-            return output
+            return output.replace('\n\n', '\n').strip()
+           
         except Exception as e:
             return f"Error: {str(e)}"