Update prompt_refiner.py

prompt_refiner.py

@@ -1,11 +1,25 @@
 import json
 import re
+from typing import Optional, Dict, Any
+from pydantic import BaseModel, Field, validator
 from huggingface_hub import InferenceClient
 from huggingface_hub.errors import HfHubHTTPError
-from variables import meta_prompts, prompt_refiner_model
+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")
+    response_content: Optional[Dict[str, Any]] = Field(None, description="Raw response content")
+
+    @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
 
 class PromptRefiner:
-    def __init__(self, api_token: str,meta_prompts):
+    def __init__(self, api_token: str, meta_prompts):
         self.client = InferenceClient(token=api_token, timeout=120)
         self.meta_prompts = meta_prompts
         
@@ -13,11 +27,9 @@ class PromptRefiner:
         try:
             selected_meta_prompt = self.meta_prompts.get(
                 meta_prompt_choice, 
-                self.meta_prompts["star"]  # Default to "star" if choice not found
+                self.meta_prompts["star"]
             )
-            #print('_'*100)
-            #print(selected_meta_prompt)
-            #print('°°'*100)
+            
             messages = [
                 {
                     "role": "system", 
@@ -28,8 +40,7 @@ class PromptRefiner:
                     "content": selected_meta_prompt.replace("[Insert initial prompt here]", prompt)
                 }
             ]
-            #print(messages)
-            #print('°°'*100)
+            
             response = self.client.chat_completion(
                 model=prompt_refiner_model,
                 messages=messages,
@@ -38,33 +49,40 @@ class PromptRefiner:
             )
             
             response_content = response.choices[0].message.content.strip()
-            
             result = self._parse_response(response_content)
             
+            # Create and validate LLMResponse
+            llm_response = LLMResponse(**result)
+            
             return (
-                result.get('initial_prompt_evaluation', ''),
-                result.get('refined_prompt', ''),
-                result.get('explanation_of_refinements', ''),
-                result
+                llm_response.initial_prompt_evaluation,
+                llm_response.refined_prompt,
+                llm_response.explanation_of_refinements,
+                llm_response.dict()
             )
 
         except HfHubHTTPError as e:
-            return (
-                "Error: Model timeout. Please try again later.",
-                "The selected model is currently experiencing high traffic.",
-                "The selected model is currently experiencing high traffic.",
-                {}
-            )
+            return self._create_error_response("Model timeout. Please try again later.")
         except Exception as e:
-            return (
-                f"Error: {str(e)}",
-                "",
-                "An unexpected error occurred.",
-                {}
-            )
+            return self._create_error_response(f"Unexpected error: {str(e)}")
+
+    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 _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)
@@ -74,19 +92,22 @@ class PromptRefiner:
                 
                 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()
+                
+                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
                 }
-                output['response_content'] = json_output
-                return 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).replace('\\n', '\n').replace('\\"', '"') if match else ""
-            output['response_content'] = response_content
+                output[key] = match.group(1) if match else ""
+            
+            output["response_content"] = response_content
             return output
 
         except (json.JSONDecodeError, ValueError) as e:
@@ -96,7 +117,7 @@ class PromptRefiner:
                 "initial_prompt_evaluation": "Error parsing response",
                 "refined_prompt": "",
                 "explanation_of_refinements": str(e),
-                'response_content': str(e)
+                "response_content": str(e)
             }
 
     def apply_prompt(self, prompt: str, model: str) -> str:
@@ -140,7 +161,6 @@ class PromptRefiner:
             )
             
             full_response = ""
-            
             for chunk in response:
                 if chunk.choices[0].delta.content is not None:
                     full_response += chunk.choices[0].delta.content