guptaaryan16/observability_env_test

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 gr10 11from observer_env.main import Action, TracingEnvironment, Observation, create_openenv_environment12 13try:14    from openenv.core.env_server.http_server import create_app15except ImportError:16    from openenv.core.env_server import create_app17 18 19DEFAULT_WORKSPACE = "real_example/microservices-demo"20WORKSPACE_DIR = os.getenv("WORKSPACE_DIR", DEFAULT_WORKSPACE if os.path.isdir(DEFAULT_WORKSPACE) else ".")21DATASET_PATH = os.getenv("DATASET_PATH", "dataset.json")22MAX_STEPS = int(os.getenv("MAX_STEPS", "20"))23 24env = TracingEnvironment(25    workspace_dir=WORKSPACE_DIR,26    dataset_path=DATASET_PATH,27    max_steps=MAX_STEPS,28)29 30 31def _task_label(task: Dict[str, Any]) -> str:32    return f"{task['id']} [{task['difficulty']}]"33 34 35TASK_LABELS: List[str] = [_task_label(task) for task in env.tasks]36TASK_LABEL_TO_INDEX: Dict[str, int] = {label: index for index, label in enumerate(TASK_LABELS)}37 38 39def _safe_json_dumps(data: Any) -> str:40    return json.dumps(data, indent=2, sort_keys=True, default=str)41 42 43def _build_action(action_type: str, payload: str) -> Action:44    if action_type in {"command", "read_code"}:45        if not payload.strip():46            raise ValueError(f"Payload is required for {action_type}")47        return Action(**{action_type: payload.strip()})48 49    if not payload.strip():50        raise ValueError(f"JSON payload is required for {action_type}")51 52    try:53        parsed = json.loads(payload)54    except json.JSONDecodeError as exc:55        raise ValueError(f"Invalid JSON payload: {exc}") from exc56 57    return Action(**{action_type: parsed})58 59 60def _task_help(task_label: str) -> str:61    task_index = TASK_LABEL_TO_INDEX.get(task_label, 0)62    task = env.tasks[task_index]63    workflow = task.get("workflow", [])64    workflow_md = "\n".join(f"- {step}" for step in workflow) if workflow else "- No workflow metadata"65    return (66        f"### Selected task\n"67        f"- **ID:** {task.get('id')}\n"68        f"- **Difficulty:** {task.get('difficulty')}\n"69        f"- **Workspace:** `{WORKSPACE_DIR}`\n"70        f"- **Dataset:** `{DATASET_PATH}`\n\n"71        f"**Description**\n{task.get('description', 'N/A')}\n\n"72        f"**Suggested workflow**\n{workflow_md}"73    )74 75 76def reset_episode(task_label: str) -> Tuple[str, str, str]:77    task_index = TASK_LABEL_TO_INDEX.get(task_label, 0)78    obs = env.reset(task_index=task_index)79    return _safe_json_dumps(obs.model_dump()), _safe_json_dumps(env.state()), _task_help(task_label)80 81 82def run_step(action_type: str, payload: str) -> Tuple[str, str, str, str]:83    action = _build_action(action_type, payload)84    obs, reward, done, info = env.step(action)85    return (86        _safe_json_dumps(obs.model_dump()),87        _safe_json_dumps(reward.model_dump()),88        str(done),89        _safe_json_dumps(info),90    )91 92 93def get_state() -> str:94    return _safe_json_dumps(env.state())95 96 97with gr.Blocks(title="Observability RCA OpenEnv Space") as demo:98    gr.Markdown(99        """100        # Observability RCA Workbench101 102        A guided Gradio workbench for log-to-code investigation.103        Select a task, reset the episode, then run actions to explore the environment.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 164openenv_app = create_app(165    create_openenv_environment,166    Action,167    Observation,168    env_name="observability_agent_training_env",169    max_concurrent_envs=1,170)171 172app = gr.mount_gradio_app(openenv_app, demo, path="/")173 174 175if __name__ == "__main__":176    import uvicorn177    uvicorn.run(app, host="0.0.0.0", port=7860)178