server/main.py

"""
Launch Desk — FastAPI Backend
==============================
Serves the frontend and exposes an SSE streaming endpoint that runs
the Launch Planner agent via the OpenAI Agents SDK.
"""

from __future__ import annotations

import json
import os
from pathlib import Path
from typing import Any

from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse
from pydantic import BaseModel, Field

from agents import Runner
from agents.items import (
    ToolCallItem,
    ToolCallOutputItem,
)
from agents.stream_events import (
    RawResponsesStreamEvent,
    RunItemStreamEvent,
)

from server.agent import agent

# ---------------------------------------------------------------------------
# App setup
# ---------------------------------------------------------------------------

app = FastAPI(title="Launch Desk")

FRONTEND_DIR = Path(__file__).resolve().parent.parent / "frontend"


# ---------------------------------------------------------------------------
# Request model
# ---------------------------------------------------------------------------

class LaunchBrief(BaseModel):
    brief: str = Field(default="", description="Product brief text")
    audience: str = Field(default="", description="Target audience")
    launch_date: str = Field(default="", description="Target launch date YYYY-MM-DD")
    constraints: str = Field(default="", description="Constraints or requirements")
    assets: str = Field(default="", description="Available assets")


# ---------------------------------------------------------------------------
# SSE helper
# ---------------------------------------------------------------------------

def _sse(data: dict) -> str:
    """Format a dict as a Server-Sent Event line."""
    return f"data: {json.dumps(data)}\n\n"


def _build_prompt(req: LaunchBrief) -> str:
    """Compose the user prompt from the structured request."""
    parts: list[str] = []
    if req.brief:
        parts.append(f"## Product Brief\n{req.brief}")
    if req.audience:
        parts.append(f"## Target Audience\n{req.audience}")
    if req.launch_date:
        parts.append(f"## Launch Date\n{req.launch_date}")
    if req.constraints:
        parts.append(f"## Constraints\n{req.constraints}")
    if req.assets:
        parts.append(f"## Available Assets\n{req.assets}")

    if not parts:
        parts.append(
            "No launch details were provided. Ask the user for a product "
            "brief, audience, and target launch date."
        )

    parts.append(
        "\n## Instructions\n"
        "1. Use extract_tasks to break the brief into prioritised tasks.\n"
        "2. Use check_readiness to assess launch readiness.\n"
        "3. Use generate_checklist to create per-owner checklists.\n"
        "4. Use draft_copy to write launch copy for email and twitter.\n"
        "5. Produce a final plan with these sections:\n"
        "   - Prioritized Plan\n"
        "   - Risk Register\n"
        "   - Owner Checklist\n"
        "   - Launch Copy\n"
        "   - Follow-up Questions"
    )
    return "\n\n".join(parts)


# ---------------------------------------------------------------------------
# SSE streaming generator
# ---------------------------------------------------------------------------

async def _stream_agent(req: LaunchBrief):
    """Async generator yielding SSE events from the agent run."""
    prompt = _build_prompt(req)

    # Track which tool calls we've seen so we can pair start/end events.
    # call_id -> tool_name
    seen_calls: dict[str, str] = {}

    try:
        result = Runner.run_streamed(agent, input=prompt)

        async for event in result.stream_events():

            # --- Raw LLM streaming events (text deltas) ---
            if isinstance(event, RawResponsesStreamEvent):
                raw = event.data
                # ResponseTextDeltaEvent has type "response.output_text.delta"
                if hasattr(raw, "type") and raw.type == "response.output_text.delta":
                    delta_text = getattr(raw, "delta", "")
                    if delta_text:
                        yield _sse({"type": "text_delta", "text": delta_text})
                continue

            # --- Run item events (tool calls, tool outputs, messages) ---
            if isinstance(event, RunItemStreamEvent):
                item = event.item

                # Tool call started
                if isinstance(item, ToolCallItem):
                    tool_name = item.tool_name or "unknown_tool"
                    call_id = item.call_id or ""
                    if call_id:
                        seen_calls[call_id] = tool_name
                    # Try to extract arguments from the raw item
                    args: Any = {}
                    raw_item = item.raw_item
                    if isinstance(raw_item, dict):
                        args = raw_item.get("arguments", {})
                    elif hasattr(raw_item, "arguments"):
                        args = raw_item.arguments
                    # Parse args if they're a JSON string
                    if isinstance(args, str):
                        try:
                            args = json.loads(args)
                        except (json.JSONDecodeError, TypeError):
                            pass
                    yield _sse({
                        "type": "tool_start",
                        "name": tool_name,
                        "args": args,
                    })
                    continue

                # Tool call completed
                if isinstance(item, ToolCallOutputItem):
                    raw_item = item.raw_item
                    call_id = item.call_id or ""
                    tool_name = seen_calls.get(call_id, "unknown_tool")
                    output = item.output
                    # Try to parse output as JSON for nicer SSE payload
                    if isinstance(output, str):
                        try:
                            output = json.loads(output)
                        except (json.JSONDecodeError, TypeError):
                            pass
                    yield _sse({
                        "type": "tool_end",
                        "name": tool_name,
                        "result": output,
                    })
                    continue

        # --- Done ---
        final_output = result.final_output
        yield _sse({"type": "done", "output": final_output})

    except Exception as exc:
        yield _sse({
            "type": "error",
            "message": f"Agent error: {exc}",
        })


# ---------------------------------------------------------------------------
# Endpoints
# ---------------------------------------------------------------------------

@app.get("/api/health")
async def health():
    return {"status": "ok"}


@app.post("/api/agent/stream")
async def agent_stream(req: LaunchBrief):
    return StreamingResponse(
        _stream_agent(req),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",
        },
    )


@app.get("/", response_class=HTMLResponse)
async def serve_frontend():
    index_path = FRONTEND_DIR / "index.html"
    if index_path.exists():
        return HTMLResponse(content=index_path.read_text(encoding="utf-8"))
    return HTMLResponse(content="<h1>Frontend not found</h1>", status_code=404)


# ---------------------------------------------------------------------------
# Entry point for `python -m server.main` or `uvicorn server.main:app`
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import uvicorn

    uvicorn.run(
        "server.main:app",
        host="0.0.0.0",
        port=8001,
        reload=False,
    )