Update prompt_refiner.py

prompt_refiner.py

@@ -1,88 +1,110 @@
 import json
 import re
-from typing import Optional, Dict, Any, Union
+from typing import Optional, Dict, Any, Tuple
 from pydantic import BaseModel, Field, validator
 from huggingface_hub import InferenceClient
 from huggingface_hub.errors import HfHubHTTPError
-from variables import *
+from variables import meta_prompts, prompt_refiner_model
 
 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: Union[str, list] = Field(..., description="Explanation of the refinements made")
+    explanation_of_refinements: str = 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')
+    @validator('initial_prompt_evaluation', 'refined_prompt', 'explanation_of_refinements')
     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):
+    def __init__(self, api_token: str, meta_prompts: dict):
         self.client = InferenceClient(token=api_token, timeout=120)
         self.meta_prompts = meta_prompts
 
-    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 refine_prompt(self, prompt: str, meta_prompt_choice: str) -> Tuple[str, str, str, dict]:
+        """Refine the given prompt using the selected meta prompt."""
+        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)
+            
+            try:
+                llm_response = LLMResponse(**result)
+                return (
+                    llm_response.initial_prompt_evaluation,
+                    llm_response.refined_prompt,
+                    llm_response.explanation_of_refinements,
+                    llm_response.dict()
+                )
+            except Exception as e:
+                print(f"Error creating LLMResponse: {e}")
+                return self._create_error_response(f"Error validating response: {str(e)}")
 
-    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
+        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 _parse_response(self, response_content: str) -> dict:
+        """Parse the LLM response content."""
         try:
-            # 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 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).strip()
+                # Clean up the JSON string
+                json_str = re.sub(r'\s+', ' ', json_str)
+                json_str = json_str.replace('•', '*')  # Replace bullet points
+                
                 try:
-                    parsed_json = json.loads(json_content)
+                    parsed_json = json.loads(json_str)
                     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)
+                    
+                    return {
+                        "initial_prompt_evaluation": parsed_json.get("initial_prompt_evaluation", ""),
+                        "refined_prompt": parsed_json.get("refined_prompt", ""),
+                        "explanation_of_refinements": parsed_json.get("explanation_of_refinements", ""),
+                        "response_content": parsed_json
+                    }
+                except json.JSONDecodeError as e:
+                    print(f"JSON parsing error: {e}")
+                    return self._create_error_dict(str(e))
+            
+            # Fallback to regex parsing if JSON extraction fails
+            return self._parse_with_regex(response_content)
 
         except Exception as e:
-            print(f"Error parsing response: {str(e)}")
+            print(f"Error parsing response: {e}")
             print(f"Raw content: {response_content}")
             return self._create_error_dict(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."""
+        """Parse content using regex patterns when JSON parsing fails."""
         output = {}
         for key in ["initial_prompt_evaluation", "refined_prompt", "explanation_of_refinements"]:
             pattern = rf'"{key}":\s*"(.*?)"(?:,|\}})'
@@ -93,7 +115,7 @@ class PromptRefiner:
         return output
 
     def _create_error_dict(self, error_message: str) -> dict:
-        """Create standardized error response dictionary."""
+        """Create a standardized error response dictionary."""
         return {
             "initial_prompt_evaluation": f"Error parsing response: {error_message}",
             "refined_prompt": "",
@@ -101,4 +123,11 @@ class PromptRefiner:
             "response_content": {"error": error_message}
         }
 
-    # Rest of your code remains the same...
+    def _create_error_response(self, error_message: str) -> Tuple[str, str, str, dict]:
+        """Create a standardized error response tuple."""
+        return (
+            f"Error: {error_message}",
+            "The selected model is currently unavailable.",
+            "An error occurred during processing.",
+            {"error": error_message}
+        )