app.py

import gradio as gr
import openai
import os
from typing import Optional

# System prompt for reformatting rough drafts into polished system prompts
SYSTEM_PROMPT = """You are a system prompt specialist that converts rough drafts into polished, clear specifications for AI code generation agents and subagents.

Your task:
1. Take the user's rough draft describing an AI code generation agent's desired functionalities
2. Transform it into a well-organized, clear system prompt written in second person ("you")
3. Ensure all desired functionalities from the original draft are preserved and clearly reflected
4. Include context that the agent will operate as part of an AI agent crew collaborating within a repository/codebase
5. Inform the agent about its multi-agent environment without prescribing specific collaborative behaviors

Requirements:
- Write the system prompt directing the agent in second person
- Organize the content logically with clear structure
- Maintain all functional requirements from the original draft
- Include information about the agent's multi-agent context in the system prompt
- Let the agent determine how to handle collaboration based on its role
- Return only the polished system prompt with no additional commentary

Output format: Return the improved system prompt as plain text with no markdown formatting or code fences."""

# Alternative system prompt for Claude Code with frontmatter
CLAUDE_CODE_SYSTEM_PROMPT = """You are a system prompt specialist that converts rough drafts into polished, clear specifications for AI code generation agents designed for use in Claude Code.

Your task:
1. Take the user's rough draft describing an AI code generation agent's desired functionalities
2. Create appropriate frontmatter following the Claude Code format with name, description, examples, tools, model, and color
3. Transform the draft into a well-organized system prompt written in second person ("you")
4. Include context that the agent will operate as part of an AI agent crew collaborating within a repository/codebase
5. Ensure the frontmatter description and examples clearly show when to use this agent

Requirements:
- Start with properly formatted frontmatter (name, description with examples, tools, model, color)
- Write the system prompt directing the agent in second person
- Organize the content logically with clear structure
- Maintain all functional requirements from the original draft
- Include information about the agent's multi-agent context
- Provide realistic usage examples in the frontmatter description
- Return only the complete prompt with frontmatter as plain text

Output format: Return the complete prompt with frontmatter and system prompt as plain text with no markdown formatting or code fences."""

def reformat_prompt(api_key: str, rough_draft: str, prompt_type: str = "standard", api_key_state: str = "") -> tuple[str, str]:
    """
    Reformat a rough draft into a polished system prompt using OpenAI API.
    
    Args:
        api_key: OpenAI API key
        rough_draft: The rough draft text to be reformatted
        prompt_type: Type of prompt to generate ("standard" or "claude_code")
        api_key_state: State variable for API key persistence
        
    Returns:
        Tuple of (reformatted system prompt or error message, updated api_key_state)
    """
    # Use provided API key or fall back to state
    effective_api_key = api_key.strip() or api_key_state.strip()
    
    if not effective_api_key:
        return "Error: Please provide your OpenAI API key.", api_key_state
    
    if not rough_draft.strip():
        return "Error: Please provide a rough draft to reformat.", effective_api_key
    
    # Select the appropriate system prompt based on type
    selected_system_prompt = CLAUDE_CODE_SYSTEM_PROMPT if prompt_type == "claude_code" else SYSTEM_PROMPT
    
    try:
        # Initialize OpenAI client with the effective API key
        client = openai.OpenAI(api_key=effective_api_key)
        
        # Make API call to reformat the prompt
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": selected_system_prompt},
                {"role": "user", "content": rough_draft}
            ],
            temperature=0.1,
            max_tokens=2000
        )
        
        reformatted_text = response.choices[0].message.content
        return reformatted_text, effective_api_key
        
    except openai.AuthenticationError:
        return "Error: Invalid API key. Please check your OpenAI API key.", api_key_state
    except openai.RateLimitError:
        return "Error: Rate limit exceeded. Please try again later.", effective_api_key
    except openai.APIError as e:
        return f"Error: OpenAI API error - {str(e)}", effective_api_key
    except Exception as e:
        return f"Error: {str(e)}", effective_api_key

def create_interface():
    """Create and configure the Gradio interface."""
    
    with gr.Blocks(
        title="Code Gen Agent Prompt Assistant",
        theme=gr.themes.Soft(),
        css="""
        .container {
            max-width: 1200px;
            margin: auto;
        }
        .header {
            text-align: center;
            margin-bottom: 2rem;
        }
        .api-key-box {
            background-color: #f8f9fa;
            padding: 1rem;
            border-radius: 8px;
            margin-bottom: 1rem;
        }
        .system-prompt-display {
            background-color: #f8f9fa;
            padding: 1rem;
            border-radius: 8px;
            font-family: monospace;
            white-space: pre-wrap;
            border: 1px solid #dee2e6;
        }
        .persistence-info {
            background-color: #e8f5e8;
            padding: 0.75rem;
            border-radius: 6px;
            margin-top: 0.5rem;
            font-size: 0.9em;
            color: #2d5a2d;
        }
        """
    ) as interface:
        
        gr.HTML("""
        <div class="header">
            <h1>Code Gen Agent Prompt Assistant</h1>
            <p>Transform rough drafts into polished system prompts for AI code generation agents operating within multi-agent crews.</p>
        </div>
        """)
        
        with gr.Tabs():
            with gr.TabItem("Prompt Reformatter"):
                with gr.Row():
                    with gr.Column():
                        # Create a state variable for API key persistence
                        api_key_state = gr.State("")
                        
                        api_key_input = gr.Textbox(
                            label="OpenAI API Key",
                            placeholder="sk-...",
                            type="password",
                            info="Your API key will be remembered during this session"
                        )
                        
                        prompt_type_radio = gr.Radio(
                            choices=[
                                ("Standard Multi-Agent Prompt", "standard"),
                                ("Claude Code with Frontmatter", "claude_code")
                            ],
                            value="standard",
                            label="Prompt Type",
                            info="Choose the format for your system prompt"
                        )
                        
                        rough_draft_input = gr.Textbox(
                            label="Rough Draft",
                            placeholder="Enter your rough draft describing the desired functionalities of the AI code generation agent...",
                            lines=10,
                            max_lines=20
                        )
                        
                        with gr.Row():
                            clear_btn = gr.Button("Clear", variant="secondary")
                            submit_btn = gr.Button("Reformat Prompt", variant="primary", scale=2)
                    
                    with gr.Column():
                        output = gr.Textbox(
                            label="Reformatted System Prompt",
                            lines=15,
                            max_lines=25,
                            show_copy_button=True,
                            interactive=False
                        )
            
            with gr.TabItem("System Prompts"):
                gr.HTML("""
                <h3>System Prompts Used in Workflow</h3>
                <p>These are the system prompts that guide the AI in reformatting your rough drafts:</p>
                """)
                
                with gr.Tabs():
                    with gr.TabItem("Standard Multi-Agent"):
                        gr.Textbox(
                            value=SYSTEM_PROMPT,
                            label="Standard Multi-Agent System Prompt",
                            lines=15,
                            max_lines=25,
                            show_copy_button=True,
                            interactive=False,
                            elem_classes=["system-prompt-display"]
                        )
                        
                        gr.HTML("""
                        <div style="margin-top: 1rem; padding: 1rem; background-color: #e3f2fd; border-radius: 8px;">
                            <h4>About This System Prompt</h4>
                            <p>This system prompt creates standard prompts that:</p>
                            <ul>
                                <li>Inform agents about their multi-agent environment context</li>
                                <li>Convert rough notes into polished system prompts for code generation agents</li>
                                <li>Focus on agents within multi-agent crews working in codebases</li>
                                <li>Let agents determine their own collaborative behaviors</li>
                            </ul>
                        </div>
                        """)
                    
                    with gr.TabItem("Claude Code with Frontmatter"):
                        gr.Textbox(
                            value=CLAUDE_CODE_SYSTEM_PROMPT,
                            label="Claude Code System Prompt",
                            lines=15,
                            max_lines=25,
                            show_copy_button=True,
                            interactive=False,
                            elem_classes=["system-prompt-display"]
                        )
                        
                        gr.HTML("""
                        <div style="margin-top: 1rem; padding: 1rem; background-color: #fff3e0; border-radius: 8px;">
                            <h4>About This System Prompt</h4>
                            <p>This system prompt creates Claude Code compatible prompts that:</p>
                            <ul>
                                <li>Include proper frontmatter with name, description, examples, tools, model, and color</li>
                                <li>Provide realistic usage examples in the description</li>
                                <li>Create agents designed for Claude Code's multi-agent environment</li>
                                <li>Format everything according to Claude Code specifications</li>
                            </ul>
                        </div>
                        """)
        
        # Event handlers
        submit_btn.click(
            fn=reformat_prompt,
            inputs=[api_key_input, rough_draft_input, prompt_type_radio, api_key_state],
            outputs=[output, api_key_state],
            show_progress=True
        )
        
        clear_btn.click(
            fn=lambda: ("", "standard", "", "", ""),
            outputs=[api_key_input, prompt_type_radio, rough_draft_input, output, api_key_state]
        )
        
    
    return interface

# Create and launch the interface
if __name__ == "__main__":
    app = create_interface()
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True
    )