app.py

"""
AuditEnv OpenEnv Playground UI
Gradio frontend that wraps the FastAPI backend server.
"""
import json
import subprocess
import time
import gradio as gr
import requests

API_BASE = "http://127.0.0.1:8000"

# ──────────────────────────────────────────────
# Boot the FastAPI server in the background
# ──────────────────────────────────────────────
def _start_api_server():
    subprocess.Popen(
        ["uvicorn", "auditenv.server:app",
         "--host", "127.0.0.1", "--port", "8000", "--app-dir", "src"],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )

def _wait_for_api(max_retries: int = 40) -> bool:
    for _ in range(max_retries):
        try:
            r = requests.get(f"{API_BASE}/health", timeout=2)
            if r.status_code == 200:
                return True
        except Exception:
            pass
        time.sleep(1)
    return False

_start_api_server()
_api_ready = _wait_for_api()

# ──────────────────────────────────────────────
# API helpers
# ──────────────────────────────────────────────
def _fmt(response: requests.Response) -> str:
    try:
        return json.dumps(response.json(), indent=2)
    except Exception:
        return response.text

def api_reset(task_id: str) -> str:
    if not _api_ready:
        return '{"error": "API server failed to start"}'
    try:
        r = requests.post(f"{API_BASE}/reset", json={"task_id": task_id}, timeout=15)
        return _fmt(r)
    except Exception as e:
        return json.dumps({"error": str(e)})

def api_step(action_type: str, task_id: str, note: str,
             invoice_id: str, vendor_name: str, amount: float,
             violation_type: str, confidence: float) -> str:
    if not _api_ready:
        return '{"error": "API server failed to start"}'
    finding = None
    if action_type == "submit_finding":
        finding = {
            "invoice_id": invoice_id or None,
            "vendor_name": vendor_name or None,
            "amount": amount if amount else None,
            "violation_type": violation_type or None,
            "confidence": confidence,
        }
    payload = {
        "action_type": action_type,
        "task_id": task_id,
        "note": note or "",
        "finding": finding,
    }
    try:
        r = requests.post(f"{API_BASE}/step", json=payload, timeout=15)
        return _fmt(r)
    except Exception as e:
        return json.dumps({"error": str(e)})

def api_get_state() -> str:
    if not _api_ready:
        return '{"error": "API server failed to start"}'
    try:
        r = requests.get(f"{API_BASE}/state", timeout=15)
        return _fmt(r)
    except Exception as e:
        return json.dumps({"error": str(e)})

def api_health() -> str:
    try:
        r = requests.get(f"{API_BASE}/health", timeout=5)
        return _fmt(r)
    except Exception as e:
        return json.dumps({"error": str(e)})

# ──────────────────────────────────────────────
# Gradio UI
# ──────────────────────────────────────────────
CSS = """
#title { text-align: center; margin-bottom: 0.5rem; }
#subtitle { text-align: center; color: #6b7280; margin-bottom: 1.5rem; }
.action-btn { min-width: 110px; }
#json-out { font-family: monospace; font-size: 0.82rem; }
"""

QUICK_START_MD = """
## Quick Start

1. **Choose a task** (Easy / Medium / Hard)
2. Click **Reset** to start a new episode
3. Click **Get State** to see the documents your agent must audit
4. Build a **finding** and click **Step** to submit it
5. The environment returns a reward from 0.0 → 1.0

---

### Action types
| Action | Description |
|---|---|
| `submit_finding` | Flag a specific document as violating policy |
| `flag_human_review` | Escalate an ambiguous item |
| `noop` | Do nothing this step |

---

### Scoring
| Event | Reward |
|---|---|
| True positive (correct flag) | +0.2 |
| Full evidence chain | +0.3 |
| False positive | −0.5 |
| False negative | −0.2 |
"""

README_MD = """
## AuditEnv

**AuditEnv** is an OpenEnv-compatible reinforcement-learning environment
for autonomous compliance auditing.

Instead of eating pac-dots, the AI learns to be a corporate compliance
auditor — finding fraud, policy violations, and access anomalies.

### Difficulty Levels
| Level | Task | Max Steps |
|---|---|---|
| Easy | Expense Report Audit | 12 |
| Medium | Access Control Review | 20 |
| Hard | Multi-System Fraud Detection | 28 |

### API Endpoints
| Method | Path | Purpose |
|---|---|---|
| `POST` | `/reset` | Start new episode |
| `POST` | `/step` | Submit action |
| `GET` | `/state` | Read current state |
| `GET` | `/health` | Health check |
| `GET` | `/docs` | Swagger UI |
"""

with gr.Blocks(css=CSS, title="OpenEnv Agentic Environment: AuditEnv") as demo:

    gr.Markdown("# 🏢 OpenEnv Agentic Environment: AuditEnv", elem_id="title")
    gr.Markdown(
        "Autonomous compliance auditing — powered by reinforcement learning",
        elem_id="subtitle",
    )

    with gr.Row():
        # ── Left sidebar ──────────────────────────────
        with gr.Column(scale=1, min_width=220):
            gr.Markdown("### Navigation")
            with gr.Accordion("Quick Start", open=True):
                gr.Markdown(QUICK_START_MD)
            with gr.Accordion("README", open=False):
                gr.Markdown(README_MD)

        # ── Main area ─────────────────────────────────
        with gr.Column(scale=3):
            with gr.Tab("🎮 Playground"):
                gr.Markdown("### Playground\nClick **Reset** to start a new episode.")

                with gr.Row():
                    task_dd = gr.Dropdown(
                        choices=["easy", "medium", "hard"],
                        value="easy",
                        label="Task Difficulty",
                        scale=1,
                    )
                    action_radio = gr.Radio(
                        choices=["submit_finding", "flag_human_review", "noop"],
                        value="noop",
                        label="Action Type",
                        scale=2,
                    )

                note_box = gr.Textbox(
                    label="Note / Message",
                    placeholder="Optional note to attach to this action…",
                )

                with gr.Accordion("Finding Details (only for submit_finding)", open=False):
                    with gr.Row():
                        invoice_id_box = gr.Textbox(label="Invoice ID", placeholder="e.g. INV-0042")
                        vendor_box = gr.Textbox(label="Vendor Name", placeholder="e.g. Acme Corp")
                    with gr.Row():
                        amount_box = gr.Number(label="Amount (USD)", value=0.0)
                        violation_dd = gr.Dropdown(
                            choices=["duplicate_receipt", "alcohol_over_limit",
                                     "late_submission", "sod_conflict",
                                     "dormant_account", "shell_company",
                                     "invoice_splitting", "round_tripping"],
                            label="Violation Type",
                        )
                    confidence_slider = gr.Slider(0.0, 1.0, value=0.8, step=0.05, label="Confidence")

                with gr.Row():
                    step_btn  = gr.Button("▶ Step",      variant="primary",   elem_classes="action-btn")
                    reset_btn = gr.Button("🔄 Reset",    variant="secondary", elem_classes="action-btn")
                    state_btn = gr.Button("📊 Get State", variant="secondary", elem_classes="action-btn")

                gr.Markdown("#### Status — Raw JSON Response")
                json_out = gr.Code(language="json", label="", elem_id="json-out", lines=20)

                # Wire buttons
                reset_btn.click(
                    fn=api_reset,
                    inputs=[task_dd],
                    outputs=json_out,
                )
                step_btn.click(
                    fn=api_step,
                    inputs=[action_radio, task_dd, note_box,
                            invoice_id_box, vendor_box, amount_box,
                            violation_dd, confidence_slider],
                    outputs=json_out,
                )
                state_btn.click(
                    fn=api_get_state,
                    inputs=[],
                    outputs=json_out,
                )

            with gr.Tab("🔧 Custom / API"):
                gr.Markdown("""
### Direct API Access

Your AuditEnv server is also available at the `/docs` endpoint for full
Swagger UI access, or hit the endpoints directly:

```
POST /reset   — start a new episode
POST /step    — submit an action
GET  /state   — read current state
GET  /health  — health check
```
""")
                health_btn = gr.Button("🩺 Ping Health Check", variant="primary")
                health_out = gr.Code(language="json", label="Response", lines=5)
                health_btn.click(fn=api_health, inputs=[], outputs=health_out)

            with gr.Tab("📖 About"):
                gr.Markdown(README_MD)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)