diff --git a/WebHostLib/api/generate.py b/WebHostLib/api/generate.py index 5a66d1e693..8d44906501 100644 --- a/WebHostLib/api/generate.py +++ b/WebHostLib/api/generate.py @@ -1,15 +1,16 @@ import json import pickle +import typing from uuid import UUID from flask import request, session, url_for from markupsafe import Markup -from pony.orm import commit +from pony.orm import commit, select -from WebHostLib import app +from WebHostLib import app, cache from WebHostLib.check import get_yaml_data, roll_options from WebHostLib.generate import get_meta -from WebHostLib.models import Generation, STATE_QUEUED, Seed, STATE_ERROR +from WebHostLib.models import Generation, STATE_QUEUED, STATE_STARTED, Seed, STATE_ERROR from . import api_endpoints @@ -74,12 +75,23 @@ def generate_api(): def wait_seed_api(seed: UUID): seed_id = seed seed = Seed.get(id=seed_id) + reply_dict: dict[str, typing.Any] = {"queue_len": get_queue_length()} if seed: - return {"text": "Generation done"}, 201 + reply_dict["text"] = "Generation done" + return reply_dict, 201 generation = Generation.get(id=seed_id) if not generation: - return {"text": "Generation not found"}, 404 + reply_dict["text"] = "Generation not found" + return reply_dict, 404 elif generation.state == STATE_ERROR: - return {"text": "Generation failed"}, 500 - return {"text": "Generation running"}, 202 + reply_dict["text"] = "Generation failed" + return reply_dict, 500 + reply_dict["text"] = "Generation running" + return reply_dict, 202 + + +@cache.memoize(timeout=5) +def get_queue_length() -> int: + return select(generation for generation in Generation if + generation.state == STATE_STARTED or generation.state == STATE_QUEUED).count() diff --git a/WebHostLib/templates/waitSeed.html b/WebHostLib/templates/waitSeed.html index f2729353a6..9812e0ec34 100644 --- a/WebHostLib/templates/waitSeed.html +++ b/WebHostLib/templates/waitSeed.html @@ -30,10 +30,21 @@ } const data = await response.json(); - waitSeedDiv.innerHTML = ` -

Generation in Progress

-

${data.text}

- `; + if (data.queue_len === 1){ + waitSeedDiv.innerHTML = ` +

Generation in Progress

+

${data.text}

+

This is the only generation in the queue.

+ `; + } + else { + waitSeedDiv.innerHTML = ` +

Generation in Progress

+

${data.text}

+

There are ${data.queue_len} generations in the queue.

+ `; + } + setTimeout(checkStatus, 1000); // Continue polling. } catch (error) {