mllm_ui.py

from flask import Flask, request, url_for, redirect, abort
import mllm_task_file
import mllm_make_task
import mllm_types
from lxml import html
import lxml.builder
import subprocess
import re, sys
import reactivedb
import argparse, yaml
from mllm_types import EnvConfig, GenerationOutput, GenerationEnv
import mllm_state
from lxml import etree
import datetime, token_auth
from mllm_state import apply_output, mark_as_done
from pathlib import Path
from urllib.parse import urlparse
import json

class _E:
    def __getattr__(self, elem):
        def elem_builder(*children, **kwargs):
            real_children = []
            
            for ch in children:
                if isinstance(ch, dict):
                    kwargs.update(ch)
                elif isinstance(ch, list):
                    real_children += ch
                else:
                    real_children.append(ch)

            for child in real_children:
                assert child is not None
                    
            kwargs = {
                k.replace('_', '-').rstrip('_'):v
                for k, v in kwargs.items()
                if v is not None
            }
            return getattr(lxml.builder.E, elem)(*real_children, **kwargs)

        return elem_builder

E = _E()

app = Flask(__name__)

def parse_args():
    parser = argparse.ArgumentParser(description='MLLM UI Server')
    parser.add_argument('--task-file-path', type=Path, required=True,
                      help='Path to the tasks file')
    parser.add_argument('--env-config-path', type=Path, required=True,
                      help='Path to env config')
    parser.add_argument('--debug', action='store_true')
    parser.add_argument('--port', type=int, default=5000)

    return parser.parse_args()

args = parse_args()

def get_envconfig():
    with open(args.env_config_path, 'r') as f: env_config_data = yaml.safe_load(f)
    return EnvConfig.parse_obj(env_config_data)

def get_db():
    env_config = get_envconfig()
    db = reactivedb.Db(env_config.state_root / "db.sqlite3")
    return db

get_db()

def get_systemd_units(task_id):
    output = subprocess.check_output(
        ['systemctl', '--user', 'list-units', '--all', '--no-pager', '--no-legend'],
        text=True
    )
    units = []
    prefix = f'mllm-task-{task_id}-'
    suffix = '.service'
    
    for line in output.splitlines():
        if not line.strip():
            continue
        columns = line.split()
        if columns:
            unit_name = columns[0]
            if unit_name.startswith(prefix) and unit_name.endswith(suffix):
                units.append(unit_name)
    return units

def get_generation_outputs(task_id):
    table = get_db().table(mllm_types.GenerationOutput)
    result = table.query(task_id=task_id)
    return list(result.values())

class CSRFError(Exception):
    """Custom exception for CSRF validation failures."""
    pass

def validate_csrf(request):
    """Validates CSRF using Origin and Referer headers."""
    origin = request.headers.get('Origin')
    referer = request.headers.get('Referer')

    if not origin or not referer:
        raise CSRFError('CSRF validation failed: Origin or Referer headers missing')

    expected_host = request.host_url
    
    def get_base_url(url_string):
        if not url_string: return None
        parsed_url = urlparse(url_string)
        scheme = parsed_url.scheme
        netloc = parsed_url.netloc
        
        if scheme == 'http' and netloc.endswith(':80'):
            netloc = netloc[:-3]
        elif scheme == 'https' and netloc.endswith(':443'):
            netloc = netloc[:-4]
        
        return f"{scheme}://{netloc}/"

    base_origin = get_base_url(origin)
    base_expected_host = get_base_url(expected_host)

    if base_origin != base_expected_host:
        raise CSRFError(f'CSRF validation failed: Origin does not match expected host. Got {base_origin} expected {base_expected_host}')

    if not referer.startswith(base_expected_host):
        raise CSRFError(f'CSRF validation failed: Referer does not match expected host. Got {referer} expected {base_expected_host}')

@app.route('/')
@token_auth.token_required
def show_tasks():
    with open(args.task_file_path, 'r') as f:
        task_file_content = f.read()
    tasks = mllm_task_file.parse_text(task_file_content)
    
    done_html_elements = []
    undone_html_elements = []
    
    for task in tasks:
        task_id = task.task_id
        if not task_id:
            continue
        content = '\n'.join(task.content)
        
        lines = content.splitlines()
        first_line = lines[0].strip().lower() if lines else ""
        if first_line.startswith("done"):
            task_div = E.div(
                E.h2(
                    E.a(f"Task [{task_id}]", href=url_for('show_task_detail', task_id=task_id))
                ),
                E.pre(content)
            )
            done_html_elements.append(task_div)
        else:
            has_running_units = bool(get_systemd_units(task_id))
            status_marker = E.span("⚙️ ", style="color: #0066cc; animation: spin 1.5s linear infinite; display:inline-block; transform-origin: center; vertical-align: middle;") if has_running_units else ""
            
            task_div = E.div(
                E.h2(
                    status_marker,
                    E.a(f"Task [{task_id}]", href=url_for('show_task_detail', task_id=task_id))
                ),
                E.pre(content),
                E.form(
                    E.input(type='submit', value='Start', style='font-size: 70%',),
                    method='post',
                    action=url_for('start_task', task_id=task_id),
                    style='margin-top: -0.2em;'
                )
            )
            undone_html_elements.append(task_div)
    
    page = E.html(
        E.head(
            E.title("Tasks Overview"),
            E.style("""
@keyframes spin {
  0% { transform: rotate(0deg); }
  100% { transform: rotate(360deg); }
}
""")
        ),
        E.body(
            E.div(
                E.a("Home", href=url_for("show_tasks"))
            ),
            E.h1("Tasks Overview"),
            E.h2("Open Tasks"),
            *undone_html_elements,
            E.h2("Done Tasks"),
            *done_html_elements
        )
    )
    return html.tostring(page, pretty_print=True, encoding='unicode')

@app.route('/task/<task_id>')
@token_auth.token_required
def show_task_detail(task_id):
    with open(args.task_file_path, 'r') as f:
        task_file_content = f.read()
    tasks = mllm_task_file.parse_text(task_file_content)
    
    task = next((t for t in tasks if t.task_id == task_id), None)
    if not task:
        return f"Task {task_id} not found", 404
    
    content = '\n'.join(task.content)
    units = get_systemd_units(task_id)
    generation_outputs = get_generation_outputs(task_id)
    
    task_div = E.div(
        E.h2(f"Task [{task_id}]"),
        E.pre(content)
    )
    
    task_div.append(E.h3("Running Units"))
    if units:
        units_list = E.ul(*[E.li(unit) for unit in units])
        task_div.append(units_list)
    else:
        task_div.append(E.p("No running units"))
    
    task_div.append(E.h3("Generation Outputs"))
    if generation_outputs:
        outputs_table = E.table(
            E.tr(
                E.th("ID"),
                E.th("Created"),
                E.th("Base Revision"),
                E.th("Meta"),
                E.th("Build Status"),
                E.th("Changed Files"),
                E.th("Applied")
            ),
            style="border-collapse: collapse; width: 100%"
        )

        generation_outputs.sort(key=lambda o: -o.id)
        for gen_output in generation_outputs:
            build_status_text = "Unknown" if gen_output.build_successful is None else \
                           "Success" if gen_output.build_successful else "Failed"

            build_status_style = {
                None: "",
                True: "color: green;",
                False: "color: red;"
            }[gen_output.build_successful]
            
            created_date = datetime.datetime.fromtimestamp(gen_output.created).strftime('%Y-%m-%d %H:%M:%S')
            meta_str = str(gen_output.meta) if gen_output.meta is not None else ""
            applied_text = 'Yes' if gen_output.applied else 'No'

            row = E.tr(
                E.td(
                    E.a(
                        str(gen_output.id),
                        href=url_for('show_generation_output', id=gen_output.id)
                    )
                ),
                E.td(created_date),
                E.td(E.span((gen_output.base_revision or "")[:8], style='font-family: monospace')),
                E.td(meta_str),
                E.td(E.span(build_status_text, style=build_status_style)),
                E.td(', '.join(sorted(gen_output.changed_files))),
                E.td(applied_text)
            )
            outputs_table.append(row)
        
        for td in outputs_table.xpath("//td"):
            td.set("style", "border: 1px solid #ddd; padding: 8px;")
        for th in outputs_table.xpath("//th"):
            th.set("style", "border: 1px solid #ddd; padding: 8px; background-color: #f2f2f2;")
        
        task_div.append(outputs_table)
    else:
        task_div.append(E.p("No generation outputs"))
    
    start_form = E.form(
        E.input(type='submit', value='Start'),
        method='post',
        action=url_for('start_task', task_id=task_id)
    )
    task_div.append(start_form)  # Add start task button to the task detail view
    
    page = E.html(
        E.head(E.title(f"Task {task_id} Details")),
        E.body(
            E.div(
                E.a("Home", href=url_for("show_tasks"))
            ),
            E.h1(f"Task {task_id} Details"),
            task_div
        )
    )
    return html.tostring(page, pretty_print=True, encoding='unicode')

@app.route('/task/<task_id>/start', methods=['POST'])
@token_auth.token_required
def start_task(task_id):
    validate_csrf(request)

    with open(args.task_file_path, 'r') as f:
        task_file_content = f.read()
    tasks = mllm_task_file.parse_text(task_file_content)
    task = next((t for t in tasks if t.task_id == task_id), None)
    if not task:
        return f"Task {task_id} not found", 404

    env_config = get_envconfig()
    raw_task_prompt = "\n".join(task.content)
    try:
        mllm_make_task.parse_task_prompt(raw_task_prompt, args.env_config_path.parent.resolve())
    except Exception as e:
        return f"Error in parsing task prompt: {e}", 400

    cwd = Path(__file__).parent

    cmd = [
        'systemd-run', '--user', '--unit=mllm-task-%s-start' % task.task_id, '--collect',
        '--working-directory=%s' % cwd,
        '--slice=mllm.slice',
        sys.executable, 'mllm_make_task.py',
        '--env-config', str(args.env_config_path),
        '--task-file', str(args.task_file_path),
        '--task-id', task_id,
        '--and-start-prompt'
    ]
    
    subprocess.check_call(cmd)
    return redirect(url_for('show_task_detail', task_id=task_id))

def render_messages(messages):
    elements = []
    if not isinstance(messages, list):
        elements.append(E.p("No messages to display"))
        return elements
    for message in messages:
        if not isinstance(message, dict):
            continue
        role = message.get('role', 'unknown')
        content = message.get('content', '')
        message_div = E.div(
            E.h4(f"Role: {role}"),
            E.pre(content, style="max-height: 10em; overflow-y: auto")
        )
        elements.append(message_div)
    return elements

def render_diff(gen_output):
    diffs = mllm_state.generate_diff_for_output(get_envconfig(), gen_output, html=True)
    return [ E.div(E.h1(fn), E.div(etree.fromstring(diff, etree.HTMLParser())) if diff else '(empty)')
             for fn, diff in sorted(diffs.items()) ]

@app.route('/generation_output/<int:id>')
@token_auth.token_required
def show_generation_output(id):
    db = get_db()
    table = db.table(mllm_types.GenerationOutput)
    gen_output = table.get(id=id)

    siblings = list(table.query(task_id=gen_output.task_id).values())
    siblings.sort(key=lambda x: x.id)
    ids = [g.id for g in siblings]
    current_index = ids.index(gen_output.id)
    prev_id = ids[current_index - 1] if current_index > 0 else None
    next_id = ids[current_index + 1] if current_index < len(ids) - 1 else None

    if gen_output.build_successful is None:
        build_status = 'unknown build status'
    elif gen_output.build_successful:
        build_status = E.div('build passed')
    else:
        build_status = E.div(
            E.div('build failed', style='color: red'),
            E.pre(gen_output.build_output))
    
    created_date = datetime.datetime.fromtimestamp(gen_output.created).strftime('%Y-%m-%d %H:%M:%S')

    if isinstance(gen_output.meta, dict):
        meta_content = E.pre(json.dumps(gen_output.meta, indent=2))
    elif gen_output.meta is not None:
        meta_content = E.p(str(gen_output.meta))
    else:
        meta_content = E.p("No metadata available")

    model_line = E.p(
        E.span("Model: ", style="font-weight: bold;"),
        meta_content)
    
    applied_text = 'Yes' if gen_output.applied else 'No'

    if not gen_output.applied:
        apply_form = E.form(
            E.input(type='submit', value='Apply'),
            method='post',
            action=url_for('apply_generation_output', id=id)
        )
    else:
        apply_form = E.p("This output has already been applied.")

    nav_links = []
    nav_links.append(E.a("Home", href=url_for("show_tasks")))
    nav_links.append(' ')
    nav_links.append(E.a(f"Task {gen_output.task_id}", href=url_for('show_task_detail', task_id=gen_output.task_id)))
    nav_links.append(' ')
    if prev_id is not None:
        nav_links.append(E.a("← Previous", href=url_for('show_generation_output', id=prev_id), id="prevOutputLink", title='Ctrl-LeftArrow'))
        nav_links.append(' ')
    if next_id is not None:
        nav_links.append(E.a("Next →", href=url_for('show_generation_output', id=next_id), id="nextOutputLink", title='Ctrl-RightArrow'))

    page = E.html(
        E.head(
            E.title(f"Generation Output {id}"),
            E.script(
                """
document.addEventListener('keydown', function(e) {
  if (e.ctrlKey && e.key === 'ArrowLeft') {
    let link = document.getElementById('prevOutputLink');
    if (link) window.location.replace(link.getAttribute('href'));
  }
  if (e.ctrlKey && e.key === 'ArrowRight') {
    let link = document.getElementById('nextOutputLink');
    if (link) window.location.replace(link.getAttribute('href'));
  }
});
                """
            )
        ),
        E.body(
            E.div(*nav_links, style="margin-bottom: 1em;"),

            E.h1(f"Generation Output {id}"),
            E.p(E.span("Created: ", style="font-weight: bold;"), created_date),
            model_line,
            E.p(f"Applied: {applied_text}"),
            build_status,
            apply_form,
            *render_diff(gen_output),
            E.h2("Messages"),
            *render_messages(gen_output.messages),
            E.h2("Details"),
            E.pre(str(gen_output))
        )
    )
    return html.tostring(page, pretty_print=True, encoding='unicode')

@app.route('/generation_output/<int:id>/apply', methods=['POST'])
@token_auth.token_required
def apply_generation_output(id):
    validate_csrf(request)

    force = request.args.get('force')
    
    db = get_db()
    table = db.table(mllm_types.GenerationOutput)
    gen_output = table.get(id=id)

    if not gen_output:
        return f"GenerationOutput {id} not found", 404

    if gen_output.applied:
        return f"GenerationOutput {id} has already been applied", 400

    env_config = get_envconfig()
    root_path = args.env_config_path.parent.resolve()

    try:
        apply_output(env_config, root_path, gen_output, force)
    except mllm_state.MergeConflictError as conflict_error:
        return render_conflict_resolution_page(conflict_error.conflict_files, gen_output)
        
    gen_output.applied = True
    table.set(gen_output)
    db.commit()

    mark_as_done(args.task_file_path, gen_output.task_id)
    
    return redirect('/')

def render_conflict_resolution_page(conflict_files, gen_output):
    conflict_list_items = [E.li(file) for file in conflict_files]
    page = E.html(
        E.head(E.title(f"Merge Conflicts for Generation Output {gen_output.id}")),
        E.body(
            E.h1(f"Merge Conflicts Detected"),
            E.p("The following files have merge conflicts:"),
            E.ul(*conflict_list_items),
            E.form(
                E.input(type='submit', value='Accept with Conflicts'),
                method='post',
                action=f"/generation_output/{gen_output.id}/apply?force=true"
            )
        )
    )
    return html.tostring(page, pretty_print=True, encoding='unicode')

if __name__ == '__main__':
    token_auth.install(app, socketio=None, app_name="mllm_ui")
    app.run(debug=args.debug, port=args.port)