cid_demo/ui.py

from dataclasses import dataclass

import gradio as gr

from cid.defaults import (
    DEFAULT_ALLOCATION_THRESHOLD,
    DEFAULT_BINDING_THRESHOLD,
    DEFAULT_CONVERGENCE_THRESHOLD,
    DEFAULT_NEED_THRESHOLD,
)

from .config import MODEL_CHOICES, MODEL_ID
from .examples import EXAMPLES, WORKSPACE_PLACEHOLDER
from .presentation import CUSTOM_CSS, FOOTER_HTML, HERO_HTML, READY_SUMMARY_HTML
from .runtime import run_cid


@dataclass(frozen=True)
class DemoInputs:
    prompt: gr.Textbox
    workspace: gr.Textbox
    calculator_tool: gr.Checkbox
    symbolic_math_tool: gr.Checkbox
    workspace_tool: gr.Checkbox
    max_steps: gr.Slider
    need_threshold: gr.Slider
    convergence_threshold: gr.Slider
    allocation_threshold: gr.Slider
    binding_threshold: gr.Slider
    model_selector: gr.Dropdown


@dataclass(frozen=True)
class DemoOutputs:
    final_display: gr.Textbox
    summary: gr.HTML
    display_evolution: gr.Markdown
    trace: gr.Dataframe


def _build_input_panel() -> tuple[DemoInputs, gr.Button]:
    gr.Markdown("### Try CID")
    model_selector = gr.Dropdown(
        choices=MODEL_CHOICES,
        value=MODEL_ID,
        label="Model",
        interactive=True,
        elem_id="model-selector",
    )
    prompt = gr.Textbox(
        label="Prompt",
        value="What is 7 + 5?",
        lines=4,
        max_lines=8,
        placeholder="Ask a short question...",
        elem_id="cid-prompt",
    )

    gr.Markdown("#### Tools")
    with gr.Row():
        calculator_tool = gr.Checkbox(
            label="Calculator",
            value=False,
            info="Deterministic numeric expressions",
        )
        symbolic_math_tool = gr.Checkbox(
            label="Symbolic math",
            value=False,
            info="Exact algebra and calculus",
        )
        workspace_tool = gr.Checkbox(
            label="Workspace",
            value=False,
            info="Task-local document search and read",
        )

    with gr.Accordion("Task-local workspace", open=False):
        workspace = gr.Textbox(
            label="Task-local documents",
            value="",
            lines=9,
            max_lines=18,
            placeholder=WORKSPACE_PLACEHOLDER,
            info=(
                "Enable the Workspace tool to edit these documents. Separate documents with a line "
                "containing ---. Start a block with 'Title:' or '#'."
            ),
            interactive=False,
            elem_id="cid-workspace",
        )

    workspace_tool.change(
        fn=lambda enabled: gr.Textbox(interactive=bool(enabled)),
        inputs=workspace_tool,
        outputs=workspace,
        show_progress="hidden",
    )

    gr.Examples(
        examples=EXAMPLES,
        inputs=[prompt, calculator_tool, symbolic_math_tool, workspace_tool, workspace],
        example_labels=[item[0] for item in EXAMPLES],
        label="Examples",
        elem_id="examples",
    )

    with gr.Accordion("Runtime controls", open=False):
        max_steps = gr.Slider(1, 256, value=128, step=1, label="Max CID model steps")
        with gr.Row():
            need_threshold = gr.Slider(
                0.05,
                0.95,
                value=DEFAULT_NEED_THRESHOLD,
                step=0.05,
                label="Need threshold",
            )
            convergence_threshold = gr.Slider(
                0.05,
                0.95,
                value=DEFAULT_CONVERGENCE_THRESHOLD,
                step=0.05,
                label="Convergence threshold",
            )
        with gr.Row():
            allocation_threshold = gr.Slider(
                0.05,
                0.95,
                value=DEFAULT_ALLOCATION_THRESHOLD,
                step=0.05,
                label="Allocation threshold",
            )
            binding_threshold = gr.Slider(
                0.05,
                0.95,
                value=DEFAULT_BINDING_THRESHOLD,
                step=0.05,
                label="Binding threshold",
            )

    run_button = gr.Button("Run CID", variant="primary", elem_id="run-button")
    inputs = DemoInputs(
        prompt=prompt,
        workspace=workspace,
        calculator_tool=calculator_tool,
        symbolic_math_tool=symbolic_math_tool,
        workspace_tool=workspace_tool,
        max_steps=max_steps,
        need_threshold=need_threshold,
        convergence_threshold=convergence_threshold,
        allocation_threshold=allocation_threshold,
        binding_threshold=binding_threshold,
        model_selector=model_selector,
    )
    return inputs, run_button


def _build_result_panel() -> tuple[gr.Textbox, gr.HTML]:
    gr.Markdown("### Result")
    final_display = gr.Textbox(
        label="Final CID Display",
        lines=7,
        interactive=False,
        elem_id="final-display",
    )
    gr.Markdown("#### Run at a glance")
    summary = gr.HTML(READY_SUMMARY_HTML, elem_id="run-stats")
    return final_display, summary


def _build_inspection_tabs() -> tuple[gr.Markdown, gr.Dataframe]:
    with gr.Tabs(elem_id="inspection"):
        with gr.Tab("Display evolution"):
            gr.Markdown(
                "The visible Display is materialized repeatedly; changes below expose revision "
                "across model steps."
            )
            display_evolution = gr.Markdown()
        with gr.Tab("Runtime trace"):
            gr.Markdown(
                "Low-level runtime events, including TCT lifecycle transitions and any tool-related activity."
            )
            trace = gr.Dataframe(
                headers=["step", "time_s", "event", "details"],
                datatype=["number", "number", "str", "str"],
                interactive=False,
                wrap=True,
                elem_id="runtime-trace",
            )
    return display_evolution, trace


def build_demo() -> gr.Blocks:
    with gr.Blocks(
        title="CID · Demo",
        css=CUSTOM_CSS,
        theme=gr.themes.Soft(
            primary_hue=gr.themes.colors.indigo,
            neutral_hue=gr.themes.colors.slate,
        ),
    ) as demo:
        gr.HTML(HERO_HTML, elem_id="cid-hero")

        with gr.Row(equal_height=False):
            with gr.Column(scale=5, min_width=360, elem_id="prompt-panel"):
                inputs, run_button = _build_input_panel()
            with gr.Column(scale=5, min_width=360, elem_id="result-panel"):
                final_display, summary = _build_result_panel()

        display_evolution, trace = _build_inspection_tabs()
        outputs = DemoOutputs(final_display, summary, display_evolution, trace)

        run_button.click(
            fn=run_cid,
            inputs=[
                inputs.prompt,
                inputs.workspace,
                inputs.calculator_tool,
                inputs.symbolic_math_tool,
                inputs.workspace_tool,
                inputs.max_steps,
                inputs.need_threshold,
                inputs.convergence_threshold,
                inputs.allocation_threshold,
                inputs.binding_threshold,
                inputs.model_selector,
            ],
            outputs=[
                outputs.final_display,
                outputs.summary,
                outputs.display_evolution,
                outputs.trace,
            ],
            concurrency_limit=1,
            api_name="run_cid",
        )

        gr.HTML(FOOTER_HTML, elem_id="cid-footer")

    return demo


demo = build_demo()