app.py

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
from huggingface_hub.errors import HfHubHTTPError

from custom_css import custom_css
from variables import *


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

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):
        self.client = InferenceClient(token=api_token, timeout=300)
        self.meta_prompts = {
            "morphosis": original_meta_prompt,
            "verse": new_meta_prompt,
            "physics": metaprompt1,
            "bolism": loic_metaprompt,
            "done": metadone,
            "star": echo_prompt_refiner,
            "math": math_meta_prompt,
            "arpe": autoregressive_metaprompt
        }

    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, 
                advanced_meta_prompt
            )
            
            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_input.text)
                }
            ]

            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()
            
            # Parse the response
            result = self._parse_response(response_content)
            
            return (
                result.get('initial_prompt_evaluation', ''),
                result.get('refined_prompt', ''),
                result.get('explanation_of_refinements', ''),
                result
            )

        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.",
                {}
            )
        except Exception as e:
            return (
                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 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]
                    
                    ### Subsection
                    [blank line]
                    Content
                    [blank line]
                    
                    3. List Formatting:
                    [blank line]
                    - List item 1
                    - List item 2
                    - List item 3
                    [blank line]
                    
                    4. JSON Output Structure:
                    {
                        "section_name": "
                            Content paragraph 1
                            
                            Content paragraph 2
                            
                            - List item 1
                            - List item 2
                        "
                    }
            
                    Transform content while maintaining clear visual separation between elements."""
                },
                {
                    "role": "user",
                    "content": prompt
                }
            ]
    
            # Use streaming for the response
            response_stream = self.client.text_generation(
                model=model,
                messages=messages,
                max_tokens=3000,
                temperature=0.8,
                stream=True
            )
            
            # Initialize an empty string to store the complete response
            full_response = ""
            
            # Process the stream
            for response_chunk in response_stream:
                if hasattr(response_chunk, 'token'):
                    chunk_text = response_chunk.token.text
                    full_response += chunk_text
                    yield full_response.replace('\n\n', '\n').strip()
                
        except Exception as e:
            yield f"Error: {str(e)}"

class GradioInterface:
    def __init__(self, prompt_refiner: PromptRefiner,custom_css):
        self.prompt_refiner = prompt_refiner
        custom_css = custom_css
        with gr.Blocks(css=custom_css, theme=gr.themes.Default()) as self.interface:
            with gr.Column(elem_classes=["container", "title-container"]):
                gr.Markdown("# PROMPT++")
                gr.Markdown("### Automating Prompt Engineering by Refining your Prompts")
                gr.Markdown("Learn how to generate an improved version of your prompts.")
        
            with gr.Column(elem_classes=["container", "input-container"]):
                prompt_text = gr.Textbox(
                    label="Type your prompt (or let it empty to see metaprompt)",
                   # elem_classes="no-background",
                    #elem_classes="container2",
                    lines=5
                )
                meta_prompt_choice = gr.Radio(
                    ["star","done","physics","morphosis", "verse", "phor","bolism","math","arpe"],
                    label="Choose Meta Prompt",
                    value="star",
                    elem_classes=["no-background", "radio-group"]
                   # elem_classes=[ "radio-group"]
                )
                refine_button = gr.Button("Refine Prompt")  
                
                # Option 1: Put Examples here (before Meta Prompt explanation)
                with gr.Row(elem_classes=["container2"]):
                    with gr.Accordion("Examples", open=False):
                        gr.Examples(
                            examples=[
                                ["Write a story on the end of prompt engineering replaced by an Ai specialized in refining prompts.", "star"],
                                ["Tell me about that guy who invented the light bulb", "physics"],
                                ["Explain the universe.", "star"],
                                ["What's the population of New York City and how tall is the Empire State Building and who was the first mayor?", "morphosis"],
                                ["List American presidents.", "verse"],                        
                                ["Explain why the experiment failed.", "morphosis"],
                                ["Is nuclear energy good?", "verse"],
                                ["How does a computer work?", "phor"],
                                ["How to make money fast?", "done"],
                                ["how can you prove IT0's lemma in stochastic calculus ?", "arpe"],                    
                            ],
                            inputs=[prompt_text, meta_prompt_choice]
                        )
                       
                    with gr.Accordion("Meta Prompt explanation", open=False):
                        gr.Markdown(explanation_markdown)
                

                
                # Option 2: Or put Examples here (after the button)
                # with gr.Accordion("Examples", open=False):
                #     gr.Examples(...)
                
            with gr.Column(elem_classes=["container", "analysis-container"]):
                gr.Markdown(' ')
                gr.Markdown("### Initial prompt analysis")
                analysis_evaluation = gr.Markdown()
                gr.Markdown("### Refined Prompt")
                refined_prompt = gr.Textbox(
                                    label="Refined Prompt",
                                    interactive=True,
                                    show_label=True,  # Must be True for copy button to show
                                    show_copy_button=True,  # Adds the copy button
                             #       elem_classes="no-background"
                                )
                gr.Markdown("### Explanation of Refinements")
                explanation_of_refinements = gr.Markdown()
            

            with gr.Column(elem_classes=["container", "model-container"]):
              #  gr.Markdown("## See MetaPrompt Impact")            
                with gr.Row():
                    apply_model = gr.Dropdown(models,
                                            value="meta-llama/Llama-3.1-8B-Instruct",
                                            label="Choose the Model",
                                            container=False,  # This removes the container around the dropdown
                                            scale=1,         # Controls the width relative to other components
                                            min_width=300    # Sets minimum width in pixels
                                         #   elem_classes="no-background"
                                        )
                    apply_button = gr.Button("Apply MetaPrompt")

          #  with gr.Column(elem_classes=["container", "results-container"]):
                gr.Markdown("### Prompts on choosen model")
                with gr.Tabs():
                    with gr.TabItem("Original Prompt Output"):
                        original_output = gr.Markdown()
                    with gr.TabItem("Refined Prompt Output"):
                        refined_output = gr.Markdown()
            with gr.Accordion("Full Response JSON", open=False, visible=True):
                full_response_json = gr.JSON()
                
            refine_button.click(
                fn=self.refine_prompt,
                inputs=[prompt_text, meta_prompt_choice],
                outputs=[analysis_evaluation, refined_prompt, explanation_of_refinements, full_response_json]
            )

                # In the __init__ method, modify the click event:
            apply_button.click(
                fn=self.apply_prompts,
                inputs=[prompt_text, refined_prompt, apply_model],
                outputs=[original_output, refined_output],
                queue=True  # Enable queuing for streaming
            )

    def refine_prompt(self, prompt: str, meta_prompt_choice: str) -> tuple:
        input_data = PromptInput(text=prompt, meta_prompt_choice=meta_prompt_choice)
        # Since result is a tuple with 4 elements based on the return value of prompt_refiner.refine_prompt
        initial_prompt_evaluation, refined_prompt, explanation_refinements, full_response = self.prompt_refiner.refine_prompt(input_data)
        
        analysis_evaluation = f"\n\n{initial_prompt_evaluation}"
        return (
            analysis_evaluation,
            refined_prompt,
            explanation_refinements,
            full_response
        )

    def apply_prompts(self, original_prompt: str, refined_prompt: str, model: str):
        def stream_generator():
            for response in self.prompt_refiner.apply_prompt(original_prompt, model):
                yield response, ""
            
            for response in self.prompt_refiner.apply_prompt(refined_prompt, model):
                current_output = yield "", response
        
        return stream_generator()
    


    def launch(self, share=False):
        self.interface.launch(share=share)


#explanation_markdown = "".join([f"- **{key}**: {value}\n" for key, value in metaprompt_explanations.items()])
'''   
meta_info=""
api_token = os.getenv('HF_API_TOKEN')
if not api_token:
    raise ValueError("HF_API_TOKEN not found in environment variables")

metadone = os.getenv('metadone')
prompt_refiner_model = os.getenv('prompt_refiner_model')
echo_prompt_refiner = os.getenv('echo_prompt_refiner')
metaprompt1 = os.getenv('metaprompt1')   
loic_metaprompt = os.getenv('loic_metaprompt')    
openai_metaprompt = os.getenv('openai_metaprompt')
original_meta_prompt = os.getenv('original_meta_prompt')    
new_meta_prompt = os.getenv('new_meta_prompt')   
advanced_meta_prompt = os.getenv('advanced_meta_prompt')
math_meta_prompt = os.getenv('metamath')
autoregressive_metaprompt = os.getenv('autoregressive_metaprompt')
'''    

if __name__ == '__main__':
    prompt_refiner = PromptRefiner(api_token)
    gradio_interface = GradioInterface(prompt_refiner,custom_css)
    gradio_interface.launch(share=True)