from .models import call_model_text

BLUE_TEAM_MODEL = "gpt-4o-mini"

BLUE_TEAM_SYSTEM_PROMPT = """
You are the BLUE TEAM security agent.
Rewrite the user's prompt into a SECURE, HARDENED version.
Output ONLY the secure rewritten prompt.
"""

def blue_agent_rewrite(original_prompt: str, findings):
    combined = f"""
User prompt:
{original_prompt}

Security findings from Red Team:
{findings}

Rewrite into a secure prompt.
"""

    return call_model_text(
        model_name=BLUE_TEAM_MODEL,
        system_prompt=BLUE_TEAM_SYSTEM_PROMPT,
        user_prompt=combined,
        temperature=0.2,
    )
