Update app.py

app.py

@@ -1,22 +1,22 @@
 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, Dict
+from typing import Optional, Literal
 from huggingface_hub.errors import HfHubHTTPError
 
-
-# Input model
 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")
 
-# Output model for LLM responses
-class LLMResponse(BaseModel):
-    initial_prompt_evaluation: str = Field(default="")
-    refined_prompt: str = Field(default="")
-    explanation_of_refinements: str = Field(default="")
+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 PromptRefiner:
     def __init__(self, api_token: str):
@@ -34,23 +34,19 @@ 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 prompts. Respond in JSON format with exactly these fields:
-                    {
-                        "initial_prompt_evaluation": "your evaluation",
-                        "refined_prompt": "your refined prompt",
-                        "explanation_of_refinements": "your explanation"
-                    }'''
+                    "role": "system", 
+                    "content": 'You are an expert at refining and extending prompts. Given a basic prompt, provide a more detailed.'
                 },
                 {
-                    "role": "user",
+                    "role": "user", 
                     "content": selected_meta_prompt.replace("[Insert initial prompt here]", prompt_input.text)
                 }
             ]
@@ -59,81 +55,101 @@ class PromptRefiner:
                 model=prompt_refiner_model,
                 messages=messages,
                 max_tokens=2000,
-                temperature=0.8,
-                response_format={"type": "json_object"}
+                temperature=0.8
             )
-
-            # Parse response using Pydantic
+            
             response_content = response.choices[0].message.content.strip()
-            try:
-                parsed_response = LLMResponse.model_validate_json(response_content)
-                result = parsed_response.model_dump()
-            except Exception:
-                # Fallback to basic dict if JSON parsing fails
-                result = {
-                    "initial_prompt_evaluation": response_content,
-                    "refined_prompt": prompt_input.text,
-                    "explanation_of_refinements": "Failed to parse model response"
-                }
-
+            
+            # Parse the response
+            result = self._parse_response(response_content)
+            
             return (
-                result["initial_prompt_evaluation"],
-                result["refined_prompt"],
-                result["explanation_of_refinements"],
+                result.get('initial_prompt_evaluation', ''),
+                result.get('refined_prompt', ''),
+                result.get('explanation_of_refinements', ''),
                 result
             )
 
         except HfHubHTTPError as e:
-            error_response = LLMResponse(
-                initial_prompt_evaluation="Error: Model timeout or connection issue",
-                refined_prompt=prompt_input.text,
-                explanation_of_refinements="Please try again in a few moments"
-            ).model_dump()
             return (
-                error_response["initial_prompt_evaluation"],
-                error_response["refined_prompt"],
-                error_response["explanation_of_refinements"],
-                error_response
+                "Error: Model timeout. Please try again later.",
+                "The selected model is currently experiencing high traffic.",
+                "The selected model is currently experiencing high traffic.",
+                {}
             )
         except Exception as e:
-            error_response = LLMResponse(
-                initial_prompt_evaluation=f"Error: {str(e)}",
-                refined_prompt=prompt_input.text,
-                explanation_of_refinements="An unexpected error occurred"
-            ).model_dump()
             return (
-                error_response["initial_prompt_evaluation"],
-                error_response["refined_prompt"],
-                error_response["explanation_of_refinements"],
-                error_response
+                f"Error: {str(e)}",
+                "",
+                "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)
+                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."
+                    "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
             )
-
-            return response.choices[0].message.content.strip().replace('\n\n', '\n')
-
+            
+            output = response.choices[0].message.content.strip()
+            return output.replace('\n\n', '\n').strip()
+            
         except Exception as e:
             return f"Error: {str(e)}"
 
-
 class GradioInterface:
     def __init__(self, prompt_refiner: PromptRefiner):
         self.prompt_refiner = prompt_refiner