guptaaryan16/observability_env

1"""Hugging Face Space entrypoint with Gradio UI + OpenEnv API endpoints."""2 3from __future__ import annotations4 5import json6import os7from typing import Any, Dict, List, Tuple8 9import gradio as gr10from fastapi.responses import RedirectResponse11 12from observer_env.main import Action, TracingEnvironment13from server.app import app as openenv_api_app14 15 16DEFAULT_WORKSPACE = "real_example/microservices-demo"17WORKSPACE_DIR = os.getenv("WORKSPACE_DIR", DEFAULT_WORKSPACE if os.path.isdir(DEFAULT_WORKSPACE) else ".")18DATASET_PATH = os.getenv("DATASET_PATH", "dataset.json")19MAX_STEPS = int(os.getenv("MAX_STEPS", "20"))20 21env = TracingEnvironment(22    workspace_dir=WORKSPACE_DIR,23    dataset_path=DATASET_PATH,24    max_steps=MAX_STEPS,25)26 27 28def _task_label(task: Dict[str, Any]) -> str:29    return f"{task['id']} [{task['difficulty']}]"30 31 32TASK_LABELS: List[str] = [_task_label(task) for task in env.tasks]33TASK_LABEL_TO_INDEX: Dict[str, int] = {label: index for index, label in enumerate(TASK_LABELS)}34 35 36def _safe_json_dumps(data: Any) -> str:37    return json.dumps(data, indent=2, sort_keys=True, default=str)38 39 40def _build_action(action_type: str, payload: str) -> Action:41    if action_type in {"command", "read_code"}:42        if not payload.strip():43            raise ValueError(f"Payload is required for {action_type}")44        return Action(**{action_type: payload.strip()})45 46    if not payload.strip():47        raise ValueError(f"JSON payload is required for {action_type}")48 49    try:50        parsed = json.loads(payload)51    except json.JSONDecodeError as exc:52        raise ValueError(f"Invalid JSON payload: {exc}") from exc53 54    return Action(**{action_type: parsed})55 56 57def _task_help(task_label: str) -> str:58    task_index = TASK_LABEL_TO_INDEX.get(task_label, 0)59    task = env.tasks[task_index]60    workflow = task.get("workflow", [])61    workflow_md = "\n".join(f"- {step}" for step in workflow) if workflow else "- No workflow metadata"62    return (63        f"### Selected task\n"64        f"- **ID:** {task.get('id')}\n"65        f"- **Difficulty:** {task.get('difficulty')}\n"66        f"- **Workspace:** `{WORKSPACE_DIR}`\n"67        f"- **Dataset:** `{DATASET_PATH}`\n\n"68        f"**Description**\n{task.get('description', 'N/A')}\n\n"69        f"**Suggested workflow**\n{workflow_md}"70    )71 72 73def reset_episode(task_label: str) -> Tuple[str, str, str]:74    task_index = TASK_LABEL_TO_INDEX.get(task_label, 0)75    obs = env.reset(task_index=task_index)76    return _safe_json_dumps(obs.model_dump()), _safe_json_dumps(env.state()), _task_help(task_label)77 78 79def run_step(action_type: str, payload: str) -> Tuple[str, str, str, str]:80    action = _build_action(action_type, payload)81    obs, reward, done, info = env.step(action)82    return (83        _safe_json_dumps(obs.model_dump()),84        _safe_json_dumps(reward.model_dump()),85        str(done),86        _safe_json_dumps(info),87    )88 89 90def get_state() -> str:91    return _safe_json_dumps(env.state())92 93 94with gr.Blocks(title="Observability RCA OpenEnv Space") as demo:95    gr.Markdown(96        """97        # Observability RCA Workbench98 99        This Space exposes both:100        - OpenEnv API endpoints (`/reset`, `/step`, `/state`, `/schema`, `/health`)101        - A guided Gradio workbench for log-to-code investigation102 103        API routes remain available alongside the UI.104        """105    )106 107    with gr.Row(equal_height=True):108        task_name = gr.Dropdown(109            choices=TASK_LABELS,110            value=TASK_LABELS[0] if TASK_LABELS else None,111            label="Task",112            info="Choose a task by name (not numeric index).",113        )114        reset_btn = gr.Button("Reset")115        state_btn = gr.Button("Get state")116 117    task_help = gr.Markdown(value=_task_help(TASK_LABELS[0]) if TASK_LABELS else "No tasks available")118 119    reset_observation = gr.Code(label="Reset observation", language="json")120    state_output = gr.Code(label="State", language="json")121 122    with gr.Tab("Action Runner"):123        action_type = gr.Dropdown(124            choices=["command", "read_code", "get_cached_result", "submit_rca"],125            value="command",126            label="Action type",127            info="Use command/read_code for exploration. JSON required for cached result and RCA submission.",128        )129        payload = gr.Textbox(130            lines=8,131            label="Payload",132            value="echo 'ERR_CART_OVERLOAD from checkoutservice'",133        )134        step_btn = gr.Button("Run step", variant="primary")135 136        with gr.Row():137            step_done = gr.Textbox(label="Done")138            step_reward = gr.Code(label="Reward", language="json")139 140        step_observation = gr.Code(label="Observation", language="json")141        step_info = gr.Code(label="Info", language="json")142 143    with gr.Tab("Payload Templates"):144        gr.Markdown(145            """146            Copy these into **Payload** depending on action type:147 148            - `command`: `rg -n "ERR_CART_OVERLOAD|SMTP_PORT|CalculateQuote" logs src`149            - `read_code`: `src/checkoutservice/main.go:220-280`150            - `get_cached_result`: `{\"query\": \"checkout error\", \"top_k\": 3}`151            - `submit_rca`: `{\"buggy_file\": \"src/checkoutservice/main.go\", \"line_number\": 245, \"explanation\": \"...\"}`152            """153        )154 155    task_name.change(fn=_task_help, inputs=[task_name], outputs=[task_help])156    reset_btn.click(fn=reset_episode, inputs=[task_name], outputs=[reset_observation, state_output, task_help])157    step_btn.click(158        fn=run_step,159        inputs=[action_type, payload],160        outputs=[step_observation, step_reward, step_done, step_info],161    )162    state_btn.click(fn=get_state, outputs=[state_output])163 164 165# Mount Gradio under /ui to keep API paths unchanged.166app = gr.mount_gradio_app(openenv_api_app, demo, path="/")167 168if __name__ == "__main__":169    import uvicorn170    # Hugging Face Spaces exposes port 7860171    uvicorn.run(app, host="0.0.0.0", port=7860)