prompt_optimizer_agent.py

"""
Prompt Optimization Agent
==================================================
Automatically analyzes and optimizes user prompts using a multi-step agent pipeline:
  1. Task Complexity Assessor → scores task complexity (1–5)
  2. Analyzer            → identifies weaknesses in the original prompt
  3. Optimizer           → rewrites the prompt (concise if complexity < 3)
  4. Evaluator           → scores improvements (optional, user-triggered)
  5. Executor            → runs original and/or optimized prompt (optional, user-triggered)
"""

import os
import json
from openai import OpenAI

# ── Configuration ──────────────────────────────────────────────────────────────
MODEL = "gpt-5-mini-2025-08-07"
TEMPERATURE = 1           # Default temperature for gpt-5-mini
client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))


# ── Helpers ────────────────────────────────────────────────────────────────────
def ask_yes_no(question: str) -> bool:
    """Prompt the user for a yes/no answer. Returns True for yes."""
    while True:
        answer = input(f"\n{question} (y/n): ").strip().lower()
        if answer in ("y", "yes"):
            return True
        if answer in ("n", "no"):
            return False
        print("  Please enter 'y' or 'n'.")


def ask_choice(question: str, choices: list[str]) -> str:
    """Prompt the user to pick from a numbered list. Returns the chosen string."""
    print(f"\n{question}")
    for i, choice in enumerate(choices, 1):
        print(f"  [{i}] {choice}")
    while True:
        answer = input("Enter number: ").strip()
        if answer.isdigit() and 1 <= int(answer) <= len(choices):
            return choices[int(answer) - 1]
        print(f"  Please enter a number between 1 and {len(choices)}.")


def section(title: str):
    print("\n" + "=" * 60)
    print(title)
    print("=" * 60)


# ── Agent Step 1: Complexity Assessor ─────────────────────────────────────────
def assess_complexity(user_prompt: str) -> dict:
    """Score the complexity of the task on a scale of 1–5."""
    system = """You are a task complexity assessor. Evaluate the complexity of the task described in the prompt.
Scoring guide:
  1 – Trivial: single fact lookup, simple definition, one-liner answer
  2 – Simple: short explanation, basic conversion, limited steps
  3 – Moderate: multi-step reasoning, structured output, some domain knowledge
  4 – Complex: deep analysis, long-form content, multiple constraints
  5 – Expert: research-level, highly technical, creative synthesis

Return ONLY a JSON object:
{
  "complexity_score": <integer 1-5>,
  "reasoning": "<one sentence justification>"
}"""

    response = client.chat.completions.create(
        model=MODEL,
        temperature=TEMPERATURE,
        response_format={"type": "json_object"},
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": f"Assess the complexity of this task:\n\n{user_prompt}"}
        ]
    )
    return json.loads(response.choices[0].message.content)


# ── Agent Step 2: Analyzer ─────────────────────────────────────────────────────
def analyze_prompt(user_prompt: str) -> dict:
    """Identify weaknesses and issues in the original prompt."""
    system = """You are a prompt engineering expert. Analyze the given prompt to first understand the task, then identify the prompt's weaknesses.
Return ONLY a JSON object with this exact structure:
{
  "issues": ["issue 1", "issue 2", ...],
  "missing_elements": ["element 1", ...],
  "clarity_score": <integer 1-10>,
  "summary": "<one sentence summary of main problems>"
}"""

    response = client.chat.completions.create(
        model=MODEL,
        temperature=TEMPERATURE,
        response_format={"type": "json_object"},
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": f"Analyze this prompt:\n\n{user_prompt}"}
        ]
    )
    return json.loads(response.choices[0].message.content)


# ── Agent Step 3: Optimizer ────────────────────────────────────────────────────
def optimize_prompt(user_prompt: str, analysis: dict, complexity_score: int) -> dict:
    """Rewrite the prompt. If complexity < 3, keep it concise."""
    concise_instruction = ""
    if complexity_score < 3:
        concise_instruction = (
            "\n\nIMPORTANT: This is a simple, low-complexity task. "
            "The optimized prompt MUST be short and concise — avoid over-engineering it. "
            "Users want fast answers; do not add unnecessary structure, role-play, or formatting instructions."
        )

    system = f"""You are a world-class prompt engineer. Given a prompt and its analysis, rewrite it to be significantly better.
Apply these techniques where relevant:
- Clear role/persona assignment
- Explicit output format instructions
- Chain-of-thought framing
- Specific constraints and scope
- Context and background information
- Examples (few-shot) when helpful
{concise_instruction}

Return ONLY a JSON object:
{{
  "optimized_prompt": "<the full rewritten prompt>",
  "techniques_applied": ["technique 1", "technique 2", ...]
}}"""

    analysis_str = json.dumps(analysis, indent=2)
    response = client.chat.completions.create(
        model=MODEL,
        temperature=TEMPERATURE,
        response_format={"type": "json_object"},
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": f"Original prompt:\n{user_prompt}\n\nAnalysis:\n{analysis_str}"}
        ]
    )
    return json.loads(response.choices[0].message.content)


# ── Agent Step 4: Evaluator (optional) ────────────────────────────────────────
def evaluate_prompt(original: str, optimized: str) -> dict:
    """Score the optimized prompt and explain the improvements."""
    system = """You are a prompt quality evaluator. Compare the original and optimized prompts.
Return ONLY a JSON object:
{
  "quality_score": <integer 1-10>,
  "improvement_score": <integer 1-10>,
  "strengths": ["strength 1", "strength 2", ...],
  "explanation": "<2-3 sentence explanation of key improvements>"
}"""

    response = client.chat.completions.create(
        model=MODEL,
        temperature=TEMPERATURE,
        response_format={"type": "json_object"},
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": f"Original:\n{original}\n\nOptimized:\n{optimized}"}
        ]
    )
    return json.loads(response.choices[0].message.content)


# ── Agent Step 5: Executor (optional) ─────────────────────────────────────────
def execute_prompt(prompt: str) -> str:
    """Run a prompt and return the plain-text response."""
    response = client.chat.completions.create(
        model=MODEL,
        temperature=TEMPERATURE,
        messages=[{"role": "user", "content": prompt}]
    )
    return response.choices[0].message.content.strip()


# ── Main Pipeline ──────────────────────────────────────────────────────────────
def run_optimizer(user_prompt: str, verbose: bool = True) -> dict:
    """Run the full prompt optimization pipeline."""

    def log(msg):
        if verbose: print(msg)

    # ── Step 1: Complexity Assessment ──────────────────────────────────────────
    section("🧩 STEP 1: Assessing task complexity...")
    complexity = assess_complexity(user_prompt)
    score = complexity["complexity_score"]
    log(f"  Complexity Score : {score}/5")
    log(f"  Reasoning        : {complexity['reasoning']}")
    if score < 3:
        log("  → Low complexity detected. Optimized prompt will be kept concise.")

    # ── Step 2: Analysis ───────────────────────────────────────────────────────
    section("🔍 STEP 2: Analyzing your prompt...")
    analysis = analyze_prompt(user_prompt)
    log(f"  Clarity Score : {analysis['clarity_score']}/10")
    log(f"  Summary       : {analysis['summary']}")
    log(f"  Issues found  : {len(analysis['issues'])}")
    for issue in analysis["issues"]:
        log(f"    • {issue}")

    # ── Step 3: Optimization ───────────────────────────────────────────────────
    section("✨ STEP 3: Optimizing your prompt...")
    optimization = optimize_prompt(user_prompt, analysis, score)
    log(f"  Techniques applied:")
    for t in optimization["techniques_applied"]:
        log(f"    ✓ {t}")

    section("📝 ORIGINAL PROMPT:")
    log(user_prompt)

    section("🚀 OPTIMIZED PROMPT:")
    log(optimization["optimized_prompt"])

    result = {
        "original_prompt": user_prompt,
        "complexity": complexity,
        "analysis": analysis,
        "optimized_prompt": optimization["optimized_prompt"],
        "techniques_applied": optimization["techniques_applied"],
        "evaluation": None,
        "outputs": {}
    }

    # ── Step 4: Optional Evaluation ────────────────────────────────────────────
    if ask_yes_no("📊 Would you like to evaluate the quality of the optimized prompt?"):
        section("📊 STEP 4: Evaluating improvements...")
        evaluation = evaluate_prompt(user_prompt, optimization["optimized_prompt"])
        log(f"  Optimized Quality Score : {evaluation['quality_score']}/10")
        log(f"  Improvement Score       : {evaluation['improvement_score']}/10")
        log(f"  Strengths:")
        for s in evaluation["strengths"]:
            log(f"    ★ {s}")
        log(f"  Explanation: {evaluation['explanation']}")
        result["evaluation"] = evaluation
    else:
        log("\n  ⏭  Skipping evaluation.")

    # ── Step 5: Optional Execution ─────────────────────────────────────────────
    if ask_yes_no("⚡ Would you like to run the prompt(s) and see the outputs?"):
        exec_choice = ask_choice(
            "Which output(s) would you like to see?",
            [
                "Optimized prompt only",
                "Both original and optimized (side-by-side comparison)",
            ]
        )

        if exec_choice == "Both original and optimized (side-by-side comparison)":
            section("▶  EXECUTING ORIGINAL PROMPT...")
            original_output = execute_prompt(user_prompt)
            log(original_output)
            result["outputs"]["original"] = original_output

        section("▶  EXECUTING OPTIMIZED PROMPT...")
        optimized_output = execute_prompt(optimization["optimized_prompt"])
        log(optimized_output)
        result["outputs"]["optimized"] = optimized_output
    else:
        log("\n  ⏭  Skipping execution.")

    return result


# ── Interactive CLI Mode ───────────────────────────────────────────────────────
def interactive_mode():
    print("\n🤖 Prompt Optimization Agent")
    print("   Powered by OpenAI | Type 'quit' to exit\n")

    while True:
        print("-" * 60)
        user_input = input("Enter your prompt to optimize:\n> ").strip()

        if user_input.lower() in ("quit", "exit", "q"):
            print("Goodbye!")
            break
        if not user_input:
            print("Please enter a prompt.")
            continue

        try:
            result = run_optimizer(user_input)

            if ask_yes_no("\n💾 Save full result to JSON file?"):
                filename = "optimized_result.json"
                with open(filename, "w") as f:
                    json.dump(result, f, indent=2)
                print(f"  ✅ Saved to {filename}")

        except Exception as e:
            print(f"❌ Error: {e}")


# ── Programmatic Usage Example ─────────────────────────────────────────────────
def example_usage():
    """Example of using the optimizer programmatically."""
    weak_prompt = "write something about climate change"
    result = run_optimizer(weak_prompt)
    return result


# ── Entry Point ────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    import sys

    if not os.environ.get("OPENAI_API_KEY"):
        print("❌ Error: OPENAI_API_KEY environment variable not set.")
        print("   Run: export OPENAI_API_KEY='your-key-here'")
        sys.exit(1)

    if len(sys.argv) > 1:
        # Pass prompt as CLI argument: python prompt_optimizer_agent.py "your prompt here"
        prompt = " ".join(sys.argv[1:])
        run_optimizer(prompt)
    else:
        interactive_mode()