app.py

import os
import json
import re
import gradio as gr
from groq import Groq
import logging
from pydantic import BaseModel, Field
from typing import Optional, Literal

from custom_css import custom_css
from variables import *

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Add system prompt generation meta prompt
SYSTEM_META_PROMPT = """
Given a task description or existing prompt, produce a detailed system prompt to guide a language model in completing the task effectively.
# Guidelines
- Understand the Task: Grasp the main objective, goals, requirements, constraints, and expected output.
- Minimal Changes: If an existing prompt is provided, improve it only if it's simple. For complex prompts, enhance clarity and add missing elements without altering the original structure.
- Reasoning Before Conclusions**: Encourage reasoning steps before any conclusions are reached. ATTENTION! If the user provides examples where the reasoning happens afterward, REVERSE the order! NEVER START EXAMPLES WITH CONCLUSIONS!
    - Reasoning Order: Call out reasoning portions of the prompt and conclusion parts (specific fields by name). For each, determine the ORDER in which this is done, and whether it needs to be reversed.
    - Conclusion, classifications, or results should ALWAYS appear last.
- Examples: Include high-quality examples if helpful, using placeholders [in brackets] for complex elements.
   - What kinds of examples may need to be included, how many, and whether they are complex enough to benefit from placeholders.
- Clarity and Conciseness: Use clear, specific language. Avoid unnecessary instructions or bland statements.
- Formatting: Use markdown features for readability. DO NOT USE ``` CODE BLOCKS UNLESS SPECIFICALLY REQUESTED.
- Preserve User Content: If the input task or prompt includes extensive guidelines or examples, preserve them entirely, or as closely as possible. If they are vague, consider breaking down into sub-steps. Keep any details, guidelines, examples, variables, or placeholders provided by the user.
- Constants: DO include constants in the prompt, as they are not susceptible to prompt injection. Such as guides, rubrics, and examples.
- Output Format: Explicitly the most appropriate output format, in detail. This should include length and syntax (e.g. short sentence, paragraph, JSON, etc.)
    - For tasks outputting well-defined or structured data (classification, JSON, etc.) bias toward outputting a JSON.
    - JSON should never be wrapped in code blocks (```) unless explicitly requested.
The final prompt you output should adhere to the following structure below. Do not include any additional commentary, only output the completed system prompt. SPECIFICALLY, do not include any additional messages at the start or end of the prompt. (e.g. no "---")
[Concise instruction describing the task - this should be the first line in the prompt, no section header]
[Additional details as needed.]
[Optional sections with headings or bullet points for detailed steps.]
# Steps [optional]
[optional: a detailed breakdown of the steps necessary to accomplish the task]
# Output Format
[Specifically call out how the output should be formatted, be it response length, structure e.g. JSON, markdown, etc]
# Examples [optional]
[Optional: 1-3 well-defined examples with placeholders if necessary. Clearly mark where examples start and end, and what the input and output are. User placeholders as necessary.]
[If the examples are shorter than what a realistic example is expected to be, make a reference with () explaining how real examples should be longer / shorter / different. AND USE PLACEHOLDERS! ]
# Notes [optional]
[optional: edge cases, details, and an area to call or repeat out specific important considerations]
""".strip()

class PromptInput(BaseModel):
    text: str = Field(..., description="The initial prompt text")
    meta_prompt_choice: Literal["superstar","star","done","physics","morphosis", "verse", "phor","bolism","math", "math_meta", "system"] = 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
    system_prompt: Optional[str] = None  # New field for system prompt

class PromptRefiner:
    def __init__(self, api_token: str):
        self.client = Groq(api_key=api_key)

    def generate_system_prompt(self, task_or_prompt: str, model: str = "llama-3.3-70b-versatile") -> str:
        """Generate a system prompt for the given task or prompt."""
        messages = [
            {
                "role": "system",
                "content": SYSTEM_META_PROMPT,
            },
            {
                "role": "user",
                "content": f"Task, Goal, or Current Prompt:\n{task_or_prompt}",
            },
        ]
        
        response = self.client.chat.completions.create(
            model=model,
            messages=messages,
            max_tokens=8000,
            temperature=0.7,
            stream=True
        )
        
        prompt = ''
        for chunk in response:
            if chunk.choices[0].delta.content is not None:
                prompt += chunk.choices[0].delta.content
                
        logger.info("Generated system prompt: %s", prompt)
        return prompt
    
    def refine_prompt(self, prompt_input: PromptInput) -> RefinementOutput:
        # Handle system prompt generation separately
        if prompt_input.meta_prompt_choice == "system":
            system_prompt = self.generate_system_prompt(prompt_input.text)
            return RefinementOutput(
                refined_prompt=system_prompt,
                explanation_of_refinements="Generated system prompt based on the task/prompt.",
                system_prompt=system_prompt
            )
            
        # Existing meta prompt selection logic      
        if prompt_input.meta_prompt_choice == "morphosis":
            selected_meta_prompt = original_meta_prompt
        elif prompt_input.meta_prompt_choice == "verse":
            selected_meta_prompt = new_meta_prompt
        elif prompt_input.meta_prompt_choice == "physics":
            selected_meta_prompt = metaprompt1
        elif prompt_input.meta_prompt_choice == "bolism":
            selected_meta_prompt = loic_metaprompt
        elif prompt_input.meta_prompt_choice == "done":
            selected_meta_prompt = metadone
        elif prompt_input.meta_prompt_choice == "star":
            selected_meta_prompt = echo_prompt_refiner
        elif prompt_input.meta_prompt_choice == "superstar":
            selected_meta_prompt = advanced_echo_prompt_refiner
        elif prompt_input.meta_prompt_choice == "math":
            selected_meta_prompt = math_meta_prompt
        elif prompt_input.meta_prompt_choice == "math_meta":
            selected_meta_prompt = math_meta
        else:
            selected_meta_prompt = advanced_meta_prompt
    
        messages = [
            {"role": "system", "content": 'You are an expert at refining and extending prompts. Given a basic prompt, provide a more detailed.'},
            {"role": "user", "content": selected_meta_prompt.replace("[Insert initial prompt here]", prompt_input.text)}
        ]
        response = self.client.chat.completions.create(
            model="llama-3.3-70b-versatile",
            messages=messages,
            max_tokens=8192,
            temperature=0.7
        )
        response_content = response.choices[0].message.content.strip()
        try:
            # Extract JSON from between <json> tags
            json_match = re.search(r'<json>\s*(.*?)\s*</json>', response_content, re.DOTALL)
            if json_match:
                json_str = json_match.group(1)
                # Remove newlines and escape quotes within the JSON string
                json_str = re.sub(r'\n\s*', ' ', json_str)
                json_str = json_str.replace('"', '\\"')
                # Wrap the entire string in quotes and parse it
                json_output = json.loads(f'"{json_str}"')
                # Ensure json_output is a dictionary
                if isinstance(json_output, str):
                    json_output = json.loads(json_output)
                # Unescape the parsed JSON
                for key, value in json_output.items():
                    if isinstance(value, str):
                        json_output[key] = value.replace('\\"', '"')
                return RefinementOutput(**json_output, raw_content=response_content)
            else:
                raise ValueError("No JSON found in the response")
        except (json.JSONDecodeError, ValueError) as e:
            print(f"Error parsing JSON: {e}")
            print(f"Raw content: {response_content}")
            # If JSON parsing fails, attempt to extract the content manually
            output = {}
            for key in ["initial_prompt_evaluation", "refined_prompt", "explanation_of_refinements"]:
                pattern = rf'"{key}":\s*"(.*?)"(?:,|\}})'
                match = re.search(pattern, response_content, re.DOTALL)
                if match:
                    output[key] = match.group(1).replace('\\n', '\n').replace('\\"', '"')
                else:
                    output[key] = "" # Set empty string if content not found
            return RefinementOutput(**output, raw_content=response_content)
                
    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.completions.create(
                model=model,
                messages=messages,
                max_tokens=8000, # Increased token limit
                temperature=0.8,
                stream=True  # Enable streaming in the API call
            )
            # Initialize an empty string to accumulate the response
            full_response = ""
            
            # Process the streaming response
            for chunk in response:
                if chunk.choices[0].delta.content is not None:
                    full_response += chunk.choices[0].delta.content
                    
            # Return the complete response
            return full_response.replace('\n\n', '\n').strip()
           
        except Exception as e:
            return 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(
                    ["superstar", "star", "done", "physics", "morphosis", "verse", "phor","bolism","math","math_meta", "system"],
                    label="Choose Meta Prompt",
                    value="superstar",
                    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.", "superstar"],
                                ["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 ?", "math_meta"], 
                                ["Optimize the prompt that users enter for image generation with Stable Diffusion XL model", "system"], 
                            ],
                            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="llama-3.1-70b-versatile",
                                            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()

            # Add new tab for system prompt output
            with gr.Column(elem_classes=["container", "system-prompt-container"]):
                with gr.Tabs():
                    with gr.TabItem("System Prompt"):
                        system_prompt_output = gr.Textbox(
                            label="Generated System Prompt",
                            interactive=True,
                            show_label=True,
                            show_copy_button=True
                        )

            # Modified click handler to include system prompt output
            refine_button.click(
                fn=self.refine_prompt,
                inputs=[prompt_text, meta_prompt_choice],
                outputs=[analysis_evaluation, refined_prompt, explanation_of_refinements, full_response_json, system_prompt_output]
            )
                
            apply_button.click(
                fn=self.apply_prompts,
                inputs=[prompt_text, refined_prompt, apply_model],
                outputs=[original_output, refined_output]
            )
            gr.HTML(
                "<p style='text-align: center; color:orange;'>⚠ This space is in progress, and we're actively working on it, so you might find some bugs! Please report any issues you have in the Community tab to help us make it better for all.</p>"
            )

    def refine_prompt(self, prompt: str, meta_prompt_choice: str) -> tuple:
        try:
            input_data = PromptInput(text=prompt, meta_prompt_choice=meta_prompt_choice)
            result = self.prompt_refiner.refine_prompt(input_data)

             # Include system prompt in output
            system_prompt = str(result.system_prompt) if result.system_prompt else ""
            
            # Ensure all values are strings or None
            analysis_evaluation = str(result.initial_prompt_evaluation) if result.initial_prompt_evaluation else ""
            refined_prompt = str(result.refined_prompt) if result.refined_prompt else ""
            explanation_refinements = str(result.explanation_of_refinements) if result.explanation_of_refinements else ""
            
            # Create response dictionary
            full_response = {
                "initial_prompt_evaluation": str(result.initial_prompt_evaluation) if result.initial_prompt_evaluation else "",
                "refined_prompt": str(result.refined_prompt) if result.refined_prompt else "",
                "explanation_of_refinements": str(result.explanation_of_refinements) if result.explanation_of_refinements else "",
                "raw_content": str(result.raw_content) if result.raw_content else "",
                "system_prompt": system_prompt
            }
            
            return (
                analysis_evaluation,
                refined_prompt,
                explanation_refinements,
                full_response,
                system_prompt
            )
        except Exception as e:
            error_response = {
                "error": str(e),
                "initial_prompt_evaluation": "",
                "refined_prompt": "",
                "explanation_of_refinements": "",
                "raw_content": "",
                "system_prompt": ""
            }
            return "", "", "", error_response, ""

    def apply_prompts(self, original_prompt: str, refined_prompt: str, model: str):
        original_output = self.prompt_refiner.apply_prompt(original_prompt, model)
        refined_output = self.prompt_refiner.apply_prompt(refined_prompt, model)
        return original_output, refined_output

    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_key = os.getenv('GROQ_API_KEY')
if not api_key:
    raise ValueError("GROQ_API_KEY not found in environment variables")

        
metadone=os.getenv('metadone')
echo_prompt_refiner = os.getenv('echo_prompt_refiner')
advanced_echo_prompt_refiner = os.getenv('advanced_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('math_meta_prompt')
math_meta = os.getenv('math_meta')
'''   


# Main code to run the application
if __name__ == '__main__':
    prompt_refiner = PromptRefiner(api_key)
    gradio_interface = GradioInterface(prompt_refiner,custom_css)
    gradio_interface.launch(share=True)