File size: 9,305 Bytes
7ed59a1
 
49175b2
51d80c4
7ed59a1
 
49175b2
51d80c4
 
 
 
49175b2
51d80c4
 
49175b2
51d80c4
 
 
 
7ed59a1
49175b2
 
 
 
 
 
 
 
 
7ed59a1
169974c
7ed59a1
 
 
49175b2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
169974c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
51d80c4
169974c
 
 
 
7ed59a1
169974c
 
 
 
 
 
 
49175b2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
import json
import re
from typing import Optional, Dict, Any, Union, List, Tuple
from pydantic import BaseModel, Field, validator
from huggingface_hub import InferenceClient
from huggingface_hub.errors import HfHubHTTPError
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: Union[str, List[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')
    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('\\"', '"').replace('•', '-') 
                   for item in v if isinstance(item, str)]
        return v

class PromptRefiner:
    def __init__(self, api_token: str, meta_prompts: dict):
        self.client = InferenceClient(token=api_token, timeout=120)
        self.meta_prompts = meta_prompts

    def _clean_json_string(self, content: str) -> str:
        """Clean and prepare JSON string for parsing."""
        content = content.replace('•', '-')  # Replace bullet points
        content = re.sub(r'\s+', ' ', content)  # Normalize whitespace
        content = content.replace('\\"', '"')  # Fix escaped quotes
        return content.strip()

    def _parse_response(self, response_content: str) -> dict:
        """Parse the LLM response with enhanced error handling."""
        try:
            # Extract content between <json> tags
            json_match = re.search(r'<json>\s*(.*?)\s*</json>', response_content, re.DOTALL)
            if json_match:
                json_str = self._clean_json_string(json_match.group(1))
                try:
                    # Try parsing the cleaned JSON
                    parsed_json = json.loads(json_str)
                    if isinstance(parsed_json, str):
                        parsed_json = json.loads(parsed_json)
                    
                    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:
                    # If JSON parsing fails, try regex parsing
                    return self._parse_with_regex(json_str)
            
            # If no JSON tags found, try regex parsing
            return self._parse_with_regex(response_content)

        except Exception as e:
            print(f"Error parsing response: {str(e)}")
            print(f"Raw content: {response_content}")
            return self._create_error_dict(str(e))

    def _parse_with_regex(self, content: str) -> dict:
        """Parse content using regex when JSON parsing fails."""
        output = {}
        
        # Handle explanation_of_refinements list format
        refinements_match = re.search(r'"explanation_of_refinements":\s*\[(.*?)\]', content, re.DOTALL)
        if refinements_match:
            refinements_str = refinements_match.group(1)
            refinements = [
                item.strip().strip('"').strip("'").replace('•', '-')
                for item in re.findall(r'[•"]([^"•]+)[•"]', refinements_str)
            ]
            output["explanation_of_refinements"] = refinements
        else:
            # Try single string format
            pattern = r'"explanation_of_refinements":\s*"(.*?)"(?:,|\})'
            match = re.search(pattern, content, re.DOTALL)
            output["explanation_of_refinements"] = match.group(1).strip() if match else ""

        # Extract other fields
        for key in ["initial_prompt_evaluation", "refined_prompt"]:
            pattern = rf'"{key}":\s*"(.*?)"(?:,|\}})'
            match = re.search(pattern, content, re.DOTALL)
            output[key] = match.group(1).strip() if match else ""
        
        output["response_content"] = content
        return output

    def _create_error_dict(self, error_message: str) -> dict:
        """Create a standardized error response dictionary."""
        return {
            "initial_prompt_evaluation": f"Error parsing response: {error_message}",
            "refined_prompt": "",
            "explanation_of_refinements": "",
            "response_content": {"error": error_message}
        }

    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)}")

        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 _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}
        )

    def apply_prompt(self, prompt: str, model: str) -> str:
        """Apply formatting to the prompt using the specified model."""
        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)}"