Update prompt_refiner.py

prompt_refiner.py

@@ -1,6 +1,6 @@
 import json
 import re
-from typing import Optional, Dict, Any
+from typing import Optional, Dict, Any, Union
 from pydantic import BaseModel, Field, validator
 from huggingface_hub import InferenceClient
 from huggingface_hub.errors import HfHubHTTPError
@@ -9,163 +9,96 @@ from variables import *
 class LLMResponse(BaseModel):
     initial_prompt_evaluation: str = Field(..., description="Evaluation of the initial prompt")
     refined_prompt: str = Field(..., description="The refined version of the prompt")
-    explanation_of_refinements: str = Field(..., description="Explanation of the refinements made")
+    explanation_of_refinements: Union[str, list] = Field(..., description="Explanation of the refinements made")
     response_content: Optional[Dict[str, Any]] = Field(None, description="Raw response content")
 
-    @validator('initial_prompt_evaluation', 'refined_prompt', 'explanation_of_refinements')
+    @validator('initial_prompt_evaluation', 'refined_prompt')
     def clean_text_fields(cls, v):
         if isinstance(v, str):
             return v.strip().replace('\\n', '\n').replace('\\"', '"')
         return v
 
+    @validator('explanation_of_refinements')
+    def clean_refinements(cls, v):
+        if isinstance(v, str):
+            return v.strip().replace('\\n', '\n').replace('\\"', '"')
+        elif isinstance(v, list):
+            return [item.strip().replace('\\n', '\n').replace('\\"', '"') if isinstance(item, str) else item for item in v]
+        return v
+
 class PromptRefiner:
     def __init__(self, api_token: str, meta_prompts):
         self.client = InferenceClient(token=api_token, timeout=120)
         self.meta_prompts = meta_prompts
-        
-    def refine_prompt(self, prompt: str, meta_prompt_choice: str) -> tuple:
-        try:
-            selected_meta_prompt = self.meta_prompts.get(
-                meta_prompt_choice, 
-                self.meta_prompts["star"]
-            )
-            
-            messages = [
-                {
-                    "role": "system", 
-                    "content": 'You are an expert at refining and extending prompts. Given a basic prompt, provide a more relevant and detailed prompt.'
-                },
-                {
-                    "role": "user", 
-                    "content": selected_meta_prompt.replace("[Insert initial prompt here]", prompt)
-                }
-            ]
-            
-            response = self.client.chat_completion(
-                model=prompt_refiner_model,
-                messages=messages,
-                max_tokens=3000,
-                temperature=0.8
-            )
-            
-            response_content = response.choices[0].message.content.strip()
-            result = self._parse_response(response_content)
-            
-            # Create and validate LLMResponse
-            llm_response = LLMResponse(**result)
-            
-            return (
-                llm_response.initial_prompt_evaluation,
-                llm_response.refined_prompt,
-                llm_response.explanation_of_refinements,
-                llm_response.dict()
-            )
 
-        except HfHubHTTPError as e:
-            return self._create_error_response("Model timeout. Please try again later.")
-        except Exception as e:
-            return self._create_error_response(f"Unexpected error: {str(e)}")
+    def _sanitize_json_string(self, json_str: str) -> str:
+        """Clean and prepare JSON string for parsing."""
+        json_str = json_str.lstrip('\ufeff').strip()
+        json_str = json_str.replace('\n', ' ')
+        json_str = re.sub(r'\s+', ' ', json_str)
+        json_str = json_str.replace('•', '*')
+        return json_str
 
-    def _create_error_response(self, error_message: str) -> tuple:
-        error_response = LLMResponse(
-            initial_prompt_evaluation=f"Error: {error_message}",
-            refined_prompt="The selected model is currently unavailable.",
-            explanation_of_refinements="An error occurred during processing.",
-            response_content={"error": error_message}
-        )
-        return (
-            error_response.initial_prompt_evaluation,
-            error_response.refined_prompt,
-            error_response.explanation_of_refinements,
-            error_response.dict()
-        )
+    def _extract_json_content(self, content: str) -> str:
+        """Extract JSON content from between <json> tags."""
+        json_match = re.search(r'<json>\s*(.*?)\s*</json>', content, re.DOTALL)
+        if json_match:
+            return self._sanitize_json_string(json_match.group(1))
+        return content
 
     def _parse_response(self, response_content: str) -> dict:
         try:
-            # First attempt: Try to extract JSON from <json> tags
-            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)
-                
-                return {
-                    "initial_prompt_evaluation": json_output.get("initial_prompt_evaluation", ""),
-                    "refined_prompt": json_output.get("refined_prompt", ""),
-                    "explanation_of_refinements": json_output.get("explanation_of_refinements", ""),
-                    "response_content": json_output
-                }
-            
-            # Second attempt: Try to extract fields using regex
-            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) if match else ""
-            
-            output["response_content"] = response_content
-            return output
+            # First attempt: Try to parse the entire content as JSON
+            cleaned_content = self._sanitize_json_string(response_content)
+            try:
+                parsed_json = json.loads(cleaned_content)
+                if isinstance(parsed_json, str):
+                    parsed_json = json.loads(parsed_json)
+                return self._normalize_json_output(parsed_json)
+            except json.JSONDecodeError:
+                # Second attempt: Try to extract JSON from <json> tags
+                json_content = self._extract_json_content(response_content)
+                try:
+                    parsed_json = json.loads(json_content)
+                    if isinstance(parsed_json, str):
+                        parsed_json = json.loads(parsed_json)
+                    return self._normalize_json_output(parsed_json)
+                except json.JSONDecodeError:
+                    # Third attempt: Try to parse using regex
+                    return self._parse_with_regex(response_content)
 
-        except (json.JSONDecodeError, ValueError) as e:
-            print(f"Error parsing response: {e}")
+        except Exception as e:
+            print(f"Error parsing response: {str(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)
-            }
+            return self._create_error_dict(str(e))
 
-    def apply_prompt(self, prompt: str, model: str) -> str:
-        try:
-            messages = [
-                {
-                    "role": "system",
-                    "content": """You are a markdown formatting expert. Format your responses with proper spacing and structure following these rules:
-                        
-                        1. Paragraph Spacing:
-                        - Add TWO blank lines between major sections (##)
-                        - Add ONE blank line between subsections (###)
-                        - Add ONE blank line between paragraphs within sections
-                        - Add ONE blank line before and after lists
-                        - Add ONE blank line before and after code blocks
-                        - Add ONE blank line before and after blockquotes
-                        
-                        2. Section Formatting:
-                        # Title
-                        
-                        ## Major Section
-                        
-                        [blank line]
-                        Content paragraph 1
-                        [blank line]
-                        Content paragraph 2
-                        [blank line]"""
-                },
-                {
-                    "role": "user",
-                    "content": prompt
-                }
-            ]
-    
-            response = self.client.chat_completion(
-                model=model,
-                messages=messages,
-                max_tokens=3000,
-                temperature=0.8,
-                stream=True
-            )
-            
-            full_response = ""
-            for chunk in response:
-                if chunk.choices[0].delta.content is not None:
-                    full_response += chunk.choices[0].delta.content
-                    
-            return full_response.replace('\n\n', '\n').strip()
-                
-        except Exception as e:
-            return f"Error: {str(e)}"
+    def _normalize_json_output(self, json_output: dict) -> dict:
+        """Normalize JSON output to expected format."""
+        return {
+            "initial_prompt_evaluation": json_output.get("initial_prompt_evaluation", ""),
+            "refined_prompt": json_output.get("refined_prompt", ""),
+            "explanation_of_refinements": json_output.get("explanation_of_refinements", ""),
+            "response_content": json_output
+        }
+
+    def _parse_with_regex(self, content: str) -> dict:
+        """Parse content using regex patterns."""
+        output = {}
+        for key in ["initial_prompt_evaluation", "refined_prompt", "explanation_of_refinements"]:
+            pattern = rf'"{key}":\s*"(.*?)"(?:,|\}})'
+            match = re.search(pattern, content, re.DOTALL)
+            output[key] = match.group(1) if match else ""
+        
+        output["response_content"] = content
+        return output
+
+    def _create_error_dict(self, error_message: str) -> dict:
+        """Create standardized error response dictionary."""
+        return {
+            "initial_prompt_evaluation": f"Error parsing response: {error_message}",
+            "refined_prompt": "",
+            "explanation_of_refinements": "",
+            "response_content": {"error": error_message}
+        }
+
+    # Rest of your code remains the same...