source/space/app.py

"""Public, proposal-only Gradio demo for BarunAction-35M candidate-v2."""

from __future__ import annotations

import hashlib
import html
import json
import os
import threading
from pathlib import Path
from typing import Any

try:
    import gradio as gr
except ImportError:  # Keeps the core request path importable for offline tests.
    gr = None  # type: ignore[assignment]

MODEL_ID = "harrrshall/BarunAction-35M"
MODEL_REVISION = "candidate-v2"
MODEL_PARAMETERS = 35_072_768
SOURCE_RELEASE = "https://github.com/harrrshall/barunaction-35m/tree/v1.0.0"

EXPECTED_CHECKPOINT_SHA256 = {
    "barun_config.json": "9b3a1d71baa95a198744d250f9629231738d942570b8685c44307fd83dd33565",
    "model.safetensors": "fdb95ccf58a095e0d321be998924318b35ee59a334f6dd97d8726d2cf80021d3",
    "tokenizer.json": "70ded9605fccd09c2340ca7e225361eab0ae8b4dbbb0d6e26343ab5183979db6",
}
CHECKPOINT_MANIFEST_SHA256 = "c743ab7c4d33ae75c6b0aa4547458a961b92766da8fcf85fd148fda2ebb5530a"

MAX_REQUEST_CHARS = 4_000
MAX_NOW_CHARS = 128
MAX_TOOLS_BYTES = 64_000
MAX_CONTEXT_BYTES = 32_000
MAX_NEW_TOKENS = 192

DEFAULT_TOOL_SCHEMAS: list[dict[str, Any]] = [
    {
        "name": "create_calendar_event",
        "description": "Create a new calendar event.",
        "arguments": {
            "title": {"type": "string", "description": "Calendar event title."},
            "datetime": {
                "type": "string",
                "description": "Calendar event date and time.",
            },
        },
        "required": ["title", "datetime"],
        "additional_arguments": False,
        "side_effecting": True,
    },
    {
        "name": "create_contact",
        "description": "Create a contact in an external address book.",
        "arguments": {
            "first_name": {"type": "string", "description": "Contact first name."},
            "last_name": {"type": "string", "description": "Contact last name."},
            "phone_number": {"type": "string", "description": "Contact phone number."},
            "email": {"type": "string", "description": "Contact email address."},
        },
        "required": ["first_name", "last_name"],
        "additional_arguments": False,
        "side_effecting": True,
    },
    {
        "name": "open_wifi_settings",
        "description": "Open the device Wi-Fi settings screen.",
        "arguments": {},
        "required": [],
        "additional_arguments": False,
        "side_effecting": False,
    },
    {
        "name": "send_email",
        "description": "Draft an email for an external mail client.",
        "arguments": {
            "to": {"type": "string", "description": "Recipient email address."},
            "subject": {"type": "string", "description": "Email subject line."},
            "body": {"type": "string", "description": "Email body text."},
        },
        "required": ["to", "subject"],
        "additional_arguments": False,
        "side_effecting": True,
    },
    {
        "name": "show_map",
        "description": "Show a map for a place or search query.",
        "arguments": {"query": {"type": "string", "description": "Place or map search query."}},
        "required": ["query"],
        "additional_arguments": False,
        "side_effecting": False,
    },
    {
        "name": "turn_off_flashlight",
        "description": "Turn off the device flashlight.",
        "arguments": {},
        "required": [],
        "additional_arguments": False,
        "side_effecting": True,
    },
    {
        "name": "turn_on_flashlight",
        "description": "Turn on the device flashlight.",
        "arguments": {},
        "required": [],
        "additional_arguments": False,
        "side_effecting": True,
    },
]

DEFAULT_TOOLS_JSON = json.dumps(DEFAULT_TOOL_SCHEMAS, indent=2, ensure_ascii=False)
DEFAULT_CONTEXT_JSON = json.dumps(
    {
        "device": {"flashlight": "off"},
        "locale": "en-IN",
        "timezone": "Asia/Kolkata",
    },
    indent=2,
)
DEFAULT_NOW = "2026-08-05T11:30:00+05:30"

EXAMPLES: dict[str, dict[str, str]] = {
    "Map · Bengaluru landmark": {
        "request": "Show me Cubbon Park in Bengaluru",
        "context": DEFAULT_CONTEXT_JSON,
        "now": DEFAULT_NOW,
    },
    "Calendar · explicit date and time": {
        "request": "Create a calendar event called Design review for 2026-08-08 at 3:30 PM",
        "context": DEFAULT_CONTEXT_JSON,
        "now": DEFAULT_NOW,
    },
    "Contact · name, phone, and email": {
        "request": "Create a contact for Mira Shah, +91 98765 43210, mira@example.com",
        "context": DEFAULT_CONTEXT_JSON,
        "now": DEFAULT_NOW,
    },
    "Email · structured arguments": {
        "request": (
            "Email dev@example.com with subject Release notes and body "
            "The candidate package is ready for review."
        ),
        "context": DEFAULT_CONTEXT_JSON,
        "now": DEFAULT_NOW,
    },
    "Device · flashlight": {
        "request": "Turn on the flashlight",
        "context": DEFAULT_CONTEXT_JSON,
        "now": DEFAULT_NOW,
    },
    "Settings · Wi-Fi": {
        "request": "Open Wi-Fi settings",
        "context": DEFAULT_CONTEXT_JSON,
        "now": DEFAULT_NOW,
    },
}

_COMPILER: Any | None = None
_COMPILER_LOCK = threading.Lock()


class PublicInputError(ValueError):
    """Safe input error whose message can be shown in the public UI."""

    def __init__(self, code: str, message: str, path: str) -> None:
        super().__init__(message)
        self.code = code
        self.message = message
        self.path = path


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _verify_checkpoint_files(checkpoint_dir: str | Path) -> None:
    root = Path(checkpoint_dir)
    for name, expected in EXPECTED_CHECKPOINT_SHA256.items():
        path = root / name
        if not path.is_file():
            raise RuntimeError(f"required checkpoint file is missing: {name}")
        if _sha256(path) != expected:
            raise RuntimeError(f"checkpoint digest mismatch: {name}")

    manifest = root / "checkpoint_manifest.json"
    if not manifest.is_file():
        raise RuntimeError("required checkpoint file is missing: checkpoint_manifest.json")
    if _sha256(manifest) != CHECKPOINT_MANIFEST_SHA256:
        raise RuntimeError("checkpoint digest mismatch: checkpoint_manifest.json")


def _runtime_device() -> str:
    device = os.getenv("BARUNACTION_DEVICE", "cpu").strip().casefold()
    if device not in {"cpu", "cuda"}:
        raise RuntimeError("BARUNACTION_DEVICE must be cpu or cuda")
    return device


def _get_compiler() -> Any:
    """Download and verify the public candidate only after the first request."""

    global _COMPILER
    if _COMPILER is not None:
        return _COMPILER

    with _COMPILER_LOCK:
        if _COMPILER is not None:
            return _COMPILER

        from huggingface_hub import snapshot_download

        checkpoint_dir = snapshot_download(
            repo_id=MODEL_ID,
            revision=MODEL_REVISION,
            repo_type="model",
            allow_patterns=[*EXPECTED_CHECKPOINT_SHA256, "checkpoint_manifest.json"],
            token=False,
        )
        _verify_checkpoint_files(checkpoint_dir)

        from barunaction import BarunActionCompiler

        compiler = BarunActionCompiler(
            checkpoint_dir,
            expected_sha256=EXPECTED_CHECKPOINT_SHA256,
            device=_runtime_device(),
        )
        _COMPILER = compiler
        return compiler


def _reject_pairs(pairs: list[tuple[str, Any]]) -> dict[str, Any]:
    result: dict[str, Any] = {}
    for key, value in pairs:
        if key in result:
            raise ValueError(f"duplicate object key {key!r}")
        result[key] = value
    return result


def _reject_constant(value: str) -> None:
    raise ValueError(f"non-finite JSON number {value!r} is not allowed")


def _parse_json_input(raw: Any, *, label: str, path: str, byte_limit: int) -> Any:
    if not isinstance(raw, str):
        raise PublicInputError("type_mismatch", f"{label} must be JSON text.", path)
    if len(raw.encode("utf-8")) > byte_limit:
        raise PublicInputError(
            "input_too_large",
            f"{label} exceeds this public demo's {byte_limit:,}-byte limit.",
            path,
        )
    try:
        return json.loads(
            raw,
            object_pairs_hook=_reject_pairs,
            parse_constant=_reject_constant,
        )
    except (json.JSONDecodeError, ValueError) as error:
        raise PublicInputError(
            "invalid_json", f"{label} is not strict JSON: {error}", path
        ) from error


def _safe_boundary(*, proposal_available: bool) -> dict[str, Any]:
    return {
        "execution_permitted": False,
        "external_side_effects": False,
        "model_proposal_available": proposal_available,
        "space_executes_tools": False,
    }


def _public_provenance(*, verified: bool) -> dict[str, Any]:
    return {
        "candidate_id": MODEL_REVISION,
        "checkpoint_sha256": dict(EXPECTED_CHECKPOINT_SHA256),
        "checkpoint_verified": verified,
        "model_id": MODEL_ID,
        "parameter_count": MODEL_PARAMETERS,
        "revision": MODEL_REVISION,
        "source_release": SOURCE_RELEASE,
    }


def _error_response(
    *,
    stage: str,
    code: str,
    message: str,
    path: str = "$",
    raw_output: str = "",
    verified: bool = False,
) -> tuple[str, None, dict[str, Any], str, dict[str, Any]]:
    safe_stage = html.escape(stage)
    safe_code = html.escape(code)
    safe_path = html.escape(path)
    safe_message = html.escape(message)
    status = (
        "### No validated proposal\n"
        f"**{safe_stage} · `{safe_code}` · `{safe_path}`**  \n"
        f"{safe_message}  \n\n"
        "**Nothing was executed.** Edit the inputs and try again."
    )
    provenance = _public_provenance(verified=verified)
    provenance["error"] = {"code": code, "path": path, "stage": stage}
    return status, None, _safe_boundary(proposal_available=False), raw_output, provenance


def _proposal_explanation(action: dict[str, Any], policy: dict[str, Any]) -> str:
    decision = str(action.get("decision", "UNKNOWN"))
    calls = action.get("calls", [])
    call_count = len(calls) if isinstance(calls, list) else 0

    if decision == "ABSTAIN":
        summary = "The model abstained and proposed no tool call."
    elif decision == "CLARIFY":
        summary = "The model requested clarification and proposed no tool call."
    elif decision == "CONFIRM":
        summary = f"The model proposed {call_count} call(s) and explicitly requested confirmation."
    else:
        summary = f"The model proposed {call_count} typed call(s)."

    gates: list[str] = []
    if policy.get("authorization_required"):
        gates.append("external authorization")
    if policy.get("confirmation_required"):
        gates.append("external confirmation")
    gate_text = " and ".join(gates) if gates else "no runtime execution grant"

    return (
        "### Validated Action IR proposal\n"
        f"{summary} The validator reports **{gate_text}**.  \n\n"
        "**Nothing was executed.** This Space has no tool handlers and always keeps "
        "`execution_permitted: false`."
    )


def compile_action(
    request: Any,
    tools_json: Any,
    context_json: Any,
    now: Any,
) -> tuple[str, dict[str, Any] | None, dict[str, Any], str, dict[str, Any]]:
    """Validate inputs, lazily run the model, and return a proposal-only view."""

    try:
        if not isinstance(request, str) or not request.strip():
            raise PublicInputError(
                "empty_request", "Request must contain non-whitespace text.", "$.request"
            )
        if len(request) > MAX_REQUEST_CHARS:
            raise PublicInputError(
                "input_too_large",
                f"Request exceeds this public demo's {MAX_REQUEST_CHARS:,}-character limit.",
                "$.request",
            )
        if not isinstance(now, str) or not now.strip():
            raise PublicInputError(
                "missing_now", "NOW must be a timezone-aware ISO-8601 timestamp.", "$.now"
            )
        if len(now) > MAX_NOW_CHARS:
            raise PublicInputError(
                "input_too_large",
                f"NOW exceeds this public demo's {MAX_NOW_CHARS}-character limit.",
                "$.now",
            )

        tools = _parse_json_input(
            tools_json,
            label="Tool schemas",
            path="$.tools",
            byte_limit=MAX_TOOLS_BYTES,
        )
        context = _parse_json_input(
            context_json,
            label="Context",
            path="$.context",
            byte_limit=MAX_CONTEXT_BYTES,
        )
        if not isinstance(tools, list):
            raise PublicInputError("type_mismatch", "Tool schemas must be a JSON array.", "$.tools")
        if not isinstance(context, dict):
            raise PublicInputError("type_mismatch", "Context must be a JSON object.", "$.context")
    except PublicInputError as error:
        return _error_response(
            stage="input",
            code=error.code,
            message=error.message,
            path=error.path,
        )

    try:
        compiler = _get_compiler()
    except Exception:  # noqa: BLE001 - public load failures must be returned, not crash the Space.
        return _error_response(
            stage="load",
            code="model_unavailable",
            message=(
                "The pinned public checkpoint could not be downloaded, hash-verified, or loaded. "
                "Please retry in a moment."
            ),
        )

    try:
        outcome = compiler.infer(
            request=request,
            tool_schemas=tools,
            context=context,
            now=now,
            max_new_tokens=MAX_NEW_TOKENS,
        )
    except Exception:  # noqa: BLE001 - inference must fail closed for every runtime failure.
        return _error_response(
            stage="inference",
            code="unexpected_runtime_error",
            message="Inference stopped safely before a validated proposal was returned.",
            verified=True,
        )

    if not outcome.ok:
        error = outcome.error
        return _error_response(
            stage=error.stage,
            code=error.code,
            message=error.message,
            path=error.path,
            raw_output=outcome.raw_output or "",
            verified=True,
        )

    action = outcome.action.to_dict()
    runtime_policy = outcome.policy.to_dict()
    safety = {
        **runtime_policy,
        "external_side_effects": False,
        "space_executes_tools": False,
    }
    provenance = _public_provenance(verified=True)
    provenance.update(
        {
            "checkpoint_format": outcome.checkpoint_format,
            "generated_tokens": outcome.generated_tokens,
            "prompt_contract_version": "barunaction-local-prompt-v1",
            "prompt_sha256": outcome.prompt_sha256,
            "prompt_tokens": outcome.prompt_tokens,
            "result_schema_version": "barunaction-inference-result-v1",
            "runtime_candidate_id": outcome.candidate_id,
        }
    )
    return (
        _proposal_explanation(action, runtime_policy),
        action,
        safety,
        outcome.raw_output or "",
        provenance,
    )


def load_example(name: str) -> tuple[str, str, str, str]:
    example = EXAMPLES.get(name, EXAMPLES[next(iter(EXAMPLES))])
    return example["request"], DEFAULT_TOOLS_JSON, example["context"], example["now"]


CSS = """
:root {
  --barun-ink: #172033;
  --barun-muted: #667085;
  --barun-indigo: #4f46e5;
  --barun-violet: #7c3aed;
  --barun-surface: rgba(255,255,255,.78);
}
.gradio-container {
  max-width: 1180px !important;
  margin: 0 auto !important;
  color: var(--barun-ink);
}
.barun-hero {
  position: relative;
  overflow: hidden;
  padding: 34px 36px;
  margin: 10px 0 20px;
  border: 1px solid rgba(99,102,241,.20);
  border-radius: 24px;
  background:
    radial-gradient(circle at 88% 10%, rgba(124,58,237,.18), transparent 32%),
    linear-gradient(135deg, rgba(238,242,255,.96), rgba(250,245,255,.92));
  box-shadow: 0 18px 50px rgba(63,55,201,.09);
}
.barun-eyebrow {
  color: var(--barun-indigo);
  font-size: .78rem;
  font-weight: 750;
  letter-spacing: .12em;
  text-transform: uppercase;
}
.barun-hero h1 {
  margin: 8px 0 4px;
  font-size: clamp(2.1rem, 5vw, 3.6rem);
  letter-spacing: -.055em;
  line-height: 1;
}
.barun-hero p { max-width: 760px; color: #475467; font-size: 1.03rem; }
.barun-pills { display: flex; flex-wrap: wrap; gap: 8px; margin-top: 18px; }
.barun-pill {
  padding: 7px 11px;
  border: 1px solid rgba(79,70,229,.18);
  border-radius: 999px;
  background: rgba(255,255,255,.72);
  color: #3730a3;
  font-size: .82rem;
  font-weight: 650;
}
.barun-safety {
  padding: 14px 17px;
  margin: 0 0 18px;
  border-left: 4px solid #16a34a;
  border-radius: 10px;
  background: rgba(240,253,244,.84);
  color: #166534;
}
.barun-footer { color: var(--barun-muted); font-size: .86rem; text-align: center; padding: 16px; }
#compile-button { min-height: 46px; font-weight: 700; }
"""


def build_demo() -> Any:
    if gr is None:
        raise RuntimeError("Gradio is required to build the public Space")

    theme = gr.themes.Soft(
        primary_hue="indigo",
        secondary_hue="violet",
        neutral_hue="slate",
        font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
        font_mono=["ui-monospace", "SFMono-Regular", "monospace"],
    )

    with gr.Blocks(theme=theme, css=CSS, title="BarunAction-35M") as demo:
        gr.HTML(
            """
            <section class="barun-hero">
              <div class="barun-eyebrow">Compact typed-action research model</div>
              <h1>BarunAction-35M</h1>
              <p>
                Turn a request and explicit tool schemas into validated Action IR—using only
                35,072,768 parameters. Explore the proposal, policy gates, and exact model output.
              </p>
              <div class="barun-pills">
                <span class="barun-pill">35.1M parameters</span>
                <span class="barun-pill">Strict JSON + schema validation</span>
                <span class="barun-pill">Hash-pinned candidate-v2</span>
                <span class="barun-pill">No real side effects</span>
              </div>
            </section>
            <div class="barun-safety">
              <strong>Proposal-only sandbox.</strong> This Space never contacts a person, edits a
              calendar, changes a device, or invokes any declared tool. Every result keeps
              <code>execution_permitted: false</code>.
            </div>
            """
        )

        with gr.Row(equal_height=False):
            with gr.Column(scale=6):
                gr.Markdown("## Compose a request")
                with gr.Row():
                    example_name = gr.Dropdown(
                        choices=list(EXAMPLES),
                        value=next(iter(EXAMPLES)),
                        label="Curated scenario",
                        scale=4,
                    )
                    load_example_button = gr.Button("Load example", scale=1)

                request = gr.Textbox(
                    value=EXAMPLES[next(iter(EXAMPLES))]["request"],
                    label="Request",
                    placeholder="Describe one personal action…",
                    lines=3,
                    max_lines=7,
                )

                with gr.Accordion("Tool schemas", open=False):
                    tools = gr.Code(
                        value=DEFAULT_TOOLS_JSON,
                        language="json",
                        label="Editable barunaction-tool-schema-v1 declarations",
                        lines=18,
                    )

                with gr.Accordion("Context and reference time", open=True):
                    context = gr.Code(
                        value=DEFAULT_CONTEXT_JSON,
                        language="json",
                        label="Context JSON",
                        lines=8,
                    )
                    now = gr.Textbox(
                        value=DEFAULT_NOW,
                        label="NOW",
                        info="Timezone-aware ISO-8601, including an explicit UTC offset.",
                    )

                compile_button = gr.Button(
                    "Compile to Action IR",
                    variant="primary",
                    elem_id="compile-button",
                )
                gr.Markdown(
                    "The first request lazily downloads and verifies the ~141 MB public "
                    "checkpoint. "
                    "Generation is deterministic and capped at 192 new tokens."
                )

            with gr.Column(scale=5):
                gr.Markdown("## Inspect the proposal")
                status = gr.Markdown(
                    "### Ready\nChoose an example or enter a request. Nothing runs until you ask "
                    "for a proposal—and declared tools are never executed."
                )
                action = gr.JSON(label="Validated Action IR")
                safety = gr.JSON(
                    value=_safe_boundary(proposal_available=False),
                    label="Safety and policy boundary",
                )
                with gr.Accordion("Raw model continuation", open=False):
                    raw_output = gr.Code(label="Unmodified continuation", language="json", lines=9)
                with gr.Accordion("Verified provenance", open=False):
                    provenance = gr.JSON(
                        value=_public_provenance(verified=False),
                        label="Model identity and request provenance",
                    )

        load_example_button.click(
            fn=load_example,
            inputs=example_name,
            outputs=[request, tools, context, now],
            api_name=False,
        )
        compile_button.click(
            fn=compile_action,
            inputs=[request, tools, context, now],
            outputs=[status, action, safety, raw_output, provenance],
            api_name="compile_action",
        )
        request.submit(
            fn=compile_action,
            inputs=[request, tools, context, now],
            outputs=[status, action, safety, raw_output, provenance],
            api_name=False,
        )

        gr.HTML(
            f"""
            <div class="barun-footer">
              <a href="https://huggingface.co/{MODEL_ID}" target="_blank">Model card</a>
              &nbsp;·&nbsp;
              <a href="{SOURCE_RELEASE}" target="_blank">Source v1.0.0</a>
              &nbsp;·&nbsp; Apache-2.0 &nbsp;·&nbsp; Harrrshall, 2026
            </div>
            """
        )

    return demo


demo = build_demo() if gr is not None else None


if __name__ == "__main__":
    if demo is None:
        raise RuntimeError("Install the Space requirements before launching the app")
    demo.queue(default_concurrency_limit=1, max_size=16).launch()