Update app.py

app.py

@@ -20,43 +20,59 @@ class RefinementOutput(BaseModel):
 
 class PromptRefiner:
     def __init__(self, api_token: str):
-        self.client = InferenceClient(token=api_token,timeout=300)
-
-    def refine_prompt(self, prompt_input: PromptInput) -> RefinementOutput:
-        if prompt_input.meta_prompt_choice == "morphosis":
-            selected_meta_prompt = original_meta_prompt
-        elif prompt_input.meta_prompt_choice == "verse":
-            selected_meta_prompt = new_meta_prompt
-        elif prompt_input.meta_prompt_choice == "physics":
-            selected_meta_prompt = metaprompt1
-        elif prompt_input.meta_prompt_choice == "bolism":
-            selected_meta_prompt = loic_metaprompt
-        elif prompt_input.meta_prompt_choice == "done":
-            selected_meta_prompt = metadone
-        elif prompt_input.meta_prompt_choice == "star":
-            selected_meta_prompt = echo_prompt_refiner
-        elif prompt_input.meta_prompt_choice == "math":
-            selected_meta_prompt = math_meta_prompt  
-        elif prompt_input.meta_prompt_choice == "arpe":
-            selected_meta_prompt = autoregressive_metaprompt
-        else:
-            selected_meta_prompt = advanced_meta_prompt
-    
-        messages = [
-            {"role": "system", "content": 'You are an expert at refining and extending prompts. Given a basic prompt, provide a more detailed.'},
-            {"role": "user", "content": selected_meta_prompt.replace("[Insert initial prompt here]", prompt_input.text)}
-        ]
+        self.client = InferenceClient(token=api_token, timeout=300)
+        self.meta_prompts = {
+            "morphosis": original_meta_prompt,
+            "verse": new_meta_prompt,
+            "physics": metaprompt1,
+            "bolism": loic_metaprompt,
+            "done": metadone,
+            "star": echo_prompt_refiner,
+            "math": math_meta_prompt,
+            "arpe": autoregressive_metaprompt
+        }
+
+    def refine_prompt(self, prompt_input: PromptInput) -> tuple:
         try:
+            # Select meta prompt using dictionary instead of if-elif chain
+            selected_meta_prompt = self.meta_prompts.get(
+                prompt_input.meta_prompt_choice, 
+                advanced_meta_prompt
+            )
+            
+            messages = [
+                {
+                    "role": "system", 
+                    "content": 'You are an expert at refining and extending prompts. Given a basic prompt, provide a more detailed.'
+                },
+                {
+                    "role": "user", 
+                    "content": selected_meta_prompt.replace("[Insert initial prompt here]", prompt_input.text)
+                }
+            ]
+
             response = self.client.chat_completion(
                 model=prompt_refiner_model,
                 messages=messages,
                 max_tokens=2000,
                 temperature=0.8
             )
+            
             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 HfHubHTTPError as e:
             return (
-                "Error: Model timeout. Please again later.",
+                "Error: Model timeout. Please try again later.",
                 "",
                 "The selected model is currently experiencing high traffic.",
                 {}
@@ -68,49 +84,67 @@ class PromptRefiner:
                 "An unexpected error occurred.",
                 {}
             )
+
+    def _parse_response(self, response_content: str) -> dict:
         try:
+            # 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)
                 json_str = re.sub(r'\n\s*', ' ', json_str)
                 json_str = json_str.replace('"', '\\"')
                 json_output = json.loads(f'"{json_str}"')
+                
                 if isinstance(json_output, str):
                     json_output = json.loads(json_output)
-                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}")
+                
+                # Clean up JSON values
+                return {
+                    key: value.replace('\\"', '"') if isinstance(value, str) else value
+                    for key, value in json_output.items()
+                }
+            
+            # 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] = ""
-            return RefinementOutput(**output, raw_content=response_content)
-                
+                output[key] = match.group(1).replace('\\n', '\n').replace('\\"', '"') if match else ""
+            
+            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)
+            }
+
     def apply_prompt(self, prompt: str, model: str) -> str:
         try:
             messages = [
-                {"role": "system", "content": "You are a helpful assistant. Answer in stylized version with latex format or markdown if relevant. Separate your answer into logical sections using level 2 headers (##) for sections and bolding (**) for subsections.Incorporate a variety of lists, headers, and text to make the answer visually appealing"},
-                {"role": "user", "content": prompt}
+                {
+                    "role": "system",
+                    "content": "You are a helpful assistant. Answer in stylized version with latex format or markdown if relevant. Separate your answer into logical sections using level 2 headers (##) for sections and bolding (**) for subsections. Incorporate a variety of lists, headers, and text to make the answer visually appealing"
+                },
+                {
+                    "role": "user",
+                    "content": prompt
+                }
             ]
+            
             response = self.client.chat_completion(
                 model=model,
                 messages=messages,
                 max_tokens=2000,
                 temperature=0.8
             )
+            
             output = response.choices[0].message.content.strip()
-            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)}"