server/app.py

"""
app.py — Executive Assistant OpenEnv Environment

FastAPI server + environment logic for email and calendar management.
"""

import sys
import os
from pathlib import Path

# Ensure server/ directory is on the path
sys.path.insert(0, str(Path(__file__).parent))

from fastapi import FastAPI, HTTPException
from models import AssistantAction, AssistantObservation, AssistantState
from fastapi.responses import HTMLResponse
from typing import Optional
import statistics

# Import scoring functions from data.py (teammate will implement these)
from data import (
    generate_scenario,
    compute_email_quality,
    check_scheduling_correctness,
    compute_conflict_resolution,
    apply_penalties,
    TASK_DEFINITIONS,
)


# ============================================================
# THE ENVIRONMENT CLASS
# ============================================================

class ExecAssistEnv:
    def __init__(self):
        self.current_scenario = None
        self.calendar_state = None
        self.episode_done = False
        self.steps_taken = 0
        self.total_score = 0.0
        self.current_task = None
        self.seed = 42

    def reset(self, task: str = "easy"):
        """Start a new episode."""
        
        if task not in TASK_DEFINITIONS:
            raise ValueError(f"Unknown task: {task}. Choose from: easy, medium, hard")

        self.current_task = task
        self.episode_done = False
        self.steps_taken = 0
        self.total_score = 0.0

        # Generate scenario (teammate implements this in data.py)
        self.current_scenario = generate_scenario(task, seed=self.seed)
        
        return {
            "observation": self._build_observation(),
            "reward": 0.0,
            "done": False,
            "info": {
                "task": task,
                "scenario_id": self.current_scenario.get("id", "unknown"),
            }
        }

    def step(self, action: dict):
        """Agent submits action — environment scores it."""
        
        if self.episode_done:
            return {
                "observation": {"message": "Episode is done. Call /reset to start again."},
                "reward": 0.0,
                "done": True,
                "info": {"total_score": self.total_score}
            }

        # Parse action
        try:
            assistant_action = AssistantAction(**action)
        except Exception as e:
            return {
                "observation": {"message": f"Invalid action format: {str(e)}"},
                "reward": -0.5,
                "done": False,
                "info": {"error": "invalid_action_format"}
            }

        # Validate basic action structure
        if not assistant_action.email_reply or len(assistant_action.email_reply.strip()) == 0:
            return {
                "observation": {"message": "Empty email reply. Penalty applied."},
                "reward": -0.3,
                "done": False,
                "info": {"error": "empty_email_reply"}
            }

        if assistant_action.calendar_action not in ["book", "propose_alternatives", "reschedule", "decline"]:
            return {
                "observation": {"message": f"Invalid calendar_action: {assistant_action.calendar_action}"},
                "reward": -0.2,
                "done": False,
                "info": {"error": "invalid_calendar_action"}
            }

        # Compute rewards using teammate's functions
        email_score = compute_email_quality(
            assistant_action.email_reply,
            self.current_scenario
        )
        
        # Convert meeting_details to dict if it exists
        meeting_details_dict = assistant_action.meeting_details.dict() if assistant_action.meeting_details else None

        scheduling_result = check_scheduling_correctness(
        meeting_details_dict,
        self.current_scenario
        )
        
        conflict_score = compute_conflict_resolution(
        assistant_action.dict(),  # ← Add .dict() here
        self.current_scenario
        )
        
         
        penalty = apply_penalties(assistant_action.dict(), self.current_scenario)  

        # Combine scores based on task difficulty
        task_def = TASK_DEFINITIONS[self.current_task]
        weights = task_def["reward_weights"]
        
        total_reward = (
            weights["email"] * email_score +
            weights["scheduling"] * scheduling_result["score"] +
            weights["conflict"] * conflict_score
        )
        
        total_reward = max(0.0, min(1.0, total_reward - penalty))
        
        self.total_score = total_reward
        self.episode_done = True
        self.steps_taken += 1

        return {
            "observation": self._build_completion_message(assistant_action, total_reward),
            "reward": round(total_reward, 4),
            "done": True,
            "info": {
                "email_score": round(email_score, 4),
                "scheduling_score": round(scheduling_result["score"], 4),
                "conflict_score": round(conflict_score, 4),
                "penalty": round(penalty, 4),
                "scheduling_checks": scheduling_result.get("checks", {}),
            }
        }

    def _build_observation(self) -> dict:
        """Build what the agent sees."""
        
        scenario = self.current_scenario
        task_def = TASK_DEFINITIONS[self.current_task]
        
        obs = {
            "task": self.current_task,
            "description": task_def["description"],
            "emails": scenario["emails"],
            "calendar": scenario["calendar"],
            "contacts": scenario.get("contacts", {}),
            "action_required": task_def["action_required"],
        }
        
        return obs

    def _build_completion_message(self, action: AssistantAction, score: float) -> dict:
        """Build feedback message after step."""
        
        if score >= 0.9:
            message = f"Excellent work! Score: {score:.2f}"
        elif score >= 0.7:
            message = f"Good response. Score: {score:.2f}"
        elif score >= 0.5:
            message = f"Acceptable. Score: {score:.2f}"
        else:
            message = f"Needs improvement. Score: {score:.2f}"
        
        return {
            "message": message,
            "email_sent": action.email_reply[:100] + "..." if len(action.email_reply) > 100 else action.email_reply,
            "calendar_action": action.calendar_action,
        }

    def state(self):
        """Return current state."""
        return {
            "current_task": self.current_task,
            "emails_pending": len(self.current_scenario.get("emails", [])) if self.current_scenario else 0,
            "episode_done": self.episode_done,
            "steps_taken": self.steps_taken,
            "total_score": self.total_score,
        }


# ============================================================
# FASTAPI SERVER
# ============================================================

app = FastAPI(
    title="ExecAssist Environment",
    description=(
        "An OpenEnv environment where AI agents learn to manage email and calendar "
        "for a busy executive. Agents must draft professional replies, schedule meetings, "
        "and resolve conflicts."
    ),
    version="1.0.0"
)

@app.get("/", response_class=HTMLResponse, include_in_schema=False)
async def root():
    """Friendly landing page for the bare URL."""
    return """
    <!DOCTYPE html>
    <html lang="en">
    <head>
      <meta charset="utf-8">
      <title>ExecAssist — OpenEnv Environment</title>
      <style>
        body {
          font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
          max-width: 720px;
          margin: 60px auto;
          padding: 0 24px;
          background: #0f0f1a;
          color: #e8e8f0;
          line-height: 1.6;
        }
        h1 { color: #5b9dff; margin-bottom: 4px; }
        .tag { color: #8888aa; font-size: 14px; }
        .grid { display: grid; grid-template-columns: 1fr 1fr; gap: 12px; margin: 24px 0; }
        .grid-3 { grid-template-columns: 1fr 1fr 1fr; }
        @media (max-width: 600px) { .grid-3 { grid-template-columns: 1fr; } }
        .card {
          background: #1a1a2e;
          border: 1px solid #2a2a44;
          border-radius: 8px;
          padding: 14px 18px;
          text-decoration: none;
          color: #e8e8f0;
          transition: all 0.15s;
        }
        .card:hover { border-color: #5b9dff; background: #1f1f38; }
        .card-title { color: #5b9dff; font-weight: 600; margin-bottom: 4px; }
        .card-desc { color: #aaaabb; font-size: 13px; }
        table { border-collapse: collapse; width: 100%; margin: 16px 0; }
        th, td { text-align: left; padding: 8px 12px; border-bottom: 1px solid #2a2a44; }
        th { color: #5b9dff; font-weight: 600; }
        code { background: #1a1a2e; padding: 2px 6px; border-radius: 4px; color: #a8d4ff; }
        h3 { color: #8ab4f8; margin-top: 32px; }
      </style>
    </head>
    <body>
      <h1>ExecAssist</h1>
      <div class="tag">An OpenEnv environment for executive assistant email + calendar tasks · Theme #3.2</div>

      <p>An AI agent reads incoming emails, checks the calendar, drafts professional replies, and books meetings, that too without double-booking or breaking working hours. Trained with TRL GRPO on Qwen2.5-0.5B.</p>

      <h3>Read the writeup</h3>
      <div class="grid grid-3">
       <a class="card" href="https://youtu.be/iit3PgGBR4Y" target="_blank">
        <div class="card-title">🎥 Video walkthrough</div>
        <div class="card-desc">90-second YouTube demo from Team PsPredator</div>
       </a>
       <a class="card" href="/blog">
        <div class="card-title">📝 Blog post</div>
        <div class="card-desc">The story: collapse, fix, results, why it matters</div>
       </a>
       <a class="card" href="https://huggingface.co/spaces/DevanshuDon/exec-assist/blob/main/README.md">
        <div class="card-title">📖 README</div>
        <div class="card-desc">Full documentation, reward design, training pipeline</div>
       </a>
      </div>

      <h3>Try the API</h3>
      <div class="grid">
        <a class="card" href="/docs">
          <div class="card-title">📘 Interactive API docs</div>
          <div class="card-desc">Swagger UI : try /reset and /step</div>
        </a>
        <a class="card" href="/health">
          <div class="card-title">💚 Health check</div>
          <div class="card-desc">Service liveness probe</div>
        </a>
        <a class="card" href="/tasks">
          <div class="card-title">📋 List tasks</div>
          <div class="card-desc">All 3 difficulty levels</div>
        </a>
        <a class="card" href="/metadata">
          <div class="card-title">🧾 Metadata</div>
          <div class="card-desc">Environment info + version</div>
        </a>
      </div>

      <h3>Training results</h3>
      <table>
        <tr><th>Task</th><th>Baseline</th><th>Trained (GRPO)</th><th>Δ</th></tr>
        <tr><td>Easy</td><td>0.345</td><td><b>0.995</b></td><td>+188%</td></tr>
        <tr><td>Medium</td><td>0.227</td><td><b>0.745</b></td><td>+228%</td></tr>
        <tr><td>Hard</td><td>0.249</td><td><b>0.737</b></td><td>+196%</td></tr>
      </table>

      <p style="color:#8888aa; font-size:13px; margin-top:32px;">
        Built for the OpenEnv Hackathon · April 2026 · by
        <a href="https://huggingface.co/DevanshuDon" style="color:#5b9dff;">@DevanshuDon</a>
      </p>
    </body>
    </html>
    """

@app.get("/blog", response_class=HTMLResponse, include_in_schema=False)
async def blog():
    """Blog post about the hackathon submission."""
    return """
    <!DOCTYPE html>
    <html lang="en">
    <head>
      <meta charset="utf-8">
      <meta name="viewport" content="width=device-width, initial-scale=1">
      <title>Teaching a 0.5B Model to Be an Executive Assistant</title>
      <style>
        body {
          font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', 'Roboto', sans-serif;
          max-width: 720px;
          margin: 60px auto;
          padding: 0 24px;
          background: #0f0f1a;
          color: #e8e8f0;
          line-height: 1.7;
        }
        h1 { color: #5b9dff; margin-bottom: 8px; font-size: 32px; }
        h2 { color: #8ab4f8; margin-top: 40px; font-size: 22px; }
        .subtitle { color: #8888aa; font-style: italic; margin-bottom: 32px; }
        p { margin: 16px 0; }
        code {
          background: #1a1a2e;
          padding: 2px 6px;
          border-radius: 4px;
          color: #a8d4ff;
          font-family: 'Monaco', 'Courier New', monospace;
          font-size: 14px;
        }
        pre {
          background: #1a1a2e;
          border: 1px solid #2a2a44;
          border-radius: 8px;
          padding: 16px;
          overflow-x: auto;
        }
        pre code { background: none; padding: 0; }
        a { color: #5b9dff; text-decoration: none; }
        a:hover { text-decoration: underline; }
        img {
          max-width: 100%;
          border-radius: 8px;
          margin: 24px 0;
          background: white;
          padding: 8px;
        }
        .footer {
          margin-top: 60px;
          padding-top: 24px;
          border-top: 1px solid #2a2a44;
          color: #8888aa;
          font-size: 14px;
        }
      </style>
    </head>
    <body>

    <h1>Teaching a 0.5B Model to Be an Executive Assistant</h1>
    <div class="subtitle">36 hours, 90 minutes of training, three model collapses, and one bar chart that nearly killed me.</div>

    <hr style="border: none; border-top: 1px solid #2a2a44; margin: 32px 0;">

    <p>Let me start with the moment I thought I'd lost the hackathon.</p>

    <p>It's around 3pm on day two. I've been training for an hour and a half on a free Colab T4, and the cell finally finishes. I run the eval. The plot pops up. Baseline bars in red, trained bars in green. Easy task: 0.493 to 0.200. Medium: 0.348 to 0.200. Hard: 0.331 to 0.186.</p>

    <p>The trained model was <em>worse</em>. Identically worse on every task, by exactly the same number every time. I sat there for a full minute trying to figure out if I was reading the chart wrong.</p>

    <p>Turns out I wasn't. This is what GRPO collapse looks like, and now I had a textbook example of it. The model had given up exploring, found a single response that scored exactly 0.2 against my reward function no matter what input you gave it, and locked itself in. All my training was technically optimization. It was just optimization toward "say the dumbest possible thing every time and never deviate."</p>

    <p>Here's how I dug out of it.</p>

    <h2>The environment</h2>

    <p>I built <strong>ExecAssist</strong>, an OpenEnv environment that simulates an executive's morning inbox. The agent gets emails (sender, subject, body, priority), a calendar (existing meetings, working hours), and contacts. It has to spit out a JSON action: an email reply, a calendar action like <code>book</code> or <code>propose_alternatives</code>, and meeting details if it's booking something.</p>

    <p>Three difficulty tiers:</p>
    <ul>
      <li><strong>Easy.</strong> One email, calendar is wide open. Don't blow the basics.</li>
      <li><strong>Medium.</strong> The requested time conflicts with an existing meeting. Spot it, propose alternatives.</li>
      <li><strong>Hard.</strong> Three emails, multi-party coordination, priority conflicts. Actually plan.</li>
    </ul>

    <p>Reward is a weighted blend of three independent graders. Email quality scores politeness markers, greeting and closing, sufficient detail. Scheduling correctness checks for double-booking, working hours, sensible duration. Conflict resolution looks for whether the agent recognized a clash and proposed real alternatives. Then four anti-hacking penalties on top: short emails, missing meeting details, generic phrasing, overly long responses.</p>

    <p>I went with multiple independent graders because I'd read this line in the hackathon guide that stuck with me: <em>"if you only have a single reward signal, it is easier for the model to hack it."</em> Building it the right way from day one felt like the obvious move. I would later be very glad I did.</p>

    <h2>What collapsed, and why</h2>

    <p>The first config looked reasonable on paper:</p>

    <pre><code>GRPOConfig(
    learning_rate=5e-6,
    num_train_epochs=1,
    num_generations=4,
    # no beta term
)</code></pre>

    <p>Looking at the training log afterward, I could see exactly what happened. Reward bounced between 0.0 and 0.4 for the first seven steps, peaked at 0.397, then crashed and pinned itself at 0.14 for the next 38 steps. The model had found a safe floor and stopped trying anything that might score higher because variance was punishing it.</p>

    <p>The diagnosis came down to three problems. No KL penalty meant nothing was anchoring the trained policy to the base model, so it was free to drift to any degenerate local minimum. The learning rate was too aggressive, which compounded the drift. And one epoch wasn't enough runway to recover even if it wanted to.</p>

    <p>The fix was three changes:</p>

    <pre><code>GRPOConfig(
    learning_rate=1e-6,           # 5x slower
    num_train_epochs=3,           # 3x longer
    num_generations=8,            # more variety per step
    beta=0.1,                     # KL penalty, the critical one
)</code></pre>

    <p>I also dropped in a "reload clean model" cell before training started. The previous run had corrupted the weights, and I didn't want to start the next attempt from a broken policy. Then I hit Run All and walked away for 90 minutes.</p>

    <h2>The second run</h2>

    <p>Came back to find the cell still chugging along. 218 of 270 steps done. Trained weights were already in memory, so I ran the eval cells anyway and held my breath while the bars rendered.</p>

    <p>Easy task: 0.345 to <strong>0.995</strong>.<br>
    Medium: 0.227 to <strong>0.745</strong>.<br>
    Hard: 0.249 to <strong>0.737</strong>.</p>

    <p>I made a noise. I'm not going to describe what kind.</p>

    <p>Nine out of ten samples on the easy task hit a perfect 1.0. The model wasn't just lucky on those runs, it had learned the structure of the task. You could see it in the variance. Baseline scores ranged from 0.0 to 0.65 on the same prompt depending on how the dice rolled. Trained scores were tight: 0.95 to 1.0 on easy, 0.68 to 0.82 on medium, 0.65 to 0.80 on hard.</p>

    <p>The training curve told the same story. First quartile mean reward: 0.390. Last quartile: 0.648. A 66% lift during training itself, on top of the much bigger gap between trained and untrained at evaluation time.</p>

    <img src="https://huggingface.co/spaces/DevanshuDon/exec-assist/resolve/main/training_results.png" alt="Training results showing reward and loss curves with baseline vs trained comparison" onerror="this.style.display='none';this.nextElementSibling.style.display='block';">
    <p style="display:none; padding: 16px; background: #1a1a2e; border-radius: 8px; color: #aaa; text-align: center;">
      📊 Training plot available in the <a href="https://huggingface.co/spaces/DevanshuDon/exec-assist/blob/main/training_results.png">repository</a>
    </p>

    <h2>The interesting part: watching the model try to cheat</h2>

    <p>Because the reward had multiple independent components instead of one big scalar, I could see exactly how the model tried to game each one during training. Going through the early-step rollouts in order:</p>

    <ul>
      <li><strong>Steps 8 to 15:</strong> the model started outputting just the JSON action and skipping the email body entirely. The short-email penalty (-0.30 if under 20 words) caught this. Within about 30 steps, every output had a real email attached.</li>
      <li><strong>Around step 25:</strong> it started scheduling meetings at 8am or 6pm. Outside working hours. The scheduling-correctness grader returned 0 for the working-hours check, and within 50 steps the model had figured that out.</li>
      <li><strong>Around step 50:</strong> generic templated phrasing showed up. "Thank you for your email. I'll check the calendar and get back to you." Polite, vague, useless. The generic-phrasing detector docked these. Specific responses came back.</li>
      <li><strong>Throughout:</strong> the model would occasionally just forget the <code>meeting_details</code> block. The -0.40 penalty made that a non-strategy fast.</li>
    </ul>

    <p>Most submissions claim they have anti-hacking penalties. Showing the penalties firing during real training, on a real curve, is the rare part. It's the difference between saying a multi-grader rubric works and demonstrating that it does.</p>

    <h2>The bigger claim</h2>

    <p>One sanity check from earlier in the day stuck with me. Same three tasks, same scoring, but evaluated against an untuned <strong>Nemotron 120B</strong> running through OpenRouter via the standard <code>inference.py</code> baseline. It averaged <strong>0.337</strong> across the three tasks.</p>

    <p>After 90 minutes of GRPO, a model <strong>240 times smaller</strong> is averaging <strong>0.83</strong> on the same environment. Free Colab T4. Zero-dollar cloud bill.</p>

    <p>That's the point of training-environment design. A reward signal that's hard to game, a model that's allowed to actually train against it, and you get a result that beats a frontier model on its native task.</p>

    <p>I think there's a research-shaped argument hiding in here. Frontier LLMs are notoriously bad at structured calendar reasoning. Try asking any production agent to find a 30-minute slot that doesn't conflict with your standups. ExecAssist isolates that specific failure mode into a tractable RL target. The result suggests that for a class of structured personal-task workflows, task-specific RL on a small model is a real alternative to scaling up. That feels like a workshop paper, maybe.</p>

    <h2>Try it yourself</h2>

    <ul>
      <li><strong>Live environment:</strong> <a href="https://devanshudon-exec-assist.hf.space/docs">devanshudon-exec-assist.hf.space/docs</a>. Hit <code>POST /reset?task=easy</code>, then <code>POST /step</code> with your action. Swagger does most of the work.</li>
      <li><strong>Baseline:</strong> <code>python inference.py</code> reproduces the untrained scores around 0.32 average.</li>
      <li><strong>Training:</strong> the Colab notebook is in the repo. Set runtime to T4, Run All, about 50 minutes including evaluation.</li>
      <li><strong>Repo:</strong> all the code, the working hyperparameters, the broken hyperparameters that caused the collapse, the results JSON, and this writeup.</li>
    </ul>

    <p>The environment is hard in interesting ways. Try to break it. I'd be curious what the model learns to game next.</p>

    <div class="footer">
      Built for the OpenEnv Hackathon (April 2026), Theme #3.2: Personalized Tasks.<br>
      Models, environment, and training results are all on the <a href="https://huggingface.co/spaces/DevanshuDon/exec-assist">HF Space</a>.
    </div>

    </body>
    </html>
    """

    
env = ExecAssistEnv()


@app.post("/reset")
def reset(task: str = "easy"):
    return env.reset(task)


@app.post("/step")
def step(action: AssistantAction):
    return env.step(action.dict())


@app.get("/state")
def state():
    return env.state()


@app.get("/tasks")
def tasks():
    return {
        task_name: {
            "description": td["description"],
            "action_required": td["action_required"],
            "reward_weights": td["reward_weights"],
        }
        for task_name, td in TASK_DEFINITIONS.items()
    }


@app.get("/health")
def health():
    return {"status": "healthy"}


# ============================================================
# OPENENV REQUIRED ENDPOINTS
# ============================================================

@app.get("/metadata")
def metadata():
    """Return environment name and description."""
    return {
        "name": "exec-assist",
        "description": (
            "Executive Assistant environment where AI agents learn to manage email "
            "and calendar for busy professionals. Agents must balance professionalism, "
            "scheduling correctness, and conflict resolution."
        ),
        "version": "1.0.0",
        "author": "Devanshu(Team PsPredator)",
        "tasks": ["easy", "medium", "hard"],
    }


@app.get("/schema")
def schema():
    """Return action, observation, and state schemas."""
    return {
        "action": {
            "type": "object",
            "properties": {
                "email_reply": {"type": "string"},
                "calendar_action": {"type": "string", "enum": ["book", "propose_alternatives", "reschedule", "decline"]},
                "meeting_details": {"type": "object"},
            },
            "required": ["email_reply", "calendar_action"],
        },
        "observation": {
            "type": "object",
            "properties": {
                "task": {"type": "string"},
                "emails": {"type": "array"},
                "calendar": {"type": "object"},
                "contacts": {"type": "object"},
            },
        },
        "state": {
            "type": "object",
            "properties": {
                "current_task": {"type": "string"},
                "emails_pending": {"type": "integer"},
                "episode_done": {"type": "boolean"},
                "steps_taken": {"type": "integer"},
                "total_score": {"type": "number"},
            },
        },
    }


@app.post("/mcp")
async def mcp_endpoint(request_body: dict = {}):
    """MCP JSON-RPC endpoint."""
    method = request_body.get("method", "")
    req_id = request_body.get("id", 1)

    if method == "initialize":
        return {
            "jsonrpc": "2.0",
            "id": req_id,
            "result": {
                "protocolVersion": "2024-11-05",
                "serverInfo": {"name": "exec-assist", "version": "1.0.0"},
                "capabilities": {"tools": {"listChanged": False}},
            },
        }

    elif method == "tools/list":
        return {
            "jsonrpc": "2.0",
            "id": req_id,
            "result": {
                "tools": [
                    {
                        "name": "reset",
                        "description": "Start new episode (easy/medium/hard)",
                        "inputSchema": {
                            "type": "object",
                            "properties": {"task": {"type": "string", "enum": ["easy", "medium", "hard"]}},
                        },
                    },
                    {
                        "name": "step",
                        "description": "Submit email reply and calendar action",
                        "inputSchema": {
                            "type": "object",
                            "properties": {
                                "email_reply": {"type": "string"},
                                "calendar_action": {"type": "string"},
                                "meeting_details": {"type": "object"},
                            },
                            "required": ["email_reply", "calendar_action"],
                        },
                    },
                    {"name": "state", "description": "Get current state", "inputSchema": {"type": "object"}},
                ],
            },
        }

    return {"jsonrpc": "2.0", "id": req_id, "result": {}}


def main():
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)


if __name__ == "__main__":
    main()