app.py

import os
import re
import json
import gradio as gr
from openai import OpenAI

client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])

BASIC_PROMPT = """You are a helpful assistant that processes emails."""

STRUCTURED_PROMPT = """You are an email triage agent.

Your job is to read an email and decide how it should be handled.

For each email:
- Identify the main intent
- Classify priority as Low, Medium, or High
- Decide which team should own it
- Suggest the next action
- Draft a reply

Format your response clearly and use labels.
"""

ADVANCED_PROMPT = """You are an enterprise email triage agent.

Your job is to analyze incoming emails and decide how they should be handled.

For each email:
1. Identify the primary intent
2. Classify priority as Low, Medium, or High
3. Assign an owner from: Sales, Support, Finance, HR, Operations, Legal
4. Determine the next action from: Respond, Escalate, Delegate, Request More Information
5. Briefly explain your reasoning
6. Draft a professional reply

Output format:

Intent:
Priority:
Owner:
Action:
Reasoning:
Draft Reply:

Rules:
- Be concise and practical
- If the email is ambiguous, say so clearly
- If there is urgency, reflect that in both priority and action
- Keep the draft reply professional and appropriate for business use
"""

ADVANCED_TOOL_CALLING_PROMPT = """You are an enterprise email triage agent.

Your job is to analyze incoming emails and decide how they should be handled.

You have access to one tool:

lookup_status(reference_id)

Use this tool only when the email refers to a specific order, shipment, incident, or ticket and the current status is needed to answer well.

If you decide to use the tool, respond with exactly this format and nothing else:

CALL_TOOL: lookup_status <reference_id>

Examples:
CALL_TOOL: lookup_status SG-18427
CALL_TOOL: lookup_status INC-2041

If no tool is needed, then provide your final answer directly.

Final output format:

Intent:
Priority:
Owner:
Action:
Reasoning:
Draft Reply:

Rules:
- Be concise and practical
- If the email is ambiguous, say so clearly
- If there is urgency, reflect that in both priority and action
- Keep the draft reply professional and appropriate for business use
- Do not call the tool unless it is genuinely needed
"""

YOUR_REFERENCE_PROMPT = """You are an enterprise email triage agent.

Your job is to analyze incoming emails and decide how they should be handled.

You have access to one tool:

lookup_status(reference_id)

Use this tool only when the email refers to a specific order, shipment, incident, or ticket and the current status is needed to answer well.

If you decide to use the tool, respond with exactly this format and nothing else:

CALL_TOOL: lookup_status <reference_id>

Examples:
CALL_TOOL: lookup_status SG-18427
CALL_TOOL: lookup_status INC-2041

If you call the tool, do not include any explanation or answer. Only return the tool call.
If no tool is needed, then provide your final answer directly.

With no tool call the Final output format should be:

Intent:
Priority:
Owner:
Action:
Reasoning:
Draft Reply:

With tool call the Final output format should be:
CALL_TOOL: lookup_status <reference_id>

Rules:
- Be concise and practical
- If the email is ambiguous, say so clearly
- If there is urgency, reflect that in both priority and action
- Keep the draft reply professional and appropriate for business use
- Do not call the tool unless it is genuinely needed
"""

SAMPLE_EMAILS = {
    "Customer Complaint": """Subject: Delayed shipment SG-18427

We placed our order nearly two weeks ago and still have not received shipment SG-18427. We have also not had any meaningful update from your team despite following up twice.

This delay is now affecting our own delivery commitments to customers. Please let us know what is going on and when we can expect resolution.

Regards,
A frustrated customer""",

    "Sales Inquiry": """Subject: Exploring vendors for regional rollout

Hello,

We are evaluating potential partners for a Southeast Asia rollout later this year. We would like to understand your capabilities, pricing approach, and implementation timeline.

Could someone from your team reach out to discuss next steps?

Best,
Regional Director""",

    "Internal Finance": """Subject: Q3 budget clarification needed

Hi Finance team,

Can you confirm whether the additional marketing spend for Q3 has been approved? We need clarity before committing to the campaign launch next week.

Thanks,
Head of Marketing""",

    "HR Issue": """Subject: Concern about manager conduct

Hello,

I would like to report a concern regarding repeated inappropriate comments made by my manager during team meetings. I am uncomfortable raising this directly with them and would appreciate guidance on how this should be handled.

Please keep this confidential.

Thank you.""",

    "Urgent Escalation": """Subject: Update on incident INC-2041

Team,

A major outage is still affecting several enterprise clients. Can someone update us on incident INC-2041 and confirm the current status?

We need immediate support and clear ownership on incident response.

Please advise urgently.""",
}


def format_output(raw_text: str) -> str:
    """
    Light formatting only.
    We do NOT enforce a schema.
    We simply make common label-based outputs easier to read.
    """
    if not raw_text:
        return "No output received."

    known_labels = [
        "Intent:",
        "Priority:",
        "Owner:",
        "Action:",
        "Action Required:",
        "Reasoning:",
        "Draft Reply:",
    ]

    lines = raw_text.splitlines()
    formatted_lines = []

    for line in lines:
        stripped = line.strip()
        matched_label = next((label for label in known_labels if stripped.startswith(label)), None)

        if matched_label:
            value = stripped[len(matched_label):].strip()
            formatted_lines.append(f"### {matched_label[:-1]}")
            if value:
                formatted_lines.append(value)
            else:
                formatted_lines.append("")
            formatted_lines.append("")
        else:
            formatted_lines.append(line)

    return "\n".join(formatted_lines).strip()

STATUS_DB = {
    "SG-18427": {
        "record_type": "shipment",
        "status": "Delayed in customs",
        "last_update": "2026-04-03 14:20 SGT",
        "eta": "2026-04-09",
        "owner_team": "Operations",
    },
    "INC-2041": {
        "record_type": "incident",
        "status": "Engineering investigating root cause",
        "last_update": "2026-04-06 10:05 SGT",
        "eta": "Next update due in 30 minutes",
        "owner_team": "Operations",
    },
    "TKT-48392": {
        "record_type": "ticket",
        "status": "Assigned to support tier 2",
        "last_update": "2026-04-04 17:40 SGT",
        "eta": "Expected resolution within 1 business day",
        "owner_team": "Support",
    },
}


def lookup_status(reference_id: str) -> dict:
    ref = reference_id.strip().upper()
    result = STATUS_DB.get(ref)

    if result:
        return {
            "found": True,
            "reference_id": ref,
            **result
        }

    return {
        "found": False,
        "reference_id": ref,
        "message": "No matching record found in the status system."
    }


def detect_tool_call(text: str):
    pattern = r"^\s*CALL_TOOL:\s*(\w+)\s+([A-Z0-9\-]+)\s*$"
    match = re.match(pattern, text.strip())
    if match:
        return match.group(1), match.group(2)
    return None, None


def render_tool_trace(tool_used, tool_name="", tool_input=None, tool_output=None):
    if not tool_used:
        return """## Tool Trace

**Tool used:** No
"""

    input_block = json.dumps(tool_input, indent=2)
    output_block = json.dumps(tool_output, indent=2)

    return f"""## Tool Trace

**Tool used:** Yes  
**Tool name:** `{tool_name}`

### Tool input
```json
{input_block}
```
### Tool output
```json
{output_block}
```
"""

def run_agent(system_prompt, email, show_raw_output):
    if not system_prompt or not system_prompt.strip():
        return "Please enter a system prompt.", "", "## Tool Trace\n\nNo run yet."

    if not email or not email.strip():
        return "Please enter an email to analyze.", "", "## Tool Trace\n\nNo run yet."

    try:
        first_response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": email},
            ],
            temperature=0.4,
        )

        first_output = first_response.choices[0].message.content or ""
        tool_name, reference_id = detect_tool_call(first_output)

        if tool_name == "lookup_status" and reference_id:
            tool_result = lookup_status(reference_id)

            tool_trace = render_tool_trace(
                tool_used=True,
                tool_name=tool_name,
                tool_input={"reference_id": reference_id},
                tool_output=tool_result,
            )

            second_response = client.chat.completions.create(
                model="gpt-4o-mini",
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": email},
                    {"role": "assistant", "content": first_output},
                    {
                        "role": "user",
                        "content": (
                            "Tool result:\n"
                            f"{json.dumps(tool_result)}\n\n"
                            "Now provide the final answer for the email."
                        ),
                    },
                ],
                temperature=0.4,
            )

            raw_output = second_response.choices[0].message.content or ""
        else:
            tool_trace = render_tool_trace(tool_used=False)
            raw_output = first_output

        formatted_output = format_output(raw_output)

        if show_raw_output:
            return formatted_output, raw_output, tool_trace
        return formatted_output, "", tool_trace

    except Exception as e:
        error_message = f"Error: {str(e)}"
        return error_message, error_message, "## Tool Trace\n\nAn error occurred."


with gr.Blocks(title="Executive Agent Lab") as demo:
    gr.HTML("""
    <!-- Google tag (gtag.js) -->
    <script async src="https://www.googletagmanager.com/gtag/js?id=G-T4PLZJZLX0"></script>
    <script>
    window.dataLayer = window.dataLayer || [];
    function gtag(){dataLayer.push(arguments);}
    gtag('js', new Date());

    gtag('config', 'G-T4PLZJZLX0');
    </script>
    """)
    gr.Markdown("# Executive Agent Lab")
    gr.Markdown(
    "*Built by Tan Hong Ming (NUS) · For learning prompt engineering, agentic workflows, and tool use*"
)
    with gr.Accordion("How to Use This Lab", open=False):
        gr.Markdown("""
            ## What this lab is about

            This lab helps you practice three ideas:

            1. Prompt engineering  
            2. Agentic workflows  
            3. Agent tool use  

            You will improve an email triage agent by editing its system prompt.

            ## Suggested workflow

            ### Step 1
            Click **Load Basic** and then **Run Agent**.

            Observe what the agent does well and what it does poorly.

            ### Step 2
            Improve the system prompt so the agent can do the following for each email:

            1. Identify the primary intent  
            2. Classify priority as Low, Medium, or High  
            3. Assign an owner from: Sales, Support, Finance, HR, Operations, Legal  
            4. Determine the next action from: Respond, Escalate, Delegate, Request More Information  
            5. Briefly explain the reasoning  
            6. Draft a professional reply  

            ### Step 3
            Test your prompt on different sample emails.

            Compare how the output changes when your instructions become more specific.

            ### Step 4
            Get the agent to use the status lookup tool when needed.

            The tool call format is:

            `CALL_TOOL: lookup_status <reference_id>`

            Examples:

            `CALL_TOOL: lookup_status SG-18427`  
            `CALL_TOOL: lookup_status INC-2041`

            The tool should only be used when the email refers to a specific order, shipment, incident, or ticket and current status information is needed.

            ### Step 5
            After you get it working, look at the source code to understand what happens under the hood.

            ## What to look for in the code

            Here is the simple flow:

            1. The app sends your system prompt and the email to the model  
            2. The model either:
            - answers directly, or
            - outputs a tool call like `CALL_TOOL: lookup_status SG-18427`
            3. The app detects that tool call in the model output
            4. The Python function runs the lookup
            5. The tool result is sent back to the model
            6. The model produces the final answer

            So when you read the code, focus on these parts:

            - where the model is called
            - how `CALL_TOOL:` is detected
            - how `lookup_status()` works
            - how the tool result is passed back into the model"""
        )

        gr.Markdown("### Reference")

        gr.Markdown(
            "Once you have tried building your own prompt, you can load a working example below."
        )

        load_reference_btn = gr.Button("Load Reference Prompt")

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("## Agent Builder")
            gr.Markdown(
                "Define the agent's role, rules, decision logic, and output format."
            )

            agent_name = gr.Textbox(
                label="Agent Name (optional)",
                placeholder="Email Triage Agent",
            )

            system_prompt = gr.Textbox(
                label="System Prompt",
                lines=18,
                value=BASIC_PROMPT,
                placeholder="Write the instructions for your agent here..."
            )

            gr.Markdown("### Starter Prompts")
            with gr.Row():
                load_basic_btn = gr.Button("Load Basic")
                load_structured_btn = gr.Button("Load Structured")
                load_advanced_btn = gr.Button("Load Advanced")
                load_advanced_tool_btn = gr.Button("Load Advanced (Tool)")

            with gr.Accordion("What makes a good agent prompt?", open=False):
                gr.Markdown("""
                    - role  
                    - task  
                    - decision criteria  
                    - routing rules  
                    - output format  
                    - tone  
                    - handling ambiguity
                    - when to use the tool and when not to use it"""
                )

        with gr.Column(scale=1):
            gr.Markdown("## Test Input")
            gr.Markdown("Load a sample email or paste your own.")

            gr.Markdown("### Sample Emails")
            with gr.Row():
                btn_customer = gr.Button("Customer Complaint")
                btn_sales = gr.Button("Sales Inquiry")
            with gr.Row():
                btn_finance = gr.Button("Internal Finance")
                btn_hr = gr.Button("HR Issue")
            with gr.Row():
                btn_urgent = gr.Button("Urgent Escalation")

            email = gr.Textbox(
                label="Email Content",
                lines=20,
                value=SAMPLE_EMAILS["Customer Complaint"],
                placeholder="Paste or type an email here..."
            )

            run_btn = gr.Button("Run Agent", variant="primary")

        with gr.Column(scale=1):
            gr.Markdown("## Agent Output")
            gr.Markdown(
                "The app lightly formats common label-based outputs, but your prompt determines the structure."
            )

            formatted_output = gr.Markdown(
                value="Run the agent to see results here."
            )
            tool_trace = gr.Markdown(
                value="## Tool Trace\n\nNo run yet."
            )


            show_raw_output = gr.Checkbox(
                label="Show raw output",
                value=False
            )

            raw_output = gr.Textbox(
                label="Raw Output",
                lines=12,
                visible=True
            )

    def load_basic():
        return BASIC_PROMPT

    def load_structured():
        return STRUCTURED_PROMPT

    def load_advanced():
        return ADVANCED_PROMPT

    def load_advanced_tool():
        return ADVANCED_TOOL_CALLING_PROMPT
    
    def load_reference():
        return YOUR_REFERENCE_PROMPT
    
    def load_email(sample_name):
        return SAMPLE_EMAILS[sample_name]

    load_basic_btn.click(fn=load_basic, outputs=system_prompt)
    load_structured_btn.click(fn=load_structured, outputs=system_prompt)
    load_advanced_btn.click(fn=load_advanced, outputs=system_prompt)
    load_advanced_tool_btn.click(fn=load_advanced_tool, outputs=system_prompt)
    load_reference_btn.click(fn=load_reference, outputs=system_prompt)

    btn_customer.click(fn=lambda: load_email("Customer Complaint"), outputs=email)
    btn_sales.click(fn=lambda: load_email("Sales Inquiry"), outputs=email)
    btn_finance.click(fn=lambda: load_email("Internal Finance"), outputs=email)
    btn_hr.click(fn=lambda: load_email("HR Issue"), outputs=email)
    btn_urgent.click(fn=lambda: load_email("Urgent Escalation"), outputs=email)

    run_btn.click(
        fn=run_agent,
        inputs=[system_prompt, email, show_raw_output],
        outputs=[formatted_output, raw_output, tool_trace],
    )

demo.launch(server_name="0.0.0.0", server_port=7860)