Update app.py

app.py

@@ -1,27 +1,24 @@
 import os
 import json
-import re
 from huggingface_hub import InferenceClient
 import gradio as gr
 from pydantic import BaseModel, Field
-from typing import Optional, Literal
+from typing import Optional, Literal, Dict
 from huggingface_hub.errors import HfHubHTTPError
 
 class PromptInput(BaseModel):
     text: str = Field(..., description="The initial prompt text")
-    meta_prompt_choice: Literal["star","done","physics","morphosis", "verse", "phor","bolism","math","arpe"] = Field(..., description="Choice of meta prompt strategy")
+    meta_prompt_choice: Literal["star", "done", "physics", "morphosis", "verse", "phor", "bolism", "math", "arpe"] = Field(..., description="Choice of meta prompt strategy")
 
-class RefinementOutput(BaseModel):
-    query_analysis: Optional[str] = None
-    initial_prompt_evaluation: Optional[str] = None
-    refined_prompt: Optional[str] = None
-    explanation_of_refinements: Optional[str] = None
-    raw_content: Optional[str] = None
+class LLMResponse(BaseModel):
+    initial_prompt_evaluation: str = Field(default="")
+    refined_prompt: str = Field(default="")
+    explanation_of_refinements: str = Field(default="")
 
 class PromptRefiner:
     def __init__(self, api_token: str):
         self.client = InferenceClient(token=api_token, timeout=300)
-        self.meta_prompts = {
+        self.meta_prompts: Dict[str, str] = {
             "morphosis": original_meta_prompt,
             "verse": new_meta_prompt,
             "physics": metaprompt1,
@@ -34,19 +31,23 @@ class PromptRefiner:
 
     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, 
+                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": "system",
+                    "content": '''You are an expert at refining prompts. Respond in JSON format with exactly these fields:
+                    {
+                        "initial_prompt_evaluation": "your evaluation of the initial prompt",
+                        "refined_prompt": "your refined version of the prompt",
+                        "explanation_of_refinements": "your explanation of the changes made"
+                    }'''
                 },
                 {
-                    "role": "user", 
+                    "role": "user",
                     "content": selected_meta_prompt.replace("[Insert initial prompt here]", prompt_input.text)
                 }
             ]
@@ -55,101 +56,81 @@ class PromptRefiner:
                 model=prompt_refiner_model,
                 messages=messages,
                 max_tokens=2000,
-                temperature=0.8
+                temperature=0.8,
+                response_format={"type": "json_object"}
             )
-            
+
+            # Parse response using Pydantic
             response_content = response.choices[0].message.content.strip()
-            
-            # Parse the response
-            result = self._parse_response(response_content)
-            
+            try:
+                parsed_response = LLMResponse.model_validate_json(response_content)
+                result = parsed_response.model_dump()
+            except Exception as e:
+                # Fallback to basic dict if JSON parsing fails
+                result = {
+                    "initial_prompt_evaluation": "Error parsing model response",
+                    "refined_prompt": response_content,
+                    "explanation_of_refinements": str(e)
+                }
+
             return (
-                result.get('initial_prompt_evaluation', ''),
-                result.get('refined_prompt', ''),
-                result.get('explanation_of_refinements', ''),
+                result["initial_prompt_evaluation"],
+                result["refined_prompt"],
+                result["explanation_of_refinements"],
                 result
             )
 
         except HfHubHTTPError as e:
+            error_response = LLMResponse(
+                initial_prompt_evaluation="Error: Model timeout",
+                refined_prompt="The model is currently experiencing high traffic",
+                explanation_of_refinements="Please try again later"
+            )
             return (
-                "Error: Model timeout. Please try again later.",
-                "The selected model is currently experiencing high traffic.",
-                "The selected model is currently experiencing high traffic.",
-                {}
+                error_response.initial_prompt_evaluation,
+                error_response.refined_prompt,
+                error_response.explanation_of_refinements,
+                error_response.model_dump()
             )
         except Exception as e:
+            error_response = LLMResponse(
+                initial_prompt_evaluation=f"Error: {str(e)}",
+                refined_prompt="",
+                explanation_of_refinements="An unexpected error occurred"
+            )
             return (
-                f"Error: {str(e)}",
-                "",
-                "An unexpected error occurred.",
-                {}
+                error_response.initial_prompt_evaluation,
+                error_response.refined_prompt,
+                error_response.explanation_of_refinements,
+                error_response.model_dump()
             )
 
-    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)
-                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)
-                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:
             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"
+                    "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."
                 },
                 {
                     "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()
-            return output.replace('\n\n', '\n').strip()
-            
+
+            return response.choices[0].message.content.strip().replace('\n\n', '\n')
+
         except Exception as e:
             return f"Error: {str(e)}"
 
+
 class GradioInterface:
     def __init__(self, prompt_refiner: PromptRefiner):
         self.prompt_refiner = prompt_refiner